import type {
  Combinator, ParseContext, ParseResult, ParserMeta, FirstSet,
  AutoNotCheck, CharRange, ChoiceStrategy, GatedArm,
} from '../types.ts'
import { union, intersects, matchesEmpty } from './first-set.ts'
import { deriveExpected } from './expect.ts'
import { rollbackTrivia, saveTriviaMark } from './trivia-skip.ts'
import { scalarOf } from './scalar.ts'

type ArmParser<T> = T extends GatedArm<infer U> ? Combinator<U> : T extends Combinator<infer U> ? Combinator<U> : never
type UnionArms<T extends (Combinator<unknown> | GatedArm<unknown>)[]> = {
  [K in keyof T]: ArmParser<T[K]>
}[number] extends Combinator<infer U> ? U : never

export function choice<T extends [Combinator<unknown> | GatedArm<unknown>, ...(Combinator<unknown> | GatedArm<unknown>)[]]>(
  ...args: T
): Combinator<UnionArms<T>> {
  // Unwrap gated arms into (parser, gate) pairs
  const parsers = args.map(a => ('gate' in a ? a.combinator : a)) as Combinator<unknown>[]
  const gates   = args.map(a => ('gate' in a ? a.gate   : null)) as (((state: unknown) => boolean) | null)[]

  const hasGates = gates.some(g => g !== null)

  // O(1) first-char dispatch is sound when arms are pairwise-disjoint in their
  // first char AND no arm can match empty. Non-nullability is what makes a GATED
  // arm keep its dispatch slot: a gated arm is normally *skipped* so a later arm
  // can retry the same position — but if every arm is non-nullable and first-sets
  // are disjoint, no later arm can match this first char, so "skip the gate then
  // retry" is exactly "dispatch to this arm, check its gate, and fail the choice
  // if the gate is false". Every OTHER first char dispatches as before and never
  // touches the gate. A nullable arm matches at ANY position (zero-width), so
  // first-char dispatch can't represent it → such choices stay on firstMatch.
  // (Requiring non-nullability universally also tightens the ungated path, which
  // was already unsound for a nullable-but-first-set-disjoint arm; in practice no
  // ungated disjoint choice has a nullable arm, so codegen output is unchanged.)
  const disjoint = areDisjoint(parsers.map(p => p._meta.firstSet))
    && parsers.every(p => !matchesEmpty(p))

  let combined: FirstSet = { kind: 'empty' }
  for (const p of parsers) combined = union(combined, p._meta.firstSet)

  const meta: ParserMeta = {
    firstSet: combined,
    canMatchNewline: parsers.some(p => p._meta.canMatchNewline),
    isTrivia: false,
    disjoint,
  }

  // A disjoint choice dispatches by first char (gated arms check their gate inside
  // the dispatched branch); a NON-disjoint gated choice falls to firstMatch — the
  // greedy/longest-first strategies are incompatible with per-arm predicates.
  const strategy = (disjoint || hasGates) ? null : detectStrategy(parsers)
  const ordered = strategy?.tag === 'firstMatch' || strategy?.tag === 'sharedPrefix'
  // `sharedPrefix` is a `firstMatch` specialization — the interpreter runs it via
  // the firstMatch loop below, and codegen falls back to firstMatch under
  // coverage/recovery/linkable — so it needs the same auto-not table computed.
  const autoNot = (!disjoint && !hasGates && ordered)
    ? computeAutoNot(parsers)
    : parsers.map(() => null)

  // Runtime state for each strategy (built once, reused on every parse call):
  let greedyLitMap: Map<string, number> | null = null
  let sortedParsers: Combinator<unknown>[] | null = null
  const asciiDispatch = disjoint ? buildAsciiDispatch(parsers) : null

  if (strategy?.tag === 'greedyClassify') {
    greedyLitMap = new Map()
    for (let i = 0; i < parsers.length; i++) {
      if (i === strategy.superIndex) continue
      const litVal = getCoreLiteralValue(parsers[i]!)
      if (litVal !== null) greedyLitMap.set(litVal, i)
    }
  } else if (strategy?.tag === 'literalsLongestFirst') {
    sortedParsers = strategy.sortedIndices.map(i => parsers[i]!)
  }

  const scalarParsers = parsers.map(scalarOf)
  const scalarEligible = !hasGates && parsers.every(p => p._parseScalar !== undefined)
    && (disjoint || ordered)
  const scalarExpected = parsers.flatMap(deriveExpected)
  const parseScalar = (input: string, pos: number, ctx: ParseContext): number => {
    if (disjoint) {
      const code = pos < input.length ? input.codePointAt(pos)! : -1
      let idx = code >= 0 && code < 128 ? asciiDispatch![code]! : -1
      if (idx < 0 && code >= 0) {
        for (let i = 0; i < parsers.length; i++) {
          if (inFirstSet(code, parsers[i]!._meta.firstSet)) { idx = i; break }
        }
      }
      if (idx >= 0) {
        const gate = gates[idx]
        if (!gate || gate(ctx.state)) return scalarParsers[idx]!(input, pos, ctx)
        ctx._fx = deriveExpected(parsers[idx]!)
        return ~pos
      }
    } else {
      for (let i = 0; i < scalarParsers.length; i++) {
        if (gates[i] && !gates[i]!(ctx.state)) continue
        const end = scalarParsers[i]!(input, pos, ctx)
        if (end >= 0 && (!autoNot[i] || !autoNotFires(input, end, autoNot[i]!))) return end
      }
    }
    ctx._fx = scalarExpected
    return ~pos
  }

  return {
    _tag: 'choice',
    _meta: meta,
    _def: {
      tag: 'choice',
      parsers,
      gates,
      disjoint,
      strategy: strategy ?? { tag: 'firstMatch' },
      autoNot,
    },
    _parseScalar: scalarEligible ? parseScalar : undefined,
    parse(input: string, pos: number, ctx: ParseContext): ParseResult<UnionArms<T>> {
      const expected: string[] = []
      let expectedAt = pos

      // ── Disjoint: O(1) first-char dispatch (arms may be gated) ────────────
      //
      // EOF IS A DISPATCH MISS, not a reason to leave the disjoint path. Every arm
      // of a disjoint choice is non-nullable (that is a precondition of `disjoint`
      // above), so at EOF no arm can match and the answer is the same "nothing
      // could have started here" the in-bounds miss below gives. Falling through to
      // firstMatch instead made this position the ONE place a gated-off arm was
      // dropped from the report — firstMatch `continue`s past it and contributes
      // nothing — while the in-bounds miss runs `parsers.flatMap` ignoring gates and
      // DOES name it. Same gate state, same "no arm can match", two different
      // answers; codegen and both table drivers only ever gave the union. Accept and
      // reject are untouched: every arm fails at EOF on either path.
      if (disjoint) {
        const code = pos < input.length ? input.codePointAt(pos)! : -1
        let idx = code >= 0 && code < 128 ? asciiDispatch![code]! : -1
        if (idx < 0 && code >= 0) {
          for (let i = 0; i < parsers.length; i++) {
            if (inFirstSet(code, parsers[i]!._meta.firstSet)) { idx = i; break }
          }
        }
        if (idx >= 0) {
          const gate = gates[idx]
          if (gate && !gate(ctx.state)) {
            // Gate blocks this arm. Disjointness + non-nullable arms guarantee no
            // OTHER arm can match this first char, so skip-and-retry is exactly
            // fail-the-choice — we must not fall through to another arm.
            return { ok: false, expected: deriveExpected(parsers[idx]!), span: { start: pos, end: pos } }
          }
          const result = parsers[idx]!.parse(input, pos, ctx)
          if (result.ok) return result as ParseResult<UnionArms<T>>
          expected.push(...result.expected)
          if (result.committed) return { ok: false, expected, span: { start: pos, end: pos }, committed: true }
          return { ok: false, expected, span: { start: pos, end: pos } }
        }
        return {
          ok: false,
          expected: parsers.flatMap(p => {
            const r = p.parse(input, pos, ctx)
            return r.ok ? [] : r.expected
          }),
          span: { start: pos, end: pos },
        }
      }

      // ── greedyClassify: run one regex, classify by string equality ─────────
      //    One parse call total. No backtracking.
      if (strategy?.tag === 'greedyClassify') {
        const superResult = parsers[strategy.superIndex]!.parse(input, pos, ctx)
        if (!superResult.ok) return superResult as ParseResult<UnionArms<T>>

        const end = superResult.span.end
        const litIdx = greedyLitMap!.get(input.slice(pos, end))
        if (litIdx !== undefined) {
          const litVal = getCoreLiteralValue(parsers[litIdx]!)!
          const value = applyTransforms(parsers[litIdx]!, litVal, { start: pos, end })
          return { ok: true, value: value as UnionArms<T>, span: { start: pos, end } }
        }
        return superResult as ParseResult<UnionArms<T>>
      }

      // ── literalsLongestFirst: sorted descending by length, no backtracking ─
      if (strategy?.tag === 'literalsLongestFirst') {
        for (const p of sortedParsers!) {
          const r = p.parse(input, pos, ctx)
          if (r.ok) return r as ParseResult<UnionArms<T>>
          expected.push(...r.expected)
        }
        return { ok: false, expected, span: { start: pos, end: pos } }
      }

      // ── firstMatch (+ gated arms): try each arm in order, skipping gated-off arms ──
      for (let i = 0; i < parsers.length; i++) {
        if (gates[i] && !gates[i]!(ctx.state)) continue   // gate blocks this arm
        // Save leaf-array lengths so a failed/rejected arm can be rolled back.
        const mark = saveTriviaMark(ctx)
        const result = parsers[i]!.parse(input, pos, ctx)
        if (!result.ok) {
          rollbackTrivia(ctx, mark)
          const at = result.span.start
          if (at > expectedAt) { expectedAt = at; expected.length = 0 }
          if (at === expectedAt) expected.push(...result.expected)
          if (result.committed) return { ok: false, expected, span: result.span, committed: true }
          continue
        }
        const checks = autoNot[i]
        if (checks && autoNotFires(input, result.span.end, checks)) {
          rollbackTrivia(ctx, mark)
          continue
        }
        return result as ParseResult<UnionArms<T>>
      }
      // A nested gated choice can decline every arm after an enclosing arm has
      // already advanced. Do not let that empty dynamic set erase both the
      // shallower arms and the choice's own static opener contract.
      if (expected.length === 0) {
        for (const parser of parsers) expected.push(...deriveExpected(parser))
      }
      return { ok: false, expected, span: { start: pos, end: pos } }
    },
  }
}

