//
//  AST.swift
//  Metrika
//
//  Author:  Simon-Pierre Boucher
//  Contact: contact@spboucher.ai
//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.
//

/// A variable specification appearing in a varlist. Factor notation is
/// preserved verbatim; expansion into indicator columns happens at plan
/// time (CLAUDE.md §4), never at parse time.
public indirect enum ZQVarSpec: Equatable, Sendable {
    /// Plain variable reference: `price`
    case simple(String)
    /// Factor-operator variable: `i.region`, `c.age`, `ib2.group`
    case factor(op: String, name: String)
    /// Interaction: `c.age#c.age`, `i.region#c.price`
    case interaction([ZQVarSpec], full: Bool)

    /// The underlying dataset column names this spec references.
    public var referencedNames: [String] {
        switch self {
        case .simple(let n): return [n]
        case .factor(_, let n): return [n]
        case .interaction(let parts, _): return parts.flatMap(\.referencedNames)
        }
    }
}

/// Expression AST used for `if` qualifiers, `gen`/`replace` right-hand
/// sides, and option arguments.
public indirect enum ZQExpression: Equatable, Sendable {
    case number(Double)
    case string(String)
    case variable(String)
    /// Stata-style missing-value literal `.`
    case missing
    case unary(op: String, ZQExpression)
    case binary(op: String, ZQExpression, ZQExpression)
    case call(name: String, arguments: [ZQExpression])
}

/// Observation range qualifier: `in 5/200`. Negative values count from the
/// end (`in -10/-1`), 1-based like Stata.
public struct ZQRange: Equatable, Sendable {
    public var lower: Int
    public var upper: Int

    public init(lower: Int, upper: Int) {
        self.lower = lower
        self.upper = upper
    }
}

/// Weight qualifier: `[aweight = expr]`.
public struct ZQWeight: Equatable, Sendable {
    public enum Kind: String, Equatable, Sendable {
        case aweight, fweight, pweight, iweight
    }

    public var kind: Kind
    public var expression: ZQExpression

    public init(kind: Kind, expression: ZQExpression) {
        self.kind = kind
        self.expression = expression
    }
}

/// A trailing option: bare flag (`robust`, `detail`) or parenthesized
/// (`cluster(id)`, `reps(100000)`, `by(region year)`).
public struct ZQOption: Equatable, Sendable {
    public var name: String
    /// Raw argument tokens re-rendered as strings; interpretation is
    /// command-specific and happens at plan/execution time.
    public var arguments: [String]

    public init(name: String, arguments: [String] = []) {
        self.name = name
        self.arguments = arguments
    }

    public var firstArgument: String? { arguments.first }
}

/// Instrumental-variables group in a varlist:
/// `ivregress 2sls y exog (endog1 endog2 = instr1 instr2)`.
public struct ZQIVSpec: Equatable, Sendable {
    public var endogenous: [String]
    public var instruments: [String]

    public init(endogenous: [String], instruments: [String]) {
        self.endogenous = endogenous
        self.instruments = instruments
    }
}

/// Assignment payload for `gen`/`replace`-style commands:
/// `gen log_rev = ln(revenue)`.
public struct ZQAssignment: Equatable, Sendable {
    public var target: String
    public var expression: ZQExpression

    public init(target: String, expression: ZQExpression) {
        self.target = target
        self.expression = expression
    }
}

/// A fully parsed ZQL command:
/// `command [varlist] [if expr] [in range] [weight] [, options]`,
/// optionally wrapped by a prefix command (`bootstrap, reps(1000): reg …`).
public struct ZQCommand: Equatable, Sendable {
    /// Canonical verb (abbreviations already resolved: `reg` → `regress`).
    public var verb: String
    /// Sub-verb for compound commands (`graph scatter …` → "scatter").
    public var subverb: String?
    public var varlist: [ZQVarSpec]
    /// Instrumented group for IV estimators, parsed from
    /// `(endog… = instruments…)` inside the varlist.
    public var ivSpec: ZQIVSpec?
    /// `gen`/`replace` assignment, mutually exclusive with a plain varlist.
    public var assignment: ZQAssignment?
    /// Raw argument such as a file path (`use sales.parquet`).
    public var argument: String?
    public var condition: ZQExpression?
    public var range: ZQRange?
    public var weight: ZQWeight?
    public var options: [ZQOption]
    /// Inner command of a prefix construct: for
    /// `bootstrap, reps(1000): reg y x`, the outer command is `bootstrap`
    /// and `body` holds the parsed `regress` command.
    public var body: ZQCommand? {
        get { bodyBox?.value }
        set { bodyBox = newValue.map(Box.init) }
    }

    private var bodyBox: Box<ZQCommand>?

    public init(
        verb: String,
        subverb: String? = nil,
        varlist: [ZQVarSpec] = [],
        ivSpec: ZQIVSpec? = nil,
        assignment: ZQAssignment? = nil,
        argument: String? = nil,
        condition: ZQExpression? = nil,
        range: ZQRange? = nil,
        weight: ZQWeight? = nil,
        options: [ZQOption] = [],
        body: ZQCommand? = nil
    ) {
        self.verb = verb
        self.subverb = subverb
        self.varlist = varlist
        self.ivSpec = ivSpec
        self.assignment = assignment
        self.argument = argument
        self.condition = condition
        self.range = range
        self.weight = weight
        self.options = options
        self.bodyBox = body.map(Box.init)
    }

    public func option(_ name: String) -> ZQOption? {
        options.first { $0.name == name }
    }

    public func hasOption(_ name: String) -> Bool {
        option(name) != nil
    }
}

/// Reference box that lets `ZQCommand` nest itself as a prefix body.
private final class Box<T>: @unchecked Sendable {
    let value: T
    init(_ value: T) { self.value = value }
}

extension Box: Equatable where T: Equatable {
    static func == (lhs: Box<T>, rhs: Box<T>) -> Bool { lhs.value == rhs.value }
}