thal 0.0.1

use super::ast::*;
use super::lex::{lex, Tok, Token};
use crate::value::PrimitiveType;
use crate::Error;

pub fn parse_str(source: &str) -> Result<Program, Error> {
    let tokens = lex(source)?;
    let mut p = Parser::new(tokens);
    p.parse_program()
}

pub fn parse_file(path: impl AsRef<std::path::Path>) -> Result<Program, Error> {
    let src = std::fs::read_to_string(path)?;
    parse_str(&src)
}

struct Parser {
    tokens: Vec<Tok>,
    pos: usize,
}

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

    fn peek(&self) -> Option<&Token> {
        self.tokens.get(self.pos).map(|t| &t.token)
    }

    fn advance(&mut self) -> Option<Tok> {
        let t = self.tokens.get(self.pos).cloned();
        if t.is_some() {
            self.pos += 1;
        }
        t
    }

    fn expect(&mut self, expected: Token) -> Result<Tok, Error> {
        let actual = self.advance().ok_or_else(|| {
            Error::Parse(format!("expected {expected:?}, got end of input"))
        })?;
        if std::mem::discriminant(&actual.token) == std::mem::discriminant(&expected) {
            Ok(actual)
        } else {
            Err(Error::Parse(format!(
                "expected {expected:?}, got {:?} at line {}:{}",
                actual.token, actual.span.line, actual.span.col
            )))
        }
    }

    fn expect_ident(&mut self) -> Result<String, Error> {
        let tok = self.advance().ok_or_else(|| {
            Error::Parse("expected identifier, got end of input".into())
        })?;
        match tok.token {
            Token::Ident(s) => Ok(s),
            other => Err(Error::Parse(format!(
                "expected identifier, got {other:?} at line {}:{}",
                tok.span.line, tok.span.col
            ))),
        }
    }

    fn parse_program(&mut self) -> Result<Program, Error> {
        let mut program = Program::default();
        while self.peek().is_some() {
            match self.peek() {
                Some(Token::Type) => program.types.push(self.parse_type_decl()?),
                Some(Token::Mixin) => program.mixins.push(self.parse_mixin_decl()?),
                Some(Token::Molecule) => program.molecules.push(self.parse_molecule_decl()?),
                Some(Token::Reaction) => program.reactions.push(self.parse_reaction_decl()?),
                Some(other) => {
                    return Err(Error::Parse(format!(
                        "unexpected top-level token {other:?}"
                    )))
                }
                None => break,
            }
        }
        Ok(program)
    }

    fn parse_mixin_decl(&mut self) -> Result<MixinDecl, Error> {
        self.expect(Token::Mixin)?;
        let name = self.expect_ident()?;
        self.expect(Token::LBrace)?;
        let mut fields = Vec::new();
        let mut primary_key: Option<Vec<String>> = None;
        let mut merge: Option<MergeFn> = None;

        while !matches!(self.peek(), Some(Token::RBrace)) {
            match self.peek() {
                Some(Token::PrimaryKey) => {
                    self.advance();
                    self.expect(Token::Colon)?;
                    primary_key = Some(self.parse_primary_key_body()?);
                }
                Some(Token::Merge) => {
                    self.advance();
                    self.expect(Token::Colon)?;
                    self.expect(Token::Pipe)?;
                    let old = self.expect_ident()?;
                    self.expect(Token::Comma)?;
                    let new = self.expect_ident()?;
                    self.expect(Token::Pipe)?;
                    let body = self.parse_expr()?;
                    merge = Some(MergeFn {
                        old_binding: old,
                        new_binding: new,
                        body,
                    });
                }
                Some(Token::Ident(_)) => {
                    fields.push(self.parse_field_decl()?);
                }
                Some(other) => {
                    return Err(Error::Parse(format!(
                        "unexpected token in mixin body: {other:?}"
                    )))
                }
                None => return Err(Error::Parse("unterminated mixin body".into())),
            }
        }
        self.expect(Token::RBrace)?;
        Ok(MixinDecl {
            name,
            fields,
            primary_key,
            merge,
        })
    }

    fn parse_type_decl(&mut self) -> Result<TypeDecl, Error> {
        self.expect(Token::Type)?;
        let name = self.expect_ident()?;
        self.expect(Token::Eq)?;
        self.expect(Token::Enum)?;
        self.expect(Token::LBrace)?;
        let mut variants = Vec::new();
        while !matches!(self.peek(), Some(Token::RBrace)) {
            variants.push(self.expect_ident()?);
            if matches!(self.peek(), Some(Token::Comma)) {
                self.advance();
            }
        }
        self.expect(Token::RBrace)?;
        Ok(TypeDecl {
            name,
            body: TypeDeclBody::Enum(variants),
        })
    }

    fn parse_molecule_decl(&mut self) -> Result<MoleculeDecl, Error> {
        self.expect(Token::Molecule)?;
        let name = self.expect_ident()?;

        let mut mixins = Vec::new();
        if matches!(self.peek(), Some(Token::With)) {
            self.advance();
            mixins.push(self.expect_ident()?);
            while matches!(self.peek(), Some(Token::Comma)) {
                self.advance();
                mixins.push(self.expect_ident()?);
            }
        }

        self.expect(Token::LBrace)?;
        let mut fields = Vec::new();
        let mut primary_key: Option<Vec<String>> = None;
        let mut merge: Option<MergeFn> = None;

        while !matches!(self.peek(), Some(Token::RBrace)) {
            match self.peek() {
                Some(Token::PrimaryKey) => {
                    self.advance();
                    self.expect(Token::Colon)?;
                    primary_key = Some(self.parse_primary_key_body()?);
                }
                Some(Token::Merge) => {
                    self.advance();
                    self.expect(Token::Colon)?;
                    self.expect(Token::Pipe)?;
                    let old = self.expect_ident()?;
                    self.expect(Token::Comma)?;
                    let new = self.expect_ident()?;
                    self.expect(Token::Pipe)?;
                    let body = self.parse_expr()?;
                    merge = Some(MergeFn {
                        old_binding: old,
                        new_binding: new,
                        body,
                    });
                }
                Some(Token::Ident(_)) => {
                    fields.push(self.parse_field_decl()?);
                }
                Some(other) => {
                    return Err(Error::Parse(format!(
                        "unexpected token in molecule body: {other:?}"
                    )))
                }
                None => return Err(Error::Parse("unterminated molecule body".into())),
            }
        }
        self.expect(Token::RBrace)?;

        // primary_key may still be None here — the verifier resolves it from
        // mixins, and errors out if no source provides one.
        Ok(MoleculeDecl {
            name,
            mixins,
            fields,
            primary_key,
            merge,
        })
    }

    fn parse_primary_key_body(&mut self) -> Result<Vec<String>, Error> {
        if matches!(self.peek(), Some(Token::LParen)) {
            self.advance();
            let mut keys = Vec::new();
            while !matches!(self.peek(), Some(Token::RParen)) {
                keys.push(self.expect_ident()?);
                if matches!(self.peek(), Some(Token::Comma)) {
                    self.advance();
                }
            }
            self.expect(Token::RParen)?;
            Ok(keys)
        } else {
            Ok(vec![self.expect_ident()?])
        }
    }

    fn parse_field_decl(&mut self) -> Result<FieldDecl, Error> {
        let name = self.expect_ident()?;
        self.expect(Token::Colon)?;
        let ty = self.parse_type_expr()?;
        let mut default = None;
        if matches!(self.peek(), Some(Token::LBrace)) {
            self.advance();
            self.expect(Token::Default)?;
            self.expect(Token::Colon)?;
            default = Some(self.parse_expr()?);
            self.expect(Token::RBrace)?;
        }
        Ok(FieldDecl { name, ty, default })
    }

    fn parse_type_expr(&mut self) -> Result<TypeExpr, Error> {
        let base = self.parse_type_expr_atom()?;
        if matches!(self.peek(), Some(Token::Question)) {
            self.advance();
            Ok(TypeExpr::Optional(Box::new(base)))
        } else {
            Ok(base)
        }
    }

    fn parse_type_expr_atom(&mut self) -> Result<TypeExpr, Error> {
        if matches!(self.peek(), Some(Token::LBracket)) {
            self.advance();
            let inner = self.parse_type_expr()?;
            self.expect(Token::RBracket)?;
            return Ok(TypeExpr::List(Box::new(inner)));
        }
        let id = self.expect_ident()?;
        Ok(match id.as_str() {
            "Int" => TypeExpr::Primitive(PrimitiveType::Int),
            "Float" => TypeExpr::Primitive(PrimitiveType::Float),
            "Bool" => TypeExpr::Primitive(PrimitiveType::Bool),
            "String" => TypeExpr::Primitive(PrimitiveType::String),
            "Bytes" => TypeExpr::Primitive(PrimitiveType::Bytes),
            "Timestamp" => TypeExpr::Primitive(PrimitiveType::Timestamp),
            "Duration" => TypeExpr::Primitive(PrimitiveType::Duration),
            "Uuid" => TypeExpr::Primitive(PrimitiveType::Uuid),
            _ => TypeExpr::Named(id),
        })
    }

    fn parse_reaction_decl(&mut self) -> Result<ReactionDecl, Error> {
        self.expect(Token::Reaction)?;
        let name = self.expect_ident()?;
        self.expect(Token::LBrace)?;
        self.expect(Token::When)?;
        self.expect(Token::Colon)?;
        let mut when = vec![self.parse_mol_pattern()?];
        while matches!(self.peek(), Some(Token::Comma)) {
            self.advance();
            when.push(self.parse_mol_pattern()?);
        }

        let rollup = if matches!(self.peek(), Some(Token::Rollup)) {
            self.advance();
            self.expect(Token::Colon)?;
            let molecule_name = self.expect_ident()?;
            self.expect(Token::LParen)?;
            let binding = self.expect_ident()?;
            self.expect(Token::RParen)?;
            self.expect(Token::By)?;
            let predicate = self.parse_expr()?;
            Some(RollupClause {
                molecule_name,
                binding,
                predicate,
            })
        } else {
            None
        };

        let where_clause = if matches!(self.peek(), Some(Token::Where)) {
            self.advance();
            self.expect(Token::Colon)?;
            Some(self.parse_expr()?)
        } else {
            None
        };

        self.expect(Token::Emit)?;
        self.expect(Token::Colon)?;
        let mut emit = vec![{
            let (molecule_name, fields) = self.parse_struct_lit()?;
            EmitClause { molecule_name, fields }
        }];
        while matches!(self.peek(), Some(Token::Comma)) {
            self.advance();
            let (molecule_name, fields) = self.parse_struct_lit()?;
            emit.push(EmitClause { molecule_name, fields });
        }
        self.expect(Token::RBrace)?;
        Ok(ReactionDecl {
            name,
            when,
            rollup,
            where_clause,
            emit,
        })
    }

    fn parse_mol_pattern(&mut self) -> Result<MoleculePattern, Error> {
        let molecule_name = self.expect_ident()?;
        self.expect(Token::LParen)?;
        let binding = self.expect_ident()?;
        self.expect(Token::RParen)?;
        Ok(MoleculePattern {
            molecule_name,
            binding,
        })
    }

fn parse_struct_lit(&mut self) -> Result<(String, Vec<FieldAssign>), Error> {
        let name = self.expect_ident()?;
        self.expect(Token::LBrace)?;
        let mut fields = Vec::new();
        while !matches!(self.peek(), Some(Token::RBrace)) {
            let fname = self.expect_ident()?;
            self.expect(Token::Colon)?;
            let value = self.parse_expr()?;
            fields.push(FieldAssign { name: fname, value });
            if matches!(self.peek(), Some(Token::Comma)) {
                self.advance();
            }
        }
        self.expect(Token::RBrace)?;
        Ok((name, fields))
    }

    fn parse_expr(&mut self) -> Result<Expr, Error> {
        self.parse_or()
    }

    fn parse_or(&mut self) -> Result<Expr, Error> {
        let mut left = self.parse_and()?;
        while matches!(self.peek(), Some(Token::Or)) {
            self.advance();
            let right = self.parse_and()?;
            left = Expr::BinaryOp(BinaryOp::Or, Box::new(left), Box::new(right));
        }
        Ok(left)
    }

    fn parse_and(&mut self) -> Result<Expr, Error> {
        let mut left = self.parse_not()?;
        while matches!(self.peek(), Some(Token::And)) {
            self.advance();
            let right = self.parse_not()?;
            left = Expr::BinaryOp(BinaryOp::And, Box::new(left), Box::new(right));
        }
        Ok(left)
    }

    fn parse_not(&mut self) -> Result<Expr, Error> {
        if matches!(self.peek(), Some(Token::Not)) {
            self.advance();
            let inner = self.parse_not()?;
            Ok(Expr::UnaryOp(UnaryOp::Not, Box::new(inner)))
        } else {
            self.parse_comparison()
        }
    }

    fn parse_comparison(&mut self) -> Result<Expr, Error> {
        let left = self.parse_add()?;
        let op = match self.peek() {
            Some(Token::EqEq) => BinaryOp::Eq,
            Some(Token::NotEq) => BinaryOp::Ne,
            Some(Token::Lt) => BinaryOp::Lt,
            Some(Token::LtEq) => BinaryOp::Le,
            Some(Token::Gt) => BinaryOp::Gt,
            Some(Token::GtEq) => BinaryOp::Ge,
            _ => return Ok(left),
        };
        self.advance();
        let right = self.parse_add()?;
        Ok(Expr::BinaryOp(op, Box::new(left), Box::new(right)))
    }

    fn parse_add(&mut self) -> Result<Expr, Error> {
        let mut left = self.parse_unary_neg()?;
        loop {
            let op = match self.peek() {
                Some(Token::Plus) => BinaryOp::Add,
                Some(Token::Minus) => BinaryOp::Sub,
                _ => break,
            };
            self.advance();
            let right = self.parse_unary_neg()?;
            left = Expr::BinaryOp(op, Box::new(left), Box::new(right));
        }
        Ok(left)
    }

    fn parse_unary_neg(&mut self) -> Result<Expr, Error> {
        if matches!(self.peek(), Some(Token::Minus)) {
            self.advance();
            let inner = self.parse_primary()?;
            Ok(Expr::UnaryOp(UnaryOp::Neg, Box::new(inner)))
        } else {
            self.parse_primary()
        }
    }

    fn parse_primary(&mut self) -> Result<Expr, Error> {
        match self.peek().cloned() {
            Some(Token::Int(n)) => {
                self.advance();
                Ok(Expr::LitInt(n))
            }
            Some(Token::Float(f)) => {
                self.advance();
                Ok(Expr::LitFloat(f))
            }
            Some(Token::Duration(ms)) => {
                self.advance();
                Ok(Expr::LitDuration(ms))
            }
            Some(Token::String(s)) => {
                self.advance();
                Ok(Expr::LitString(s))
            }
            Some(Token::True) => {
                self.advance();
                Ok(Expr::LitBool(true))
            }
            Some(Token::False) => {
                self.advance();
                Ok(Expr::LitBool(false))
            }
            Some(Token::Null) => {
                self.advance();
                Ok(Expr::LitNull)
            }
            Some(Token::LParen) => {
                self.advance();
                let inner = self.parse_expr()?;
                self.expect(Token::RParen)?;
                Ok(inner)
            }
            Some(Token::LBracket) => {
                self.advance();
                if matches!(self.peek(), Some(Token::RBracket)) {
                    self.advance();
                    return Ok(Expr::ListLit(Vec::new()));
                }
                let first = self.parse_expr()?;
                if matches!(self.peek(), Some(Token::For)) {
                    self.advance();
                    let var = self.expect_ident()?;
                    self.expect(Token::In)?;
                    let src = self.parse_expr()?;
                    let where_clause = if matches!(self.peek(), Some(Token::Where)) {
                        self.advance();
                        Some(Box::new(self.parse_expr()?))
                    } else {
                        None
                    };
                    self.expect(Token::RBracket)?;
                    Ok(Expr::ListComp {
                        var,
                        src: Box::new(src),
                        body: Box::new(first),
                        where_clause,
                    })
                } else {
                    let mut items = vec![first];
                    while matches!(self.peek(), Some(Token::Comma)) {
                        self.advance();
                        if matches!(self.peek(), Some(Token::RBracket)) {
                            break;
                        }
                        items.push(self.parse_expr()?);
                    }
                    self.expect(Token::RBracket)?;
                    Ok(Expr::ListLit(items))
                }
            }
            Some(Token::Ident(name)) => {
                self.advance();
                self.parse_postfix(name)
            }
            Some(other) => Err(Error::Parse(format!(
                "unexpected token in expression: {other:?}"
            ))),
            None => Err(Error::Parse("unexpected end of input in expression".into())),
        }
    }

    fn parse_postfix(&mut self, head: String) -> Result<Expr, Error> {
        // possibilities:
        //  IDENT '(' args ')'      -> Call (only top-level idents, not paths)
        //  IDENT '{' fields '}'    -> StructLit
        //  IDENT ('.' IDENT)+      -> FieldAccess chain
        //  IDENT                   -> Ident
        match self.peek() {
            Some(Token::LParen) => {
                self.advance();
                let mut args = Vec::new();
                while !matches!(self.peek(), Some(Token::RParen)) {
                    args.push(self.parse_expr()?);
                    if matches!(self.peek(), Some(Token::Comma)) {
                        self.advance();
                    }
                }
                self.expect(Token::RParen)?;
                Ok(Expr::Call(head, args))
            }
            Some(Token::LBrace) => {
                // re-enter struct_lit with name already consumed
                self.expect(Token::LBrace)?;
                let mut fields = Vec::new();
                while !matches!(self.peek(), Some(Token::RBrace)) {
                    let fname = self.expect_ident()?;
                    self.expect(Token::Colon)?;
                    let value = self.parse_expr()?;
                    fields.push(FieldAssign { name: fname, value });
                    if matches!(self.peek(), Some(Token::Comma)) {
                        self.advance();
                    }
                }
                self.expect(Token::RBrace)?;
                Ok(Expr::StructLit(head, fields))
            }
            Some(Token::Dot) => {
                let mut expr = Expr::Ident(head);
                while matches!(self.peek(), Some(Token::Dot)) {
                    self.advance();
                    let field = self.expect_ident()?;
                    expr = Expr::FieldAccess(Box::new(expr), field);
                }
                Ok(expr)
            }
            _ => Ok(Expr::Ident(head)),
        }
    }
}