Imhotep/packages/imhotep-playwright/src/fol-compiler.ts

/**
 * FOL compiler for Playwright public API assertions.
 *
 * Converts accumulated FluentAssertion / FluentRelation / FluentQuantifier
 * objects into FormulaNode trees suitable for evaluateLogic().
 *
 * Every assertion compiles to a quantified formula over element domains,
 * even singleton assertions, so that the logic engine can evaluate them
 * uniformly.
 */

import type {
  FormulaNode,
  DomainRef,
  VariableRef,
  PredicateCall,
  TermNode,
} from 'imhotep-core'

import type {
  FluentAssertion,
  FluentRelation,
  FluentQuantifier,
} from 'imhotep-dsl'

import { FluentQuantifier as FQ, FluentRelation as FR } from 'imhotep-dsl'

// ---------------------------------------------------------------------------
// AST Adapter: grammar.ts shape → logic-ast.ts shape
// ---------------------------------------------------------------------------

/**
 * Adapts a formula from grammar.ts AST shape to logic-ast.ts AST shape.
 * Supports both shapes for backward compatibility:
 * - If the node has a `kind` property, it's already in logic-ast shape.
 * - If the node has a `type` property, it's in grammar.ts shape and needs conversion.
 */
function adaptGrammarFormulaToLogicAst(node: any): FormulaNode {
  // Already in logic-ast.ts shape — pass through
  if (node && typeof node === 'object' && 'kind' in node) {
    return node as FormulaNode
  }

  // Not an object or null — cannot adapt
  if (!node || typeof node !== 'object') {
    throw new TypeError('Cannot adapt non-object formula node')
  }

  const type = node.type
  switch (type) {
    case 'ForAll': {
      return {
        type: 'FormulaNode',
        kind: 'forall',
        bindings: [
          {
            type: 'TupleBinding',
            variables: [node.variable?.name ?? 'x'],
            domain: adaptDomainRef(node.domain),
          },
        ],
        body: adaptGrammarFormulaToLogicAst(node.body),
      } as FormulaNode
    }

    case 'Exists': {
      return {
        type: 'FormulaNode',
        kind: 'exists',
        bindings: [
          {
            type: 'TupleBinding',
            variables: [node.variable?.name ?? 'x'],
            domain: adaptDomainRef(node.domain),
          },
        ],
        body: adaptGrammarFormulaToLogicAst(node.body),
      } as FormulaNode
    }

    case 'And': {
      return {
        type: 'FormulaNode',
        kind: 'and',
        left: adaptGrammarFormulaToLogicAst(node.left),
        right: adaptGrammarFormulaToLogicAst(node.right),
      } as FormulaNode
    }

    case 'Or': {
      return {
        type: 'FormulaNode',
        kind: 'or',
        left: adaptGrammarFormulaToLogicAst(node.left),
        right: adaptGrammarFormulaToLogicAst(node.right),
      } as FormulaNode
    }

    case 'Not': {
      return {
        type: 'FormulaNode',
        kind: 'not',
        operand: adaptGrammarFormulaToLogicAst(node.operand),
      } as FormulaNode
    }

    case 'Implies': {
      return {
        type: 'FormulaNode',
        kind: 'implies',
        antecedent: adaptGrammarFormulaToLogicAst(node.left),
        consequent: adaptGrammarFormulaToLogicAst(node.right),
      } as FormulaNode
    }

    case 'PredicateCall': {
      return {
        type: 'FormulaNode',
        kind: 'predicate',
        predicate: node.name ?? 'unknown',
        args: (node.args ?? []).map((arg: any) => adaptTermNode(arg)),
      } as FormulaNode
    }

    case 'VariableRef': {
      return node as FormulaNode
    }

    default: {
      throw new TypeError(`Unknown grammar AST node type: ${type}`)
    }
  }
}

/**
 * Adapts a DomainRef from grammar.ts shape to logic-ast.ts shape.
 * grammar.ts uses `kind`, logic-ast.ts uses `domain`.
 */
