use std::collections::BTreeMap;
use std::fmt;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ValueType {
    Number,
    Boolean,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub(crate) enum Value {
    Number(f64),
    Boolean(bool),
}

impl Value {
    fn value_type(self) -> ValueType {
        match self {
            Self::Number(_) => ValueType::Number,
            Self::Boolean(_) => ValueType::Boolean,
        }
    }

    pub(crate) fn as_number(self) -> Option<f64> {
        match self {
            Self::Number(value) => Some(value),
            Self::Boolean(_) => None,
        }
    }

    pub(crate) fn as_bool(self) -> Option<bool> {
        match self {
            Self::Boolean(value) => Some(value),
            Self::Number(_) => None,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct CompileError {
    position: usize,
    message: String,
}

impl CompileError {
    fn new(position: usize, message: impl Into<String>) -> Self {
        Self {
            position,
            message: message.into(),
        }
    }
}

impl fmt::Display for CompileError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            formatter,
            "numeric expression compile error at byte {}: {}",
            self.position, self.message
        )
    }
}

impl std::error::Error for CompileError {}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct EvalError {
    message: String,
}

impl EvalError {
    pub(crate) fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
        }
    }
}

impl fmt::Display for EvalError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "numeric expression eval error: {}", self.message)
    }
}

impl std::error::Error for EvalError {}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum UnaryOp {
    Negate,
    Not,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum BinaryOp {
    Add,
    Subtract,
    Multiply,
    Divide,
    Remainder,
    Equal,
    NotEqual,
    Less,
    LessEqual,
    Greater,
    GreaterEqual,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Builtin {
    Round,
    Floor,
    Ceil,
    Abs,
    Min,
    Max,
    Sqrt,
    Pow,
    Log,
    Exp,
    Clamp,
    Between,
    Nz,
    SafeDiv,
    Iff,
}

#[derive(Debug, Clone, PartialEq)]
enum Instruction {
    Push(Value),
    LoadVariable(u16),
    LoadLocal(u16),
    StoreLocal(u16),
    Unary(UnaryOp),
    Binary(BinaryOp),
    Call { builtin: Builtin, argc: u8 },
    JumpIfFalse(usize),
    JumpIfTrue(usize),
    Jump(usize),
    Return,
}

#[derive(Debug, Clone)]
pub(crate) struct Program {
    instructions: Vec<Instruction>,
    variables: Vec<String>,
    variable_types: Vec<ValueType>,
    local_count: usize,
    result_type: ValueType,
}

impl Program {
    pub(crate) fn variables(&self) -> &[String] {
        &self.variables
    }

    pub(crate) fn evaluate<F>(
        &self,
        scratch: &mut Scratch,
        mut resolve: F,
    ) -> Result<Value, EvalError>
    where
        F: FnMut(usize, &str, ValueType) -> Result<Value, EvalError>,
    {
        scratch.prepare(self);
        let mut pc = 0usize;
        while let Some(instruction) = self.instructions.get(pc) {
            match *instruction {
                Instruction::Push(value) => scratch.stack.push(value),
                Instruction::LoadVariable(index) => {
                    let index = usize::from(index);
                    let cached = scratch.variables[index];
                    let value = match cached {
                        Some(value) => value,
                        None => {
                            let expected_type = self.variable_types[index];
                            let value = resolve(index, &self.variables[index], expected_type)?;
                            if value.value_type() != expected_type {
                                return Err(EvalError::new(format!(
                                    "variable {} expected {:?}, got {:?}",
                                    self.variables[index],
                                    expected_type,
                                    value.value_type()
                                )));
                            }
                            scratch.variables[index] = Some(value);
                            value
                        }
                    };
                    scratch.stack.push(value);
                }
                Instruction::LoadLocal(index) => {
                    let index = usize::from(index);
                    let value = scratch.locals[index].ok_or_else(|| {
                        EvalError::new(format!("local slot {index} was not initialized"))
                    })?;
                    scratch.stack.push(value);
                }
                Instruction::StoreLocal(index) => {
                    let value = pop(&mut scratch.stack)?;
                    scratch.locals[usize::from(index)] = Some(value);
                }
                Instruction::Unary(operator) => {
                    let value = pop(&mut scratch.stack)?;
                    scratch.stack.push(eval_unary(operator, value)?);
                }
                Instruction::Binary(operator) => {
                    let rhs = pop(&mut scratch.stack)?;
                    let lhs = pop(&mut scratch.stack)?;
                    scratch.stack.push(eval_binary(operator, lhs, rhs)?);
                }
                Instruction::Call { builtin, argc } => {
                    let argc = usize::from(argc);
                    if scratch.stack.len() < argc {
                        return Err(EvalError::new("stack underflow during builtin call"));
                    }
                    let start = scratch.stack.len() - argc;
                    let value = eval_builtin(builtin, &scratch.stack[start..])?;
                    scratch.stack.truncate(start);
                    scratch.stack.push(value);
                }
                Instruction::JumpIfFalse(target) => {
                    let condition = pop_bool(&mut scratch.stack)?;
                    if !condition {
                        pc = target;
                        continue;
                    }
                }
                Instruction::JumpIfTrue(target) => {
                    let condition = pop_bool(&mut scratch.stack)?;
                    if condition {
                        pc = target;
                        continue;
                    }
                }
                Instruction::Jump(target) => {
                    pc = target;
                    continue;
                }
                Instruction::Return => {
                    let result = pop(&mut scratch.stack)?;
                    if result.value_type() != self.result_type {
                        return Err(EvalError::new(format!(
                            "program expected {:?}, got {:?}",
                            self.result_type,
                            result.value_type()
                        )));
                    }
                    if !scratch.stack.is_empty() {
                        return Err(EvalError::new("program returned with a dirty stack"));
                    }
                    return Ok(result);
                }
            }
            pc += 1;
        }
        Err(EvalError::new("program terminated without return"))
    }
}

#[derive(Debug, Default)]
pub(crate) struct Scratch {
    stack: Vec<Value>,
    variables: Vec<Option<Value>>,
    locals: Vec<Option<Value>>,
}

impl Scratch {
    fn prepare(&mut self, program: &Program) {
        self.stack.clear();
        self.stack.reserve(
            program
                .instructions
                .len()
                .saturating_sub(self.stack.capacity()),
        );
        self.variables.clear();
        self.variables.resize(program.variables.len(), None);
        self.locals.clear();
        self.locals.resize(program.local_count, None);
    }
}

fn pop(stack: &mut Vec<Value>) -> Result<Value, EvalError> {
    stack.pop().ok_or_else(|| EvalError::new("stack underflow"))
}

fn pop_bool(stack: &mut Vec<Value>) -> Result<bool, EvalError> {
    pop(stack)?
        .as_bool()
        .ok_or_else(|| EvalError::new("boolean operand required"))
}

fn number(value: Value) -> Result<f64, EvalError> {
    value
        .as_number()
        .ok_or_else(|| EvalError::new("numeric operand required"))
}

fn eval_unary(operator: UnaryOp, value: Value) -> Result<Value, EvalError> {
    match operator {
        UnaryOp::Negate => Ok(Value::Number(-number(value)?)),
        UnaryOp::Not => {
            Ok(Value::Boolean(!value.as_bool().ok_or_else(|| {
                EvalError::new("boolean operand required for !")
            })?))
        }
    }
}

fn eval_binary(operator: BinaryOp, lhs: Value, rhs: Value) -> Result<Value, EvalError> {
    match operator {
        BinaryOp::Add => Ok(Value::Number(number(lhs)? + number(rhs)?)),
        BinaryOp::Subtract => Ok(Value::Number(number(lhs)? - number(rhs)?)),
        BinaryOp::Multiply => Ok(Value::Number(number(lhs)? * number(rhs)?)),
        BinaryOp::Divide => Ok(Value::Number(number(lhs)? / number(rhs)?)),
        BinaryOp::Remainder => Ok(Value::Number(number(lhs)? % number(rhs)?)),
        BinaryOp::Equal => Ok(Value::Boolean(match (lhs, rhs) {
            (Value::Number(lhs), Value::Number(rhs)) => float_equal(lhs, rhs),
            (Value::Boolean(lhs), Value::Boolean(rhs)) => lhs == rhs,
            _ => {
                return Err(EvalError::new(
                    "comparison operands must have the same type",
                ));
            }
        })),
        BinaryOp::NotEqual => Ok(Value::Boolean(match (lhs, rhs) {
            (Value::Number(lhs), Value::Number(rhs)) => float_not_equal(lhs, rhs),
            (Value::Boolean(lhs), Value::Boolean(rhs)) => lhs != rhs,
            _ => {
                return Err(EvalError::new(
                    "comparison operands must have the same type",
                ));
            }
        })),
        BinaryOp::Less => {
            let (lhs, rhs) = (number(lhs)?, number(rhs)?);
            Ok(Value::Boolean(
                (rhs - lhs) / float_comparison_scale(lhs, rhs) > f64::EPSILON,
            ))
        }
        BinaryOp::LessEqual => {
            let (lhs, rhs) = (number(lhs)?, number(rhs)?);
            Ok(Value::Boolean(
                (rhs - lhs) / float_comparison_scale(lhs, rhs) > -f64::EPSILON,
            ))
        }
        BinaryOp::Greater => {
            let (lhs, rhs) = (number(lhs)?, number(rhs)?);
            Ok(Value::Boolean(
                (lhs - rhs) / float_comparison_scale(lhs, rhs) > f64::EPSILON,
            ))
        }
        BinaryOp::GreaterEqual => {
            let (lhs, rhs) = (number(lhs)?, number(rhs)?);
            Ok(Value::Boolean(
                (lhs - rhs) / float_comparison_scale(lhs, rhs) > -f64::EPSILON,
            ))
        }
    }
}

fn float_comparison_scale(lhs: f64, rhs: f64) -> f64 {
    if lhs * rhs == 0.0 {
        1.0
    } else {
        lhs.abs().max(rhs.abs())
    }
}

fn float_equal(lhs: f64, rhs: f64) -> bool {
    (lhs - rhs).abs() / float_comparison_scale(lhs, rhs) <= f64::EPSILON
}

fn float_not_equal(lhs: f64, rhs: f64) -> bool {
    (lhs - rhs).abs() / float_comparison_scale(lhs, rhs) > f64::EPSILON
}

fn eval_builtin(builtin: Builtin, args: &[Value]) -> Result<Value, EvalError> {
    let numeric = |index: usize| -> Result<f64, EvalError> {
        args.get(index)
            .copied()
            .ok_or_else(|| EvalError::new("missing builtin argument"))
            .and_then(number)
    };
    Ok(match builtin {
        Builtin::Round => Value::Number(numeric(0)?.round()),
        Builtin::Floor => Value::Number(numeric(0)?.floor()),
        Builtin::Ceil => Value::Number(numeric(0)?.ceil()),
        Builtin::Abs => Value::Number(numeric(0)?.abs()),
        Builtin::Min => Value::Number(numeric(0)?.min(numeric(1)?)),
        Builtin::Max => Value::Number(numeric(0)?.max(numeric(1)?)),
        Builtin::Sqrt => Value::Number(numeric(0)?.sqrt()),
        Builtin::Pow => Value::Number(numeric(0)?.powf(numeric(1)?)),
        Builtin::Log => Value::Number(numeric(0)?.ln()),
        Builtin::Exp => Value::Number(numeric(0)?.exp()),
        Builtin::Clamp => Value::Number(numeric(0)?.clamp(numeric(1)?, numeric(2)?)),
        Builtin::Between => {
            let value = numeric(0)?;
            Value::Boolean(value >= numeric(1)? && value <= numeric(2)?)
        }
        Builtin::Nz => {
            let value = numeric(0)?;
            Value::Number(if value.is_finite() {
                value
            } else {
                numeric(1)?
            })
        }
        Builtin::SafeDiv => {
            let lhs = numeric(0)?;
            let rhs = numeric(1)?;
            let fallback = if args.len() == 3 { numeric(2)? } else { 0.0 };
            Value::Number(if rhs.abs() <= f64::EPSILON {
                fallback
            } else {
                lhs / rhs
            })
        }
        Builtin::Iff => {
            let condition = args
                .first()
                .and_then(|value| value.as_bool())
                .ok_or_else(|| EvalError::new("iff condition must be boolean"))?;
            if condition { args[1] } else { args[2] }
        }
    })
}

#[derive(Debug, Clone, PartialEq)]
enum Expr {
    Number(f64, usize),
    Boolean(bool, usize),
    Variable(String, usize),
    Unary {
        operator: UnaryOp,
        operand: Box<Expr>,
        position: usize,
    },
    Binary {
        operator: ParsedBinaryOp,
        lhs: Box<Expr>,
        rhs: Box<Expr>,
        position: usize,
    },
    Call {
        name: String,
        args: Vec<Expr>,
        position: usize,
    },
    If {
        condition: Box<Expr>,
        when_true: Box<Expr>,
        when_false: Box<Expr>,
        position: usize,
    },
}

impl Expr {
    fn position(&self) -> usize {
        match self {
            Self::Number(_, position)
            | Self::Boolean(_, position)
            | Self::Variable(_, position) => *position,
            Self::Unary { position, .. }
            | Self::Binary { position, .. }
            | Self::Call { position, .. }
            | Self::If { position, .. } => *position,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ParsedBinaryOp {
    Add,
    Subtract,
    Multiply,
    Divide,
    Remainder,
    Equal,
    NotEqual,
    Less,
    LessEqual,
    Greater,
    GreaterEqual,
    And,
    Or,
}

#[derive(Debug, Clone, PartialEq)]
struct LetStatement {
    name: String,
    value: Expr,
    position: usize,
}

#[derive(Debug, Clone, PartialEq)]
struct ParsedProgram {
    statements: Vec<LetStatement>,
    result: Expr,
}

#[derive(Debug, Clone, PartialEq)]
struct Token {
    kind: TokenKind,
    position: usize,
}

#[derive(Debug, Clone, PartialEq)]
enum TokenKind {
    Number(f64),
    Identifier(String),
    True,
    False,
    Let,
    Const,
    If,
    Else,
    Plus,
    Minus,
    Star,
    Slash,
    Percent,
    Bang,
    AndAnd,
    OrOr,
    Equal,
    EqualEqual,
    BangEqual,
    Less,
    LessEqual,
    Greater,
    GreaterEqual,
    LeftParen,
    RightParen,
    LeftBrace,
    RightBrace,
    Comma,
    Semicolon,
    Eof,
}

struct Lexer<'a> {
    source: &'a str,
    cursor: usize,
}

impl<'a> Lexer<'a> {
    fn new(source: &'a str) -> Self {
        Self { source, cursor: 0 }
    }

    fn tokenize(mut self) -> Result<Vec<Token>, CompileError> {
        let mut tokens = Vec::new();
        loop {
            self.skip_space_and_comments()?;
            let position = self.cursor;
            let Some(ch) = self.peek() else {
                tokens.push(Token {
                    kind: TokenKind::Eof,
                    position,
                });
                return Ok(tokens);
            };
            let kind = match ch {
                '0'..='9' | '.' if self.number_starts_here() => self.number()?,
                'a'..='z' | 'A'..='Z' | '_' => self.identifier(),
                '+' => {
                    self.advance();
                    TokenKind::Plus
                }
                '-' => {
                    self.advance();
                    TokenKind::Minus
                }
                '*' => {
                    self.advance();
                    TokenKind::Star
                }
                '/' => {
                    self.advance();
                    TokenKind::Slash
                }
                '%' => {
                    self.advance();
                    TokenKind::Percent
                }
                '!' => {
                    self.advance();
                    if self.consume('=') {
                        TokenKind::BangEqual
                    } else {
                        TokenKind::Bang
                    }
                }
                '&' => {
                    self.advance();
                    if !self.consume('&') {
                        return Err(CompileError::new(position, "single & is not supported"));
                    }
                    TokenKind::AndAnd
                }
                '|' => {
                    self.advance();
                    if !self.consume('|') {
                        return Err(CompileError::new(position, "single | is not supported"));
                    }
                    TokenKind::OrOr
                }
                '=' => {
                    self.advance();
                    if self.consume('=') {
                        TokenKind::EqualEqual
                    } else {
                        TokenKind::Equal
                    }
                }
                '<' => {
                    self.advance();
                    if self.consume('=') {
                        TokenKind::LessEqual
                    } else {
                        TokenKind::Less
                    }
                }
                '>' => {
                    self.advance();
                    if self.consume('=') {
                        TokenKind::GreaterEqual
                    } else {
                        TokenKind::Greater
                    }
                }
                '(' => {
                    self.advance();
                    TokenKind::LeftParen
                }
                ')' => {
                    self.advance();
                    TokenKind::RightParen
                }
                '{' => {
                    self.advance();
                    TokenKind::LeftBrace
                }
                '}' => {
                    self.advance();
                    TokenKind::RightBrace
                }
                ',' => {
                    self.advance();
                    TokenKind::Comma
                }
                ';' => {
                    self.advance();
                    TokenKind::Semicolon
                }
                '\'' | '"' => {
                    return Err(CompileError::new(
                        position,
                        "string expressions are not supported by the numeric VM",
                    ));
                }
                '[' | ']' | ':' => {
                    return Err(CompileError::new(
                        position,
                        "maps, arrays, and indexing are not supported by the numeric VM",
                    ));
                }
                _ => {
                    return Err(CompileError::new(
                        position,
                        format!("unsupported character {ch:?}"),
                    ));
                }
            };
            tokens.push(Token { kind, position });
        }
    }

    fn skip_space_and_comments(&mut self) -> Result<(), CompileError> {
        loop {
            while self.peek().is_some_and(char::is_whitespace) {
                self.advance();
            }
            if self.rest().starts_with("//") {
                while self.peek().is_some_and(|ch| ch != '\n') {
                    self.advance();
                }
                continue;
            }
            if self.rest().starts_with("/*") {
                let start = self.cursor;
                self.cursor += 2;
                while !self.rest().starts_with("*/") {
                    if self.peek().is_none() {
                        return Err(CompileError::new(start, "unterminated block comment"));
                    }
                    self.advance();
                }
                self.cursor += 2;
                continue;
            }
            return Ok(());
        }
    }

    fn number_starts_here(&self) -> bool {
        self.peek().is_some_and(|ch| ch.is_ascii_digit())
            || (self.peek() == Some('.')
                && self
                    .rest()
                    .chars()
                    .nth(1)
                    .is_some_and(|ch| ch.is_ascii_digit()))
    }

    fn number(&mut self) -> Result<TokenKind, CompileError> {
        let start = self.cursor;
        let mut seen_dot = false;
        let mut seen_exponent = false;
        while let Some(ch) = self.peek() {
            match ch {
                '0'..='9' | '_' => {
                    self.advance();
                }
                '.' if !seen_dot && !seen_exponent => {
                    seen_dot = true;
                    self.advance();
                }
                'e' | 'E' if !seen_exponent => {
                    seen_exponent = true;
                    self.advance();
                    if self.peek().is_some_and(|next| matches!(next, '+' | '-')) {
                        self.advance();
                    }
                }
                _ => break,
            }
        }
        let raw = self.source[start..self.cursor].replace('_', "");
        let value = raw
            .parse::<f64>()
            .map_err(|_| CompileError::new(start, format!("invalid numeric literal {raw:?}")))?;
        Ok(TokenKind::Number(value))
    }

    fn identifier(&mut self) -> TokenKind {
        let start = self.cursor;
        self.advance();
        while self
            .peek()
            .is_some_and(|ch| ch == '_' || ch.is_ascii_alphanumeric())
        {
            self.advance();
        }
        match &self.source[start..self.cursor] {
            "true" => TokenKind::True,
            "false" => TokenKind::False,
            "let" => TokenKind::Let,
            "const" => TokenKind::Const,
            "if" => TokenKind::If,
            "else" => TokenKind::Else,
            identifier => TokenKind::Identifier(identifier.to_string()),
        }
    }

    fn rest(&self) -> &'a str {
        &self.source[self.cursor..]
    }

    fn peek(&self) -> Option<char> {
        self.rest().chars().next()
    }

    fn advance(&mut self) -> Option<char> {
        let ch = self.peek()?;
        self.cursor += ch.len_utf8();
        Some(ch)
    }

    fn consume(&mut self, expected: char) -> bool {
        if self.peek() == Some(expected) {
            self.advance();
            true
        } else {
            false
        }
    }
}

struct Parser {
    tokens: Vec<Token>,
    cursor: usize,
}

impl Parser {
    fn new(tokens: Vec<Token>) -> Self {
        Self { tokens, cursor: 0 }
    }

    fn parse(mut self) -> Result<ParsedProgram, CompileError> {
        let mut statements = Vec::new();
        while matches!(self.current().kind, TokenKind::Let | TokenKind::Const) {
            let position = self.advance().position;
            let name = match self.advance().kind {
                TokenKind::Identifier(name) => name,
                _ => return Err(self.error("expected identifier after let/const")),
            };
            self.expect(TokenKind::Equal, "expected = in local declaration")?;
            let value = self.expression(0)?;
            self.expect(TokenKind::Semicolon, "expected ; after local declaration")?;
            statements.push(LetStatement {
                name,
                value,
                position,
            });
        }
        if matches!(self.current().kind, TokenKind::Eof) {
            return Err(self.error("numeric VM program requires a result expression"));
        }
        let result = self.expression(0)?;
        while matches!(self.current().kind, TokenKind::Semicolon) {
            self.advance();
        }
        if !matches!(self.current().kind, TokenKind::Eof) {
            return Err(self.error("unexpected token after result expression"));
        }
        Ok(ParsedProgram { statements, result })
    }

    fn expression(&mut self, min_precedence: u8) -> Result<Expr, CompileError> {
        let mut lhs = self.prefix()?;
        loop {
            let Some((operator, precedence)) = self.binary_operator() else {
                break;
            };
            if precedence < min_precedence {
                break;
            }
            let position = self.advance().position;
            let rhs = self.expression(precedence + 1)?;
            lhs = Expr::Binary {
                operator,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
                position,
            };
        }
        Ok(lhs)
    }

    fn prefix(&mut self) -> Result<Expr, CompileError> {
        let token = self.advance();
        match token.kind {
            TokenKind::Number(value) => Ok(Expr::Number(value, token.position)),
            TokenKind::True => Ok(Expr::Boolean(true, token.position)),
            TokenKind::False => Ok(Expr::Boolean(false, token.position)),
            TokenKind::Identifier(name) => {
                if matches!(self.current().kind, TokenKind::LeftParen) {
                    self.advance();
                    let mut args = Vec::new();
                    if !matches!(self.current().kind, TokenKind::RightParen) {
                        loop {
                            args.push(self.expression(0)?);
                            if matches!(self.current().kind, TokenKind::Comma) {
                                self.advance();
                            } else {
                                break;
                            }
                        }
                    }
                    self.expect(TokenKind::RightParen, "expected ) after function arguments")?;
                    Ok(Expr::Call {
                        name,
                        args,
                        position: token.position,
                    })
                } else {
                    Ok(Expr::Variable(name, token.position))
                }
            }
            TokenKind::Minus => Ok(Expr::Unary {
                operator: UnaryOp::Negate,
                operand: Box::new(self.expression(7)?),
                position: token.position,
            }),
            TokenKind::Plus => self.expression(7),
            TokenKind::Bang => Ok(Expr::Unary {
                operator: UnaryOp::Not,
                operand: Box::new(self.expression(7)?),
                position: token.position,
            }),
            TokenKind::LeftParen => {
                let expression = self.expression(0)?;
                self.expect(TokenKind::RightParen, "expected )")?;
                Ok(expression)
            }
            TokenKind::If => self.if_expression(token.position),
            _ => Err(CompileError::new(
                token.position,
                "expected numeric or boolean expression",
            )),
        }
    }

    fn if_expression(&mut self, position: usize) -> Result<Expr, CompileError> {
        let condition = self.expression(0)?;
        self.expect(TokenKind::LeftBrace, "expected { after if condition")?;
        let when_true = self.expression(0)?;
        while matches!(self.current().kind, TokenKind::Semicolon) {
            self.advance();
        }
        self.expect(TokenKind::RightBrace, "expected } after true branch")?;
        self.expect(TokenKind::Else, "expected else after if branch")?;
        self.expect(TokenKind::LeftBrace, "expected { after else")?;
        let when_false = self.expression(0)?;
        while matches!(self.current().kind, TokenKind::Semicolon) {
            self.advance();
        }
        self.expect(TokenKind::RightBrace, "expected } after false branch")?;
        Ok(Expr::If {
            condition: Box::new(condition),
            when_true: Box::new(when_true),
            when_false: Box::new(when_false),
            position,
        })
    }

    fn binary_operator(&self) -> Option<(ParsedBinaryOp, u8)> {
        Some(match self.current().kind {
            TokenKind::OrOr => (ParsedBinaryOp::Or, 1),
            TokenKind::AndAnd => (ParsedBinaryOp::And, 2),
            TokenKind::EqualEqual => (ParsedBinaryOp::Equal, 3),
            TokenKind::BangEqual => (ParsedBinaryOp::NotEqual, 3),
            TokenKind::Less => (ParsedBinaryOp::Less, 4),
            TokenKind::LessEqual => (ParsedBinaryOp::LessEqual, 4),
            TokenKind::Greater => (ParsedBinaryOp::Greater, 4),
            TokenKind::GreaterEqual => (ParsedBinaryOp::GreaterEqual, 4),
            TokenKind::Plus => (ParsedBinaryOp::Add, 5),
            TokenKind::Minus => (ParsedBinaryOp::Subtract, 5),
            TokenKind::Star => (ParsedBinaryOp::Multiply, 6),
            TokenKind::Slash => (ParsedBinaryOp::Divide, 6),
            TokenKind::Percent => (ParsedBinaryOp::Remainder, 6),
            _ => return None,
        })
    }

    fn expect(&mut self, expected: TokenKind, message: &str) -> Result<(), CompileError> {
        if std::mem::discriminant(&self.current().kind) == std::mem::discriminant(&expected) {
            self.advance();
            Ok(())
        } else {
            Err(self.error(message))
        }
    }

    fn current(&self) -> &Token {
        &self.tokens[self.cursor]
    }

    fn advance(&mut self) -> Token {
        let token = self.tokens[self.cursor].clone();
        if !matches!(token.kind, TokenKind::Eof) {
            self.cursor += 1;
        }
        token
    }

    fn error(&self, message: impl Into<String>) -> CompileError {
        CompileError::new(self.current().position, message)
    }
}

struct Compiler<'a, F>
where
    F: Fn(&str) -> Option<ValueType>,
{
    resolve_type: &'a F,
    instructions: Vec<Instruction>,
    variables: Vec<String>,
    variable_types: Vec<ValueType>,
    variable_indices: BTreeMap<String, u16>,
    locals: BTreeMap<String, (u16, ValueType)>,
}

impl<'a, F> Compiler<'a, F>
where
    F: Fn(&str) -> Option<ValueType>,
{
    fn new(resolve_type: &'a F) -> Self {
        Self {
            resolve_type,
            instructions: Vec::new(),
            variables: Vec::new(),
            variable_types: Vec::new(),
            variable_indices: BTreeMap::new(),
            locals: BTreeMap::new(),
        }
    }

    fn compile(mut self, parsed: ParsedProgram) -> Result<Program, CompileError> {
        for statement in parsed.statements {
            if self.locals.contains_key(&statement.name) {
                return Err(CompileError::new(
                    statement.position,
                    format!("duplicate local variable {}", statement.name),
                ));
            }
            let value_type = self.expression(&statement.value)?;
            let slot = u16::try_from(self.locals.len())
                .map_err(|_| CompileError::new(statement.position, "too many local variables"))?;
            self.instructions.push(Instruction::StoreLocal(slot));
            self.locals.insert(statement.name, (slot, value_type));
        }
        let result_type = self.expression(&parsed.result)?;
        self.instructions.push(Instruction::Return);
        Ok(Program {
            instructions: self.instructions,
            variables: self.variables,
            variable_types: self.variable_types,
            local_count: self.locals.len(),
            result_type,
        })
    }

    fn expression(&mut self, expression: &Expr) -> Result<ValueType, CompileError> {
        match expression {
            Expr::Number(value, _) => {
                self.instructions
                    .push(Instruction::Push(Value::Number(*value)));
                Ok(ValueType::Number)
            }
            Expr::Boolean(value, _) => {
                self.instructions
                    .push(Instruction::Push(Value::Boolean(*value)));
                Ok(ValueType::Boolean)
            }
            Expr::Variable(name, position) => {
                if let Some((slot, value_type)) = self.locals.get(name).copied() {
                    self.instructions.push(Instruction::LoadLocal(slot));
                    return Ok(value_type);
                }
                let value_type = (self.resolve_type)(name).ok_or_else(|| {
                    CompileError::new(
                        *position,
                        format!("identifier {name} is not numeric or boolean"),
                    )
                })?;
                let index = if let Some(index) = self.variable_indices.get(name).copied() {
                    index
                } else {
                    let index = u16::try_from(self.variables.len())
                        .map_err(|_| CompileError::new(*position, "too many external variables"))?;
                    self.variables.push(name.clone());
                    self.variable_types.push(value_type);
                    self.variable_indices.insert(name.clone(), index);
                    index
                };
                self.instructions.push(Instruction::LoadVariable(index));
                Ok(value_type)
            }
            Expr::Unary {
                operator,
                operand,
                position,
            } => {
                let operand_type = self.expression(operand)?;
                let required = match operator {
                    UnaryOp::Negate => ValueType::Number,
                    UnaryOp::Not => ValueType::Boolean,
                };
                require_type(operand_type, required, *position)?;
                self.instructions.push(Instruction::Unary(*operator));
                Ok(required)
            }
            Expr::Binary {
                operator,
                lhs,
                rhs,
                position,
            } => self.binary(*operator, lhs, rhs, *position),
            Expr::Call {
                name,
                args,
                position,
            } => self.call(name, args, *position),
            Expr::If {
                condition,
                when_true,
                when_false,
                position,
            } => {
                let condition_type = self.expression(condition)?;
                require_type(condition_type, ValueType::Boolean, *position)?;
                let false_jump = self.instructions.len();
                self.instructions.push(Instruction::JumpIfFalse(usize::MAX));
                let true_type = self.expression(when_true)?;
                let end_jump = self.instructions.len();
                self.instructions.push(Instruction::Jump(usize::MAX));
                let false_target = self.instructions.len();
                patch_jump(&mut self.instructions, false_jump, false_target)?;
                let false_type = self.expression(when_false)?;
                require_type(false_type, true_type, when_false.position())?;
                let end_target = self.instructions.len();
                patch_jump(&mut self.instructions, end_jump, end_target)?;
                Ok(true_type)
            }
        }
    }

    fn binary(
        &mut self,
        operator: ParsedBinaryOp,
        lhs: &Expr,
        rhs: &Expr,
        position: usize,
    ) -> Result<ValueType, CompileError> {
        if matches!(operator, ParsedBinaryOp::And | ParsedBinaryOp::Or) {
            let lhs_type = self.expression(lhs)?;
            require_type(lhs_type, ValueType::Boolean, position)?;
            let branch = self.instructions.len();
            self.instructions.push(match operator {
                ParsedBinaryOp::And => Instruction::JumpIfFalse(usize::MAX),
                ParsedBinaryOp::Or => Instruction::JumpIfTrue(usize::MAX),
                _ => unreachable!(),
            });
            let rhs_type = self.expression(rhs)?;
            require_type(rhs_type, ValueType::Boolean, rhs.position())?;
            let end_jump = self.instructions.len();
            self.instructions.push(Instruction::Jump(usize::MAX));
            let short_target = self.instructions.len();
            self.instructions
                .push(Instruction::Push(Value::Boolean(matches!(
                    operator,
                    ParsedBinaryOp::Or
                ))));
            let end_target = self.instructions.len();
            patch_jump(&mut self.instructions, branch, short_target)?;
            patch_jump(&mut self.instructions, end_jump, end_target)?;
            return Ok(ValueType::Boolean);
        }

        let lhs_type = self.expression(lhs)?;
        let rhs_type = self.expression(rhs)?;
        let (required, result, instruction) = match operator {
            ParsedBinaryOp::Add => (ValueType::Number, ValueType::Number, BinaryOp::Add),
            ParsedBinaryOp::Subtract => (ValueType::Number, ValueType::Number, BinaryOp::Subtract),
            ParsedBinaryOp::Multiply => (ValueType::Number, ValueType::Number, BinaryOp::Multiply),
            ParsedBinaryOp::Divide => (ValueType::Number, ValueType::Number, BinaryOp::Divide),
            ParsedBinaryOp::Remainder => {
                (ValueType::Number, ValueType::Number, BinaryOp::Remainder)
            }
            ParsedBinaryOp::Less => (ValueType::Number, ValueType::Boolean, BinaryOp::Less),
            ParsedBinaryOp::LessEqual => {
                (ValueType::Number, ValueType::Boolean, BinaryOp::LessEqual)
            }
            ParsedBinaryOp::Greater => (ValueType::Number, ValueType::Boolean, BinaryOp::Greater),
            ParsedBinaryOp::GreaterEqual => (
                ValueType::Number,
                ValueType::Boolean,
                BinaryOp::GreaterEqual,
            ),
            ParsedBinaryOp::Equal => {
                require_type(rhs_type, lhs_type, position)?;
                self.instructions.push(Instruction::Binary(BinaryOp::Equal));
                return Ok(ValueType::Boolean);
            }
            ParsedBinaryOp::NotEqual => {
                require_type(rhs_type, lhs_type, position)?;
                self.instructions
                    .push(Instruction::Binary(BinaryOp::NotEqual));
                return Ok(ValueType::Boolean);
            }
            ParsedBinaryOp::And | ParsedBinaryOp::Or => unreachable!(),
        };
        require_type(lhs_type, required, lhs.position())?;
        require_type(rhs_type, required, rhs.position())?;
        self.instructions.push(Instruction::Binary(instruction));
        Ok(result)
    }

    fn call(
        &mut self,
        name: &str,
        args: &[Expr],
        position: usize,
    ) -> Result<ValueType, CompileError> {
        let (builtin, expected_args, result_type) = match name {
            "round" => (Builtin::Round, &[ValueType::Number][..], ValueType::Number),
            "floor" => (Builtin::Floor, &[ValueType::Number][..], ValueType::Number),
            "ceil" => (Builtin::Ceil, &[ValueType::Number][..], ValueType::Number),
            "abs" => (Builtin::Abs, &[ValueType::Number][..], ValueType::Number),
            "min" => (
                Builtin::Min,
                &[ValueType::Number, ValueType::Number][..],
                ValueType::Number,
            ),
            "max" => (
                Builtin::Max,
                &[ValueType::Number, ValueType::Number][..],
                ValueType::Number,
            ),
            "sqrt" => (Builtin::Sqrt, &[ValueType::Number][..], ValueType::Number),
            "pow" => (
                Builtin::Pow,
                &[ValueType::Number, ValueType::Number][..],
                ValueType::Number,
            ),
            "log" => (Builtin::Log, &[ValueType::Number][..], ValueType::Number),
            "exp" => (Builtin::Exp, &[ValueType::Number][..], ValueType::Number),
            "clamp" => (
                Builtin::Clamp,
                &[ValueType::Number, ValueType::Number, ValueType::Number][..],
                ValueType::Number,
            ),
            "between" => (
                Builtin::Between,
                &[ValueType::Number, ValueType::Number, ValueType::Number][..],
                ValueType::Boolean,
            ),
            "nz" => (
                Builtin::Nz,
                &[ValueType::Number, ValueType::Number][..],
                ValueType::Number,
            ),
            "safe_div" if args.len() == 2 => (
                Builtin::SafeDiv,
                &[ValueType::Number, ValueType::Number][..],
                ValueType::Number,
            ),
            "safe_div" if args.len() == 3 => (
                Builtin::SafeDiv,
                &[ValueType::Number, ValueType::Number, ValueType::Number][..],
                ValueType::Number,
            ),
            "iff" if args.len() == 3 => {
                let condition_type = self.expression(&args[0])?;
                require_type(condition_type, ValueType::Boolean, args[0].position())?;
                let true_type = self.expression(&args[1])?;
                let false_type = self.expression(&args[2])?;
                require_type(false_type, true_type, args[2].position())?;
                self.instructions.push(Instruction::Call {
                    builtin: Builtin::Iff,
                    argc: 3,
                });
                return Ok(true_type);
            }
            _ => {
                return Err(CompileError::new(
                    position,
                    format!("function {name} is not supported by the numeric VM"),
                ));
            }
        };
        if args.len() != expected_args.len() {
            return Err(CompileError::new(
                position,
                format!(
                    "function {name} expects {} arguments, got {}",
                    expected_args.len(),
                    args.len()
                ),
            ));
        }
        for (argument, expected_type) in args.iter().zip(expected_args) {
            let actual_type = self.expression(argument)?;
            require_type(actual_type, *expected_type, argument.position())?;
        }
        self.instructions.push(Instruction::Call {
            builtin,
            argc: u8::try_from(args.len())
                .map_err(|_| CompileError::new(position, "too many function arguments"))?,
        });
        Ok(result_type)
    }
}

fn patch_jump(
    instructions: &mut [Instruction],
    index: usize,
    target: usize,
) -> Result<(), CompileError> {
    let Some(instruction) = instructions.get_mut(index) else {
        return Err(CompileError::new(0, "invalid jump patch index"));
    };
    match instruction {
        Instruction::JumpIfFalse(value)
        | Instruction::JumpIfTrue(value)
        | Instruction::Jump(value) => {
            *value = target;
            Ok(())
        }
        _ => Err(CompileError::new(0, "instruction is not a jump")),
    }
}

fn require_type(
    actual: ValueType,
    expected: ValueType,
    position: usize,
) -> Result<(), CompileError> {
    if actual == expected {
        Ok(())
    } else {
        Err(CompileError::new(
            position,
            format!("expected {expected:?}, got {actual:?}"),
        ))
    }
}

pub(crate) fn compile<F>(source: &str, resolve_type: F) -> Result<Program, CompileError>
where
    F: Fn(&str) -> Option<ValueType>,
{
    let tokens = Lexer::new(source).tokenize()?;
    let parsed = Parser::new(tokens).parse()?;
    Compiler::new(&resolve_type).compile(parsed)
}

#[cfg(test)]
mod tests {
    use std::hint::black_box;
    use std::time::Instant;

    use rhai::{Engine, Scope};

    use super::{EvalError, Scratch, Value, ValueType, compile};

    fn evaluate(source: &str, values: &[(&str, Value)]) -> Value {
        let program = compile(source, |name| {
            values
                .iter()
                .find_map(|(key, value)| (*key == name).then_some(value.value_type()))
        })
        .expect("compile");
        let mut scratch = Scratch::default();
        program
            .evaluate(&mut scratch, |_index, name, _expected| {
                values
                    .iter()
                    .find_map(|(key, value)| (*key == name).then_some(*value))
                    .ok_or_else(|| EvalError::new(format!("missing {name}")))
            })
            .expect("evaluate")
    }

    #[test]
    fn evaluates_numeric_boolean_and_if_expressions() {
        let source = r#"
            let clamped = clamp(close, 2.0, 20.0);
            let ratio = safe_div(clamped, previous, 0.0);
            if ratio > 1.0 && !paused { min(ratio, 2.0) } else { 0.5 }
        "#;
        assert_eq!(
            evaluate(
                source,
                &[
                    ("close", Value::Number(12.0)),
                    ("previous", Value::Number(10.0)),
                    ("paused", Value::Boolean(false)),
                ],
            ),
            Value::Number(1.2)
        );
    }

    #[test]
    fn floating_comparisons_match_rhai_epsilon_semantics() {
        let adjacent = 11.699999999999998_f64;
        assert_eq!(
            evaluate("value == 11.7", &[("value", Value::Number(adjacent))]),
            Value::Boolean(true)
        );
        assert_eq!(
            evaluate("value != 11.7", &[("value", Value::Number(adjacent))]),
            Value::Boolean(false)
        );
        assert_eq!(
            evaluate("value <= 11.7", &[("value", Value::Number(adjacent))]),
            Value::Boolean(true)
        );
        assert_eq!(
            evaluate("value >= 11.7", &[("value", Value::Number(adjacent))]),
            Value::Boolean(true)
        );
        assert_eq!(
            evaluate("value < 11.7", &[("value", Value::Number(adjacent))]),
            Value::Boolean(false)
        );
        assert_eq!(
            evaluate("value > 11.7", &[("value", Value::Number(adjacent))]),
            Value::Boolean(false)
        );
    }

    #[test]
    fn short_circuit_does_not_resolve_unused_variable() {
        let program = compile("false && missing", |name| {
            (name == "missing").then_some(ValueType::Boolean)
        })
        .expect("compile");
        let mut scratch = Scratch::default();
        let result = program
            .evaluate(&mut scratch, |_index, name, _expected| {
                Err(EvalError::new(format!("unexpected resolution of {name}")))
            })
            .expect("short circuit");
        assert_eq!(result, Value::Boolean(false));
    }

    #[test]
    fn supports_numeric_builtins_and_local_reuse() {
        let source = r#"
            const scale = pow(2.0, 3.0);
            let adjusted = nz(value, 4.0) * scale;
            between(adjusted, 31.0, 33.0)
        "#;
        assert_eq!(
            evaluate(source, &[("value", Value::Number(4.0))]),
            Value::Boolean(true)
        );
    }

    #[test]
    fn rejects_dynamic_language_features() {
        assert!(compile("factors[\"close\"]", |_| Some(ValueType::Number)).is_err());
        assert!(compile("contains(symbol, \"SZ\")", |_| None).is_err());
        assert!(compile("custom_score(close)", |_| Some(ValueType::Number)).is_err());
    }

    #[test]
    fn scratch_reuses_allocations_across_program_runs() {
        let program = compile("left + right * 2.0", |_| Some(ValueType::Number)).expect("compile");
        let mut scratch = Scratch::default();
        for _ in 0..10 {
            let value = program
                .evaluate(&mut scratch, |_index, name, _expected| {
                    Ok(Value::Number(if name == "left" { 1.0 } else { 2.0 }))
                })
                .expect("evaluate");
            assert_eq!(value, Value::Number(5.0));
        }
    }

    #[test]
    fn matches_rhai_for_representative_numeric_boolean_corpus() {
        let source = r#"
            let clamped = clamp(signal_close, 2000.0, 3000.0);
            let ratio = safe_div(clamped - 2000.0, 1000.0, 0.0);
            let trend = ma10 > ma30 ? 1.0 : 0.3;
            let exposure = drawdown >= 0.08 ? min(trend, 0.2) : trend;
            ready && !paused ? exposure * (1.0 + ratio) : 0.0
        "#;
        let normalized = source
            .replace(
                "ma10 > ma30 ? 1.0 : 0.3",
                "if ma10 > ma30 { 1.0 } else { 0.3 }",
            )
            .replace(
                "drawdown >= 0.08 ? min(trend, 0.2) : trend",
                "if drawdown >= 0.08 { min(trend, 0.2) } else { trend }",
            )
            .replace(
                "ready && !paused ? exposure * (1.0 + ratio) : 0.0",
                "if ready && !paused { exposure * (1.0 + ratio) } else { 0.0 }",
            );
        let program = compile(&normalized, |name| match name {
            "ready" | "paused" => Some(ValueType::Boolean),
            _ => Some(ValueType::Number),
        })
        .expect("compile VM");
        let mut engine = Engine::new();
        engine.register_fn("min", |lhs: f64, rhs: f64| lhs.min(rhs));
        engine.register_fn("clamp", |value: f64, low: f64, high: f64| {
            value.clamp(low, high)
        });
        engine.register_fn("safe_div", |lhs: f64, rhs: f64, fallback: f64| {
            if rhs.abs() <= f64::EPSILON {
                fallback
            } else {
                lhs / rhs
            }
        });
        let ast = engine.compile(&normalized).expect("compile Rhai");
        let mut scratch = Scratch::default();
        for index in 0..256 {
            let signal_close = 1800.0 + index as f64 * 7.0;
            let ma10 = 10.0 + (index % 9) as f64;
            let ma30 = 12.0 + (index % 7) as f64;
            let drawdown = (index % 13) as f64 / 100.0;
            let ready = index % 5 != 0;
            let paused = index % 17 == 0;
            let values = [
                ("signal_close", Value::Number(signal_close)),
                ("ma10", Value::Number(ma10)),
                ("ma30", Value::Number(ma30)),
                ("drawdown", Value::Number(drawdown)),
                ("ready", Value::Boolean(ready)),
                ("paused", Value::Boolean(paused)),
            ];
            let vm_value = program
                .evaluate(&mut scratch, |_slot, name, _expected| {
                    values
                        .iter()
                        .find_map(|(key, value)| (*key == name).then_some(*value))
                        .ok_or_else(|| EvalError::new(format!("missing {name}")))
                })
                .expect("evaluate VM")
                .as_number()
                .expect("numeric VM result");
            let mut scope = Scope::new();
            scope.push("signal_close", signal_close);
            scope.push("ma10", ma10);
            scope.push("ma30", ma30);
            scope.push("drawdown", drawdown);
            scope.push("ready", ready);
            scope.push("paused", paused);
            let rhai_value = engine
                .eval_ast_with_scope::<f64>(&mut scope, &ast)
                .expect("evaluate Rhai");
            assert_eq!(vm_value.to_bits(), rhai_value.to_bits(), "case {index}");
        }
    }

    #[test]
    #[ignore = "manual release-mode performance evidence"]
    fn benchmark_numeric_vm_against_rhai() {
        let source = "!is_st && !is_star_st && close > ma5 && ma5 > ma10 && ma10 > ma30 && volume_ma5 < volume_ma100";
        let program = compile(source, |name| match name {
            "is_st" | "is_star_st" => Some(ValueType::Boolean),
            _ => Some(ValueType::Number),
        })
        .expect("compile VM");
        let values = [
            ("is_st", Value::Boolean(false)),
            ("is_star_st", Value::Boolean(false)),
            ("close", Value::Number(12.0)),
            ("ma5", Value::Number(11.5)),
            ("ma10", Value::Number(11.0)),
            ("ma30", Value::Number(10.0)),
            ("volume_ma5", Value::Number(8_000.0)),
            ("volume_ma100", Value::Number(10_000.0)),
        ];
        let iterations = 2_000_000usize;
        let mut scratch = Scratch::default();
        let vm_started = Instant::now();
        let mut vm_true = 0usize;
        for _ in 0..iterations {
            let result = program
                .evaluate(&mut scratch, |_slot, name, _expected| {
                    values
                        .iter()
                        .find_map(|(key, value)| (*key == name).then_some(*value))
                        .ok_or_else(|| EvalError::new(format!("missing {name}")))
                })
                .expect("VM evaluation")
                .as_bool()
                .expect("boolean VM result");
            vm_true += usize::from(black_box(result));
        }
        let vm_elapsed = vm_started.elapsed();

        let engine = Engine::new();
        let ast = engine.compile(source).expect("compile Rhai");
        let mut scope = Scope::new();
        scope.push("is_st", false);
        scope.push("is_star_st", false);
        scope.push("close", 12.0_f64);
        scope.push("ma5", 11.5_f64);
        scope.push("ma10", 11.0_f64);
        scope.push("ma30", 10.0_f64);
        scope.push("volume_ma5", 8_000.0_f64);
        scope.push("volume_ma100", 10_000.0_f64);
        let rhai_started = Instant::now();
        let mut rhai_true = 0usize;
        for _ in 0..iterations {
            let result = engine
                .eval_ast_with_scope::<bool>(&mut scope, &ast)
                .expect("Rhai evaluation");
            rhai_true += usize::from(black_box(result));
        }
        let rhai_elapsed = rhai_started.elapsed();
        assert_eq!(vm_true, rhai_true);
        println!(
            "numeric_vm_benchmark iterations={} vm_ns_per_eval={:.3} rhai_ns_per_eval={:.3} speedup={:.3}",
            iterations,
            vm_elapsed.as_nanos() as f64 / iterations as f64,
            rhai_elapsed.as_nanos() as f64 / iterations as f64,
            rhai_elapsed.as_secs_f64() / vm_elapsed.as_secs_f64(),
        );
    }
}