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

import Foundation

/// Parser for one ZQL command line:
///
///     command [varlist] [if expr] [in range] [weight] [, options]
///
/// Prefix commands (`bootstrap, reps(1000): reg y x`) recurse on the text
/// after the colon. File-taking verbs (`use`, `save`, …) treat everything
/// before the first top-level comma as a raw path.
public struct ZQCommandParser: Sendable {
    public let verbTable: ZQVerbTable

    public init(verbTable: ZQVerbTable = .builtin) {
        self.verbTable = verbTable
    }

    /// Parses a standalone expression (used by the data browser's filter
    /// bar, which compiles its text to an `if` condition — CLAUDE.md §7).
    public func parseExpression(_ text: String) throws(ZQParseError) -> ZQExpression {
        let lexer = ZQLexer()
        let tokens = try lexer.tokenize(text)
        var cursor = ExpressionParser(tokens: tokens)
        let expression = try cursor.parseExpression()
        guard cursor.isKind(.endOfLine) else {
            throw ZQParseError(
                message: "unexpected \(cursor.current.kind.describe) after expression",
                column: cursor.current.column
            )
        }
        return expression
    }

    /// Parses a single line. Returns nil for blank lines and comments.
    public func parse(_ line: String) throws(ZQParseError) -> ZQCommand? {
        let trimmed = line.trimmingCharacters(in: .whitespaces)
        if trimmed.isEmpty || trimmed.hasPrefix("//") || trimmed.hasPrefix("*") {
            return nil
        }

        // Extract the verb word to decide the parsing strategy.
        let verbWord = String(trimmed.prefix { !$0.isWhitespace && $0 != "," })
        let verbColumn = columnOf(word: verbWord, in: line)
        guard let verb = verbTable.resolve(verbWord.lowercased()) else {
            throw ZQParseError(
                message: "unknown command '\(verbWord)'",
                column: verbColumn,
                suggestion: verbTable.closest(to: verbWord.lowercased())
            )
        }

        if verbTable.prefixVerbs.contains(verb) {
            return try parsePrefixCommand(verb: verb, line: line, trimmed: trimmed)
        }
        if verbTable.fileVerbs.contains(verb) {
            return try parseFileCommand(verb: verb, trimmed: trimmed)
        }
        return try parseStandardCommand(verb: verb, line: line, trimmed: trimmed)
    }

    // MARK: - Prefix commands

    private func parsePrefixCommand(
        verb: String, line: String, trimmed: String
    ) throws(ZQParseError) -> ZQCommand {
        // Split on the first top-level colon: `bootstrap, reps(1000): reg y x`
        guard let colonIndex = topLevelColonIndex(in: trimmed) else {
            throw ZQParseError(
                message: "prefix command '\(verb)' requires ': command'",
                column: trimmed.count + 1
            )
        }
        let head = String(trimmed[trimmed.startIndex..<colonIndex])
        let tail = String(trimmed[trimmed.index(after: colonIndex)...])
            .trimmingCharacters(in: .whitespaces)

        var command = try parseStandardCommand(verb: verb, line: head, trimmed: head)
        guard let body = try parse(tail) else {
            throw ZQParseError(
                message: "prefix command '\(verb)' has an empty body",
                column: trimmed.count + 1
            )
        }
        command.body = body
        return command
    }

    /// Index of the first colon that is not nested inside parentheses,
    /// brackets, or a string literal.
    private func topLevelColonIndex(in text: String) -> String.Index? {
        var depth = 0
        var inString = false
        var index = text.startIndex
        while index < text.endIndex {
            let c = text[index]
            if c == "\"" { inString.toggle() }
            if !inString {
                switch c {
                case "(", "[": depth += 1
                case ")", "]": depth -= 1
                case ":" where depth == 0: return index
                default: break
                }
            }
            index = text.index(after: index)
        }
        return nil
    }

    // MARK: - File commands