// ---------------------------------------------------------------------------
// Strategy detection
// ---------------------------------------------------------------------------

function detectStrategy(parsers: Combinator<unknown>[]): ChoiceStrategy {
  // greedyClassify: exactly one regex arm whose regex matches every literal arm's
  // value exactly (and can potentially match more). All other arms must be literals.
  const regexIndices: number[] = []
  const literalIndices: number[] = []
  for (let i = 0; i < parsers.length; i++) {
    if (getCoreRegexDef(parsers[i]!) !== null) regexIndices.push(i)
    else if (getCoreLiteralValue(parsers[i]!) !== null) literalIndices.push(i)
  }

  if (
    regexIndices.length === 1 &&
    literalIndices.length === parsers.length - 1 &&
    literalIndices.length > 0
  ) {
    const superIndex = regexIndices[0]!
    const regexDef = getCoreRegexDef(parsers[superIndex]!)!
    const flags = 'y' + regexDef.flags.replace(/[gy]/g, '')
    const re = new RegExp(regexDef.source, flags)
    const allSubsumed = literalIndices.every(i => {
      const litVal = getCoreLiteralValue(parsers[i]!)!
      re.lastIndex = 0
      const m = re.exec(litVal)
      return m !== null && m[0] === litVal
    })
    if (allSubsumed) return { tag: 'greedyClassify', superIndex }
  }

  // literalsLongestFirst: every arm is a literal — try longest first, no backtracking
  if (parsers.length === literalIndices.length) {
    const sortedIndices = [...literalIndices].sort((a, b) =>
      getCoreLiteralValue(parsers[b]!)!.length - getCoreLiteralValue(parsers[a]!)!.length
    )
    return { tag: 'literalsLongestFirst', sortedIndices }
  }

  // sharedPrefix: EVERY arm is a bare sequence beginning with the SAME concrete
  // leading literal/regex — the compiler parses that left factor once. Detected
  // last (its shape is disjoint from the literal-only strategies above).
  const shared = detectSharedPrefix(parsers)
  if (shared !== null) return shared

  return { tag: 'firstMatch' }
}