function adaptDomainRef(domain: any): DomainRef {
  if (!domain || typeof domain !== 'object') {
    return { type: 'DomainRef', domain: 'elements' } as DomainRef
  }
  const domainName = domain.domain || domain.kind || 'elements'
  return {
    type: 'DomainRef',
    domain: domainName,
    selector: domain.selector,
    parentVar: domain.parentVar,
  } as DomainRef
}

/**
 * Adapts a term node (VariableRef, DomainRef, or literal) to logic-ast.ts shape.
 * VariableRef is already compatible. DomainRef needs kind→domain mapping.
 * String/number literals pass through (engine has fallback handler).
 */
function adaptTermNode(arg: any): any {
  if (!arg || typeof arg !== 'object') {
    return arg
  }
  if (arg.type === 'VariableRef') {
    return arg
  }
  if (arg.type === 'DomainRef') {
    return adaptDomainRef(arg)
  }
  return arg
}

// ---------------------------------------------------------------------------
// Domain helpers
// ---------------------------------------------------------------------------

function makeDomain(selector: string): DomainRef {
  return {
    type: 'DomainRef',
    domain: 'elements',
    selector,
  }
}

function makeVar(name: string): VariableRef {
  return {
    type: 'VariableRef',
    name,
  }
}

// ---------------------------------------------------------------------------
// Predicate compilation
// ---------------------------------------------------------------------------

function makePredicate(
  name: string,
  args: TermNode[],
  options?: Record<string, unknown>,
): PredicateCall {
  const call: PredicateCall = {
    type: 'FormulaNode',
    kind: 'predicate',
    predicate: name,
    args,
  }
  if (options && Object.keys(options).length > 0) {
    ;(call as any).options = options
  }
  return call
}

// ---------------------------------------------------------------------------
// Relation → Formula
// ---------------------------------------------------------------------------

function compileRelation(
  relation: FluentRelation,
  subjectVar: VariableRef,
  refVar: VariableRef,
): FormulaNode {
  const getAssertionSpace = (): 'layout' | 'visual' | undefined => {
    const maybeAssertion = (relation as any).assertion
    if (maybeAssertion && typeof maybeAssertion.getSpace === 'function') {
      return maybeAssertion.getSpace()
    }
    return undefined
  }

  // Handle compound relations (.and / .or chaining)
  const compoundParts = (relation as any)._compoundParts as Array<{ relation: string; referenceSelector: string; options: Record<string, unknown> }> | undefined
  const compoundOperator = (relation as any)._compoundOperator as 'and' | 'or' | undefined

  if (compoundParts && compoundParts.length > 1 && compoundOperator) {
    const assertionSpace = getAssertionSpace()
    // Each compound part may have a different reference selector.
    // Create a separate reference variable + forall binding per part.
    let result: FormulaNode | null = null
    for (let i = 0; i < compoundParts.length; i++) {
      const part = compoundParts[i]
      const partOptions = {
        ...part.options,
        ...(part.options?.space === undefined && assertionSpace !== undefined ? { space: assertionSpace } : {}),
      }
      const partRefVar = makeVar(`$ref_${i}`)
      const partBody = compileSingleRelation(part.relation, part.referenceSelector, partOptions, subjectVar, partRefVar)
      // Wrap in forall over the part's reference domain
      const partFormula: FormulaNode = {
        type: 'FormulaNode',
        kind: 'forall',
        bindings: [{
          type: 'TupleBinding',
          variables: [partRefVar.name],
          domain: makeDomain(part.referenceSelector),
        }],
        body: partBody,
      }
      if (result === null) {
        result = partFormula
      } else {
        result = {
          type: 'FormulaNode',
          kind: compoundOperator,
          left: result,
          right: partFormula,
        }
      }
    }
    return result!
  }

  const assertionSpace = getAssertionSpace()
  const relationOptions: Record<string, unknown> = {
    ...(relation.options as Record<string, unknown>),
    ...(((relation.options as Record<string, unknown>)?.space === undefined && assertionSpace !== undefined)
      ? { space: assertionSpace }
      : {}),
  }
  return compileSingleRelation(relation.relation, relation.referenceSelector, relationOptions, subjectVar, refVar)
}

function compileSingleRelation(
  name: string,
  _referenceSelector: string,
  relOptions: Record<string, unknown>,
  subjectVar: VariableRef,
  refVar: VariableRef,
): FormulaNode {
  // Build options map for the predicate evaluator.
  const ropts = relOptions as {
    minGap?: number | string | { value: number; unit: string }
    maxGap?: number | string | { value: number; unit: string }
    tolerance?: number | string
    axis?: string
    min?: number | string | { value: number; unit: string }
    max?: number | string | { value: number; unit: string }
    dimension?: string
    value?: number | string | { value: number; unit: string }
    mode?: 'atLeast' | 'atMost' | 'exact' | 'between'
    space?: 'layout' | 'visual'
    inStackingContext?: boolean
  }
  const predicateOptions: Record<string, unknown> = {}
  if (ropts.minGap !== undefined) predicateOptions.minGap = parseNumeric(ropts.minGap)
  if (ropts.maxGap !== undefined) predicateOptions.maxGap = parseNumeric(ropts.maxGap)
  if (ropts.tolerance !== undefined) predicateOptions.tolerance = parseNumeric(ropts.tolerance)
  if (ropts.axis !== undefined) predicateOptions.axis = ropts.axis
  if (ropts.space !== undefined) predicateOptions.space = ropts.space
  if (ropts.min !== undefined) predicateOptions.min = parseNumeric(ropts.min)
  if (ropts.max !== undefined) predicateOptions.max = parseNumeric(ropts.max)
  if (ropts.dimension !== undefined) predicateOptions.dimension = ropts.dimension
  if (ropts.value !== undefined) {
    const val = parseNumeric(ropts.value)
    // Size proxies store the bound in 'value' keyed by mode.
    if (ropts.mode === 'atMost') predicateOptions.max = val
    else if (ropts.mode === 'between') {
      // Between stores min/max directly on options already.
    }
    else predicateOptions.min = val
  }
  if (ropts.mode !== undefined) predicateOptions.mode = ropts.mode

  if (name.startsWith('size.')) {
    // Size assertions are unary — only subject.
    return makePredicate(name.replace('size.', ''), [subjectVar], predicateOptions)
  }

  const base = makePredicate(name, [subjectVar, refVar], predicateOptions)
  if (ropts.inStackingContext === true) {
    return {
      type: 'FormulaNode',
      kind: 'and',
      left: base,
      right: makePredicate('inStackingContext', [subjectVar, refVar]),
    }
  }

  return base
}

function parseNumeric(v: string | number | { value: number; unit?: string } | undefined): number | { value: number; unit: string } | undefined {
  if (v === undefined) return undefined
  if (typeof v === 'number') return v
  if (typeof v === 'object' && typeof v.value === 'number') {
    const unit = typeof v.unit === 'string' ? v.unit.toLowerCase() : 'px'
    return unit === 'px' ? v.value : { value: v.value, unit }
  }
  if (typeof v !== 'string') return undefined
  const trimmed = v.trim()
  const m = trimmed.match(/^(-?\d*\.?\d+)\s*([a-z%]+)?$/i)
  if (!m) return undefined
  const n = parseFloat(m[1] ?? '')
  if (isNaN(n)) return undefined
  const unit = (m[2] ?? 'px').toLowerCase()
  if (unit !== 'px') {
    return { value: n, unit }
  }
  return n
}

// ---------------------------------------------------------------------------
// Assertion → Formula (with quantifier support)
// ---------------------------------------------------------------------------

interface CompiledAssertion {
  formula: FormulaNode
  subjectSelector: string
  referenceSelector: string | null
}

function compileSingleAssertion(
  assertion: FluentRelation | FluentAssertion | FluentQuantifier,
): CompiledAssertion | null {
  // FluentQuantifier (forAll / exists from static methods)
  // Duck-type check for cross-package compatibility (avoids instanceof
  // breakage when packages are symlinked or duplicated).
  if (
    assertion &&
    typeof assertion === 'object' &&
    'bindings' in assertion &&
    Array.isArray((assertion as any).bindings) &&
    'toFormula' in assertion &&
    typeof (assertion as any).toFormula === 'function'
  ) {
    const q = assertion as FluentQuantifier
    const rawFormula = q.toFormula()
    return {
      formula: adaptGrammarFormulaToLogicAst(rawFormula),
      subjectSelector: q.bindings[0]?.selector ?? '',
      referenceSelector: null,
    }
  }

  // FluentCardinality (exactlyOne, atLeastN, atMostN)
  // These are evaluated outside the FOL engine using selector match counts.
  if (
    assertion &&
    typeof assertion === 'object' &&
    'kind' in assertion &&
    ['exactlyOne', 'atLeastN', 'atMostN'].includes((assertion as any).kind)
  ) {
    return null
  }

  // FluentRelation (terminal like .leftOf('.b'))
  // Duck-type check for cross-package compatibility.
  if (
    assertion &&
    typeof assertion === 'object' &&
    'relation' in assertion &&
    typeof (assertion as any).relation === 'string'
  ) {
    const rel = assertion as FluentRelation
    const subjectSel = rel.assertion.getSubject()
    const refSel = rel.referenceSelector
    const quantifier = rel.assertion.getQuantifier()

    const isCompound = !!(rel as any)._compoundParts && (rel as any)._compoundParts.length > 1

    const subjectVar = makeVar('$subject')
    const refVar = makeVar('$reference')
    const body = compileRelation(rel, subjectVar, refVar)

    const subjectBinding = {
      type: 'TupleBinding' as const,
      variables: ['$subject'],
      domain: makeDomain(subjectSel),
    }
    // For compound relations, each part already has its own reference forall.
    // Skip the outer refBinding to avoid an unused quantifier.
    const refBinding = refSel && !isCompound
      ? {
          type: 'TupleBinding' as const,
          variables: ['$reference'],
          domain: makeDomain(refSel),
        }
      : null

    let formula: FormulaNode

    switch (quantifier) {
      case 'any': {
        const anyBody = refBinding
          ? ({
              type: 'FormulaNode' as const,
              kind: 'forall',
              bindings: [refBinding],
              body,
            } as FormulaNode)
          : body
        formula = {
          type: 'FormulaNode',
          kind: 'exists',
          bindings: [subjectBinding],
          body: anyBody,
        }
        break
      }
      case 'none': {
        const noneBody = refBinding
          ? ({
              type: 'FormulaNode' as const,
              kind: 'exists',
              bindings: [refBinding],
              body,
            } as FormulaNode)
          : body
        formula = {
          type: 'FormulaNode',
          kind: 'forall',
          bindings: [subjectBinding],
          body: {
            type: 'FormulaNode',
            kind: 'not',
            operand: noneBody,
          },
        }
        break
      }
      case 'all':
      default: {
        // Default and 'all' mean "every subject satisfies against every reference".
        const bindings: any[] = [subjectBinding]
        if (refBinding) bindings.push(refBinding)
        formula = {
          type: 'FormulaNode',
          kind: 'forall',
          bindings,
          body,
        }
        break
      }
    }

    return {
      formula,
      subjectSelector: subjectSel,
      referenceSelector: refSel,
    }
  }

  // FluentAssertion without relation (incomplete — should not happen in practice)
  return null
}

// ---------------------------------------------------------------------------
// Public: compile a list of assertions to a single FormulaNode
// ---------------------------------------------------------------------------

export function compileAssertionsToFOL(
  assertions: (FluentRelation | FluentAssertion | FluentQuantifier)[],
): FormulaNode | null {
  const compiled: FormulaNode[] = []

  for (const assertion of assertions) {
    const c = compileSingleAssertion(assertion)
    if (c) {
      compiled.push(c.formula)
    }
  }

  if (compiled.length === 0) {
    return null
  }

  if (compiled.length === 1) {
    return compiled[0]
  }

  // Combine with And
  let result = compiled[0]
  for (let i = 1; i < compiled.length; i++) {
    result = {
      type: 'FormulaNode',
      kind: 'and',
      left: result,
      right: compiled[i],
    }
  }
  return result
}

// Re-export for public.ts
export type { CompiledAssertion }