    private func parseFileCommand(
        verb: String, trimmed: String
    ) throws(ZQParseError) -> ZQCommand {
        // Everything after the verb and before the first top-level comma is
        // a raw path (quoted paths keep spaces); options follow the comma.
        let afterVerb = trimmed.drop { !$0.isWhitespace }
        var path = String(afterVerb).trimmingCharacters(in: .whitespaces)
        var options: [ZQOption] = []

        if let commaIndex = topLevelCommaIndex(in: path) {
            let optionText = String(path[path.index(after: commaIndex)...])
            path = String(path[path.startIndex..<commaIndex])
                .trimmingCharacters(in: .whitespaces)
            let lexer = ZQLexer()
            let tokens = try lexer.tokenize(optionText)
            var cursor = ExpressionParser(tokens: tokens)
            options = try parseOptions(&cursor)
        }

        // `use file` / `import file` accept an optional leading `using`.
        for keyword in ["using ", "using\t"] where path.hasPrefix(keyword) {
            path = String(path.dropFirst(keyword.count))
                .trimmingCharacters(in: .whitespaces)
        }
        if path.hasPrefix("\""), path.hasSuffix("\""), path.count >= 2 {
            path = String(path.dropFirst().dropLast())
        }

        return ZQCommand(
            verb: verb,
            argument: path.isEmpty ? nil : path,
            options: options
        )
    }

    private func topLevelCommaIndex(in text: String) -> String.Index? {
        var depth = 0
        var inString = false
        var index = text.startIndex
        while index < text.endIndex {
            let c = text[index]
            if c == "\"" { inString.toggle() }
            if !inString {
                switch c {
                case "(", "[": depth += 1
                case ")", "]": depth -= 1
                case "," where depth == 0: return index
                default: break
                }
            }
            index = text.index(after: index)
        }
        return nil
    }

    // MARK: - Standard commands

    private func parseStandardCommand(
        verb: String, line: String, trimmed: String
    ) throws(ZQParseError) -> ZQCommand {
        let lexer = ZQLexer()
        let tokens = try lexer.tokenize(trimmed)
        var cursor = ExpressionParser(tokens: tokens)

        // Consume the (already resolved) verb token.
        _ = cursor.advance()
        var command = ZQCommand(verb: verb)

        // Compound verbs: `graph scatter y x`, `ivregress 2sls …`. A
        // digit-led sub-command like `2sls` lexes as number+identifier;
        // reassemble when the tokens are column-adjacent.
        if verbTable.compoundVerbs.contains(verb) {
            switch cursor.current.kind {
            case .identifier(let sub):
                _ = cursor.advance()
                command.subverb = sub
            case .number(let value) where value == value.rounded() && value >= 0:
                let numberToken = cursor.advance()
                let prefix = String(Int(value))
                guard case .identifier(let rest) = cursor.current.kind,
                      cursor.current.column == numberToken.column + prefix.count else {
                    throw ZQParseError(
                        message: "'\(verb)' requires a sub-command",
                        column: numberToken.column
                    )
                }
                _ = cursor.advance()
                command.subverb = prefix + rest
            default:
                throw ZQParseError(
                    message: "'\(verb)' requires a sub-command",
                    column: cursor.current.column
                )
            }
        }

        // `set seed 42` → subverb "seed", argument "42"
        if verb == "set" {
            guard case .identifier(let sub) = cursor.current.kind else {
                throw ZQParseError(
                    message: "'set' requires a parameter name",
                    column: cursor.current.column
                )
            }
            _ = cursor.advance()
            command.subverb = sub
            if case .number(let value) = cursor.current.kind {
                _ = cursor.advance()
                command.argument = renderNumber(value)
            } else if case .identifier(let word) = cursor.current.kind {
                _ = cursor.advance()
                command.argument = word
            }
            try expectEnd(&cursor)
            return command
        }

        // `display 2 + 2` — a bare expression command.
        if verb == "display" {
            let expression = try cursor.parseExpression()
            command.assignment = ZQAssignment(target: "", expression: expression)
            try expectEnd(&cursor)
            return command
        }

        // Assignment verbs: `gen log_rev = ln(revenue)`
        if verbTable.assignmentVerbs.contains(verb) {
            guard case .identifier(let target) = cursor.current.kind else {
                throw ZQParseError(
                    message: "'\(verb)' requires a variable name",
                    column: cursor.current.column
                )
            }
            _ = cursor.advance()
            try cursor.expect(.op("="))
            let expression = try cursor.parseExpression()
            command.assignment = ZQAssignment(target: target, expression: expression)
        } else {
            let (varlist, ivSpec) = try parseVarlist(&cursor)
            command.varlist = varlist
            command.ivSpec = ivSpec
        }

        // Qualifiers may appear in any sensible order; Stata fixes the
        // order as if → in → weight, which we enforce implicitly.
        if consumeKeyword(&cursor, "if") {
            command.condition = try cursor.parseExpression()
        }
        if consumeKeyword(&cursor, "in") {
            command.range = try parseRange(&cursor)
        }
        if cursor.isKind(.lbracket) {
            command.weight = try parseWeight(&cursor)
        }
        if cursor.consumeIfKind(.comma) {
            command.options = try parseOptions(&cursor)
        }
        try expectEnd(&cursor)
        return command
    }