/**
 * A stable structural key for a bare leading terminal (concrete literal/regex),
 * or `null` when the term is NOT a shareable concrete prefix. Deliberately
 * conservative: only an UNWRAPPED literal (case-sensitive) or regex qualifies —
 * a transform/label/ref/node wrapper changes the value/capture shape and must NOT
 * be left-factored (correctness first; the strategy is opt-in like greedyClassify).
 */
function bareLeadingTermKey(term: Combinator<unknown>): string | null {
  const d = term._def
  if (d.tag === 'literal' && !d.caseInsensitive) return `L:${d.value}`
  // `JSON.stringify` of the pair, not a delimiter join. `source` is arbitrary regex text
  // and can contain any character, so a collision here would judge two structurally
  // DIFFERENT regexes identical and left-factor arms that must not share a prefix — a
  // correctness bug in the generated parser, not a missed optimization. A JSON array is
  // injective by construction, so no argument about which delimiter cannot occur in a
  // regex source is needed. It is also printable: this was a RAW 0x00 byte, which made
  // the whole file binary to `git diff` and invisible to `grep -rn` (see
  // `scripts/check-control-bytes.mjs`). Runs once per choice when its strategy is chosen,
  // never per match.
  if (d.tag === 'regex') return `R:${JSON.stringify([d.source, d.flags])}`
  // Two lead shapes look eligible and are deliberately NOT. Both were tried and
  // measured; see test/unit/routed-fallback.test.ts, which pins the exclusions.
  //
  // `routed()` — a bare routed() IS a single-leaf terminal (it reads the
  // dispatch-consumed token, pushes exactly ONE leaf, skips no trivia, runs no
  // sub-parse), so it is replay-safe and admitting it here is correct. It is
  // excluded because it does not PAY. The strategy trades a duplicated lead scan for
  // a prescan plus a prefix-matched flag; routed()'s emission is a context read and
  // one comparison, so there is nothing worth factoring out. Measured on
  // `dispatch(name, when(k, choice(sequence(routed(), tail_i)...)))`, emitted bytes
  // with the strategy MINUS without it: +242 (2 arms), +260 (4), +296 (8) — it costs
  // more at every width and never crosses over. The same measurement on the
  // regex-lead shape the strategy was built for: -468 (2 arms), -1597 (4), -3846 (8).
  // A routed lead also shares only the ROUTED token; when the expensive shared work
  // is the next term (a prelude), that term is what needs factoring, not this one.
  //
  // `routed(fallback)` — not a leaf at all: on the fallback path it runs a whole
  // sub-parse, so replaying it as a once-recognized leaf would be wrong, and two
  // routed() with DIFFERENT fallbacks are not interchangeable to begin with.
  //
  // A `lazy` ref lead — excluded for correctness, not for want of a key. The
  // once-only prescan runs with `ctx.capturing = false`, which suppresses capture
  // only for code emitted HERE (emitLeafCapture, codegen.ts). A ref compiles to a
  // call into a function body generated under its OWN ctx (emitLazy), so its
  // captures and trivia writes would happen during the prescan and then AGAIN in
  // each arm. Replaying a ref needs recorded-and-spliced capture state, not the
  // variable reuse a terminal replay gets away with.
  return null
}

/**
 * Detect the pure single-group shared-prefix shape: at least two arms, EVERY arm
 * a bare `sequence` with ≥2 terms, and every arm's FIRST term the same concrete
 * literal/regex (identical source/flags or literal string). Returns the plan with
 * a representative prefix combinator (the first arm's leading term — all members'
 * leading terms are structurally identical, so any one emits identical code),
 * else `null`. Mixed / multi-group / wrapped shapes are conservatively skipped and
 * fall through to `firstMatch`.
 */
function detectSharedPrefix(parsers: Combinator<unknown>[]): ChoiceStrategy | null {
  if (parsers.length < 2) return null
  let key: string | null = null
  let prefix: Combinator<unknown> | null = null
  const members: number[] = []
  for (let i = 0; i < parsers.length; i++) {
    const term0 = leadingTermOfArm(parsers[i]!)
    if (term0 === null) return null
    const k = bareLeadingTermKey(term0)
    if (k === null) return null
    if (key === null) { key = k; prefix = term0 }
    else if (k !== key) return null
    members.push(i)
  }
  if (prefix === null || members.length < 2) return null
  return { tag: 'sharedPrefix', prefix, members }
}

/**
 * Peel wrappers that neither consume input nor skip trivia before their inner
 * sequence -- `node`, `parser`/`grammar`, `transform`, `label` -- down to the core
 * `sequence` and return its FIRST term (the candidate shared prefix), or null when
 * the arm is not a wrapped-or-bare sequence with >=2 terms. Because none of these
 * wrappers advance the position before the sequence's first term, the term is parsed
 * at the arm's entry position in every arm, so the once-recognized prefix can be
 * replayed there. `attempt`/`optional`/`many`/`choice`/etc. are NOT peeled, so such
 * arms are conservatively excluded.
 */
export function leadingTermOfArm(arm: Combinator<unknown>): Combinator<unknown> | null {
  let d = arm._def
  for (;;) {
    if (d.tag === 'node' || d.tag === 'grammar' || d.tag === 'transform' || d.tag === 'label') {
      d = (d as { parser: Combinator<unknown> }).parser._def
      continue
    }
    break
  }
  if (d.tag !== 'sequence' || d.parsers.length < 2) return null
  return d.parsers[0]!
}

// ---------------------------------------------------------------------------
// Auto-not analysis (firstMatch fallback only)
// ---------------------------------------------------------------------------

function computeAutoNot(parsers: Combinator<unknown>[]): (AutoNotCheck[] | null)[] {
  return parsers.map((p, i) => {
    const litVal = getCoreLiteralValue(p)
    if (litVal === null) return null

    const checks: AutoNotCheck[] = []
    for (let j = i + 1; j < parsers.length; j++) {
      const other = parsers[j]!
      const otherLit = getCoreLiteralValue(other)
      if (otherLit !== null && otherLit.startsWith(litVal) && otherLit.length > litVal.length) {
        checks.push({ kind: 'startsWith', value: otherLit.slice(litVal.length) })
        continue
      }
      const regexDef = getCoreRegexDef(other)
      if (regexDef !== null) {
        const contSet = continuationFirstSet(litVal, regexDef.source, regexDef.flags)
        if (contSet !== null) checks.push({ kind: 'firstSet', set: contSet })
      }
    }
    return checks.length > 0 ? checks : null
  })
}

function autoNotFires(input: string, end: number, checks: AutoNotCheck[]): boolean {
  for (const check of checks) {
    if (check.kind === 'firstSet') {
      const code = end < input.length ? (input.codePointAt(end) ?? -1) : -1
      if (inFirstSet(code, check.set)) return true
    } else {
      if (input.startsWith(check.value, end)) return true
    }
  }
  return false
}

// ---------------------------------------------------------------------------
// Helpers: unwrap transform/sequence wrappers
// ---------------------------------------------------------------------------

/** Walk transform wrappers to find an inner literal's string value. */
export function getCoreLiteralValue(p: Combinator<unknown>): string | null {
  const def = p._def
  if (def.tag === 'literal' && !def.caseInsensitive) return def.value
  if (def.tag === 'transform') return getCoreLiteralValue(def.parser)
  return null
}

/** Walk transform wrappers to find an inner regex's source/flags. */
export function getCoreRegexDef(p: Combinator<unknown>): { source: string; flags: string } | null {
  const def = p._def
  if (def.tag === 'regex') return { source: def.source, flags: def.flags }
  if (def.tag === 'transform') return getCoreRegexDef(def.parser)
  if (def.tag === 'label') return getCoreRegexDef(def.parser)
  return null
}

/**
 * Apply a parser's transform chain to an already-known value, without re-parsing.
 * Used by greedyClassify to avoid a second parse call for the winning literal arm.
 */
function applyTransforms(p: Combinator<unknown>, value: unknown, span: { start: number; end: number }): unknown {
  const def = p._def
  if (def.tag === 'transform') {
    const inner = applyTransforms(def.parser, value, span)
    return def.fn(inner, span)
  }
  return value
}

// ---------------------------------------------------------------------------
// Continuation first-set (for firstMatch auto-not analysis)
// ---------------------------------------------------------------------------

function continuationFirstSet(lit: string, source: string, flags: string): FirstSet | null {
  const re = new RegExp(source, 'y' + flags.replace(/[gy]/g, ''))
  re.lastIndex = 0
  const base = re.exec(lit)
  if (!base || base[0] !== lit) return null

  const contCodes: number[] = []
  for (let code = 1; code < 128; code++) {
    re.lastIndex = 0
    const m = re.exec(lit + String.fromCharCode(code))
    if (m && m[0].length > lit.length) contCodes.push(code)
  }
  if (contCodes.length === 0) return null
  return codesToFirstSet(contCodes)
}

function codesToFirstSet(codes: number[]): FirstSet {
  codes.sort((a, b) => a - b)
  const ranges: CharRange[] = []
  let lo = codes[0]!, hi = codes[0]!
  for (let i = 1; i < codes.length; i++) {
    if (codes[i] === hi + 1) { hi = codes[i]! }
    else { ranges.push({ lo, hi }); lo = hi = codes[i]! }
  }
  ranges.push({ lo, hi })
  return { kind: 'ranges', ranges }
}

function inFirstSet(code: number, fs: FirstSet): boolean {
  if (fs.kind === 'any') return true
  if (fs.kind === 'empty') return false
  for (const r of fs.ranges) if (code >= r.lo && code <= r.hi) return true
  return false
}

function areDisjoint(sets: FirstSet[]): boolean {
  if (sets.some(s => s.kind === 'any')) return false
  for (let i = 0; i < sets.length; i++)
    for (let j = i + 1; j < sets.length; j++)
      if (intersects(sets[i]!, sets[j]!)) return false
  return true
}

// ASCII first-char → arm INDEX (not the parser) so a gated arm's gate can be
// looked up (gates[idx]) at dispatch time. -1 means "no arm keys off this char".
function buildAsciiDispatch(parsers: Combinator<unknown>[]): number[] {
  const table = Array<number>(128).fill(-1)
  for (let i = 0; i < parsers.length; i++) {
    const fs = parsers[i]!._meta.firstSet
    if (fs.kind !== 'ranges') continue
    for (const { lo, hi } of fs.ranges) {
      for (let code = Math.max(0, lo); code <= Math.min(127, hi); code++) {
        table[code] = i
      }
    }
  }
  return table
}