    // MARK: - Varlist

    private func parseVarlist(
        _ cursor: inout ExpressionParser
    ) throws(ZQParseError) -> ([ZQVarSpec], ZQIVSpec?) {
        var specs: [ZQVarSpec] = []
        var ivSpec: ZQIVSpec?
        loop: while true {
            switch cursor.current.kind {
            case .identifier(let name):
                if name == "if" || name == "in" { break loop }
                specs.append(try parseVarSpec(&cursor))
            case .lparen:
                // IV group: (endog… = instruments…)
                guard ivSpec == nil else {
                    throw ZQParseError(
                        message: "only one (endogenous = instruments) group is allowed",
                        column: cursor.current.column
                    )
                }
                ivSpec = try parseIVGroup(&cursor)
            default:
                break loop
            }
        }
        return (specs, ivSpec)
    }

    private func parseIVGroup(
        _ cursor: inout ExpressionParser
    ) throws(ZQParseError) -> ZQIVSpec {
        try cursor.expect(.lparen)
        var endogenous: [String] = []
        while case .identifier(let name) = cursor.current.kind {
            endogenous.append(name)
            _ = cursor.advance()
        }
        guard !endogenous.isEmpty else {
            throw ZQParseError(
                message: "expected endogenous variable names before '='",
                column: cursor.current.column
            )
        }
        try cursor.expect(.op("="))
        var instruments: [String] = []
        while case .identifier(let name) = cursor.current.kind {
            instruments.append(name)
            _ = cursor.advance()
        }
        guard !instruments.isEmpty else {
            throw ZQParseError(
                message: "expected instrument names after '='",
                column: cursor.current.column
            )
        }
        try cursor.expect(.rparen)
        return ZQIVSpec(endogenous: endogenous, instruments: instruments)
    }

    private func parseVarSpec(
        _ cursor: inout ExpressionParser
    ) throws(ZQParseError) -> ZQVarSpec {
        var terms: [ZQVarSpec] = [try parseFactorTerm(&cursor)]
        var full = false
        while case .op(let symbol) = cursor.current.kind, symbol == "#" || symbol == "##" {
            full = full || symbol == "##"
            _ = cursor.advance()
            terms.append(try parseFactorTerm(&cursor))
        }
        return terms.count == 1 ? terms[0] : .interaction(terms, full: full)
    }

    private func parseFactorTerm(
        _ cursor: inout ExpressionParser
    ) throws(ZQParseError) -> ZQVarSpec {
        guard case .identifier(let first) = cursor.current.kind else {
            throw ZQParseError(
                message: "expected a variable name, found \(cursor.current.kind.describe)",
                column: cursor.current.column
            )
        }
        _ = cursor.advance()
        if cursor.isKind(.op(".")) {
            _ = cursor.advance()
            guard case .identifier(let name) = cursor.current.kind else {
                throw ZQParseError(
                    message: "expected a variable name after '\(first).'",
                    column: cursor.current.column
                )
            }
            _ = cursor.advance()
            return .factor(op: first, name: name)
        }
        return .simple(first)
    }

    // MARK: - Qualifiers

    private func consumeKeyword(_ cursor: inout ExpressionParser, _ keyword: String) -> Bool {
        if case .identifier(let word) = cursor.current.kind, word == keyword {
            _ = cursor.advance()
            return true
        }
        return false
    }

    private func parseRange(
        _ cursor: inout ExpressionParser
    ) throws(ZQParseError) -> ZQRange {
        let lower = try parseRangeBound(&cursor)
        try cursor.expect(.op("/"))
        let upper = try parseRangeBound(&cursor)
        return ZQRange(lower: lower, upper: upper)
    }

    private func parseRangeBound(
        _ cursor: inout ExpressionParser
    ) throws(ZQParseError) -> Int {
        var sign = 1
        if cursor.consumeIfKind(.op("-")) { sign = -1 }
        // `f` and `l` mean first and last observation.
        if case .identifier(let word) = cursor.current.kind {
            if word == "f" { _ = cursor.advance(); return 1 }
            if word == "l" { _ = cursor.advance(); return -1 }
        }
        guard case .number(let value) = cursor.current.kind,
              value == value.rounded() else {
            throw ZQParseError(
                message: "expected an integer observation number",
                column: cursor.current.column
            )
        }
        _ = cursor.advance()
        return sign * Int(value)
    }

    private func parseWeight(
        _ cursor: inout ExpressionParser
    ) throws(ZQParseError) -> ZQWeight {
        try cursor.expect(.lbracket)
        guard case .identifier(let kindWord) = cursor.current.kind else {
            throw ZQParseError(
                message: "expected a weight type (aweight, fweight, pweight, iweight)",
                column: cursor.current.column
            )
        }
        let normalized: ZQWeight.Kind? = switch kindWord {
        case "aweight", "aw", "w", "weight": .aweight
        case "fweight", "fw": .fweight
        case "pweight", "pw": .pweight
        case "iweight", "iw": .iweight
        default: nil
        }
        guard let kind = normalized else {
            throw ZQParseError(
                message: "unknown weight type '\(kindWord)'",
                column: cursor.current.column
            )
        }
        _ = cursor.advance()
        try cursor.expect(.op("="))
        let expression = try cursor.parseExpression()
        try cursor.expect(.rbracket)
        return ZQWeight(kind: kind, expression: expression)
    }

    // MARK: - Options

    private func parseOptions(
        _ cursor: inout ExpressionParser
    ) throws(ZQParseError) -> [ZQOption] {
        var options: [ZQOption] = []
        while case .identifier(let name) = cursor.current.kind {
            _ = cursor.advance()
            var arguments: [String] = []
            if cursor.consumeIfKind(.lparen) {
                var depth = 1
                var pieces: [String] = []
                while depth > 0 {
                    let token = cursor.current
                    switch token.kind {
                    case .lparen:
                        depth += 1
                        pieces.append("(")
                    case .rparen:
                        depth -= 1
                        if depth > 0 { pieces.append(")") }
                    case .endOfLine:
                        throw ZQParseError(
                            message: "unclosed '(' in option '\(name)'",
                            column: token.column
                        )
                    case .identifier(let word):
                        pieces.append(word)
                    case .number(let value):
                        pieces.append(renderNumber(value))
                    case .string(let text):
                        pieces.append("\"\(text)\"")
                    case .comma:
                        pieces.append(",")
                    case .colon:
                        pieces.append(":")
                    case .lbracket:
                        pieces.append("[")
                    case .rbracket:
                        pieces.append("]")
                    case .op(let symbol):
                        pieces.append(symbol)
                    }
                    _ = cursor.advance()
                }
                // Arguments split on top-level spaces in the original text
                // are reassembled here token-by-token: identifiers and
                // numbers stand alone, operator glue joins with neighbors.
                arguments = assembleOptionArguments(pieces)
            }
            options.append(ZQOption(name: name, arguments: arguments))
            // Optional commas between options are tolerated.
            _ = cursor.consumeIfKind(.comma)
        }
        return options
    }

    /// Joins raw option-argument pieces into logical arguments: glue
    /// operators (`.`, `=`, comparison, arithmetic) bind to their
    /// neighbors, and standalone words split on whitespace boundaries.
    private func assembleOptionArguments(_ pieces: [String]) -> [String] {
        let glue: Set<String> = [
            ".", "=", "==", "!=", "<", "<=", ">", ">=",
            "+", "-", "*", "/", "^", "#", "##", "&", "|", "!",
        ]
        var arguments: [String] = []
        var pending = ""
        var previousWasGlue = false

        for piece in pieces {
            if piece == "," {
                if !pending.isEmpty { arguments.append(pending); pending = "" }
                previousWasGlue = false
                continue
            }
            if glue.contains(piece) {
                pending += piece
                previousWasGlue = true
            } else if previousWasGlue || pending.isEmpty {
                pending += piece
                previousWasGlue = false
            } else {
                arguments.append(pending)
                pending = piece
                previousWasGlue = false
            }
        }
        if !pending.isEmpty { arguments.append(pending) }
        return arguments
    }

    // MARK: - Helpers

    private func expectEnd(_ cursor: inout ExpressionParser) throws(ZQParseError) {
        guard cursor.isKind(.endOfLine) else {
            throw ZQParseError(
                message: "unexpected \(cursor.current.kind.describe) after end of command",
                column: cursor.current.column
            )
        }
    }

    private func renderNumber(_ value: Double) -> String {
        if value == value.rounded(), abs(value) < 1e15 {
            return String(Int64(value))
        }
        return String(value)
    }

    private func columnOf(word: String, in line: String) -> Int {
        guard let range = line.range(of: word) else { return 1 }
        return line.distance(from: line.startIndex, to: range.lowerBound) + 1
    }
}