use crate::ast::*;
use crate::lexer::{SpannedToken, Token};

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

#[derive(Debug, Clone)]
pub struct ParseError {
    pub span: Span,
    pub msg: String,
}

impl std::fmt::Display for ParseError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "parse error at {}..{}: {}", self.span.start, self.span.end, self.msg)
    }
}

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

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

    pub fn parse_program(&mut self) -> Result<Program, ParseError> {
        let mut statements = Vec::new();
        while !self.at_eof() {
            statements.push(self.parse_statement()?);
        }
        Ok(Program { statements })
    }

    fn parse_statement(&mut self) -> Result<Statement, ParseError> {
        // Variable assignment: ident = expr;
        if self.peek_is_ident() && self.peek_ahead_is(1, &Token::Assign) {
            return self.parse_assignment();
        }

        // Module definition
        if self.peek_is(&Token::Module) {
            return self.parse_module_def();
        }

        // For loop
        if self.peek_is(&Token::For) {
            return self.parse_for_loop();
        }

        // If/else
        if self.peek_is(&Token::If) {
            return self.parse_if_else();
        }

        // Boolean ops or module calls
        self.parse_module_call_or_boolean()
    }

    fn parse_assignment(&mut self) -> Result<Statement, ParseError> {
        let start = self.current_span();
        let name = self.expect_ident()?;
        self.expect(&Token::Assign)?;
        let value = self.parse_expr()?;
        self.expect(&Token::Semi)?;
        Ok(Statement::Assignment(Assignment {
            name,
            value,
            span: Span { start: start.start, end: self.prev_span().end },
        }))
    }

    fn parse_module_def(&mut self) -> Result<Statement, ParseError> {
        let start = self.current_span();
        self.expect(&Token::Module)?;
        let name = self.expect_ident()?;
        self.expect(&Token::LParen)?;
        let params = self.parse_params()?;
        self.expect(&Token::RParen)?;
        let body = self.parse_block()?;
        Ok(Statement::ModuleDef(ModuleDef {
            name,
            params,
            body,
            span: Span { start: start.start, end: self.prev_span().end },
        }))
    }

    fn parse_params(&mut self) -> Result<Vec<Param>, ParseError> {
        let mut params = Vec::new();
        while !self.peek_is(&Token::RParen) && !self.at_eof() {
            let name = self.expect_ident()?;
            let default = if self.try_consume(&Token::Assign) {
                Some(self.parse_expr()?)
            } else {
                None
            };
            params.push(Param { name, default });
            if !self.try_consume(&Token::Comma) {
                break;
            }
        }
        Ok(params)
    }

    fn parse_module_call_or_boolean(&mut self) -> Result<Statement, ParseError> {
        let start = self.current_span();
        let name = self.expect_ident()?;

        // Check for boolean operations
        let boolean_kind = match name.as_str() {
            "union" => Some(BooleanKind::Union),
            "difference" => Some(BooleanKind::Difference),
            "intersection" => Some(BooleanKind::Intersection),
            _ => None,
        };

        self.expect(&Token::LParen)?;
        let args = self.parse_arguments()?;
        self.expect(&Token::RParen)?;

        // Children: either a block `{ ... }` or a single trailing statement, or `;`
        let children = if self.try_consume(&Token::Semi) {
            Vec::new()
        } else if self.peek_is(&Token::LBrace) {
            self.parse_block()?
        } else {
            vec![self.parse_statement()?]
        };

        let span = Span { start: start.start, end: self.prev_span().end };

        if let Some(op) = boolean_kind {
            Ok(Statement::BooleanOp(BooleanOp { op, children, span }))
        } else {
            Ok(Statement::ModuleCall(ModuleCall { name, args, children, span }))
        }
    }

    fn parse_block(&mut self) -> Result<Vec<Statement>, ParseError> {
        self.expect(&Token::LBrace)?;
        let mut stmts = Vec::new();
        while !self.peek_is(&Token::RBrace) && !self.at_eof() {
            stmts.push(self.parse_statement()?);
        }
        self.expect(&Token::RBrace)?;
        Ok(stmts)
    }

    fn parse_arguments(&mut self) -> Result<Vec<Argument>, ParseError> {
        let mut args = Vec::new();
        while !self.peek_is(&Token::RParen) && !self.at_eof() {
            // Try named argument: ident = expr
            let arg = if self.peek_is_ident() && self.peek_ahead_is(1, &Token::Assign) {
                let name = self.expect_ident()?;
                self.expect(&Token::Assign)?;
                let value = self.parse_expr()?;
                Argument { name: Some(name), value }
            } else {
                Argument { name: None, value: self.parse_expr()? }
            };
            args.push(arg);
            if !self.try_consume(&Token::Comma) {
                break;
            }
        }
        Ok(args)
    }

    // Expression parsing with precedence climbing
    fn parse_expr(&mut self) -> Result<Expr, ParseError> {
        self.parse_ternary()
    }

    fn parse_for_loop(&mut self) -> Result<Statement, ParseError> {
        let start = self.current_span();
        self.expect(&Token::For)?;
        self.expect(&Token::LParen)?;
        let var = self.expect_ident()?;
        self.expect(&Token::Assign)?;

        // Range: [start : end], [start : step : end], or [a, b, c]
        self.expect(&Token::LBracket)?;
        let range = self.parse_for_range()?;
        self.expect(&Token::RBracket)?;
        self.expect(&Token::RParen)?;

        let body = if self.peek_is(&Token::LBrace) {
            self.parse_block()?
        } else {
            vec![self.parse_statement()?]
        };

        Ok(Statement::ForLoop(ForLoop {
            var,
            range,
            body,
            span: Span { start: start.start, end: self.prev_span().end },
        }))
    }

    fn parse_for_range(&mut self) -> Result<ForRange, ParseError> {
        let first = self.parse_expr()?;

        if self.try_consume(&Token::Colon) {
            let second = self.parse_expr()?;
            if self.try_consume(&Token::Colon) {
                // [start : step : end]
                let third = self.parse_expr()?;
                Ok(ForRange::Range { start: first, step: Some(second), end: third })
            } else {
                // [start : end]
                Ok(ForRange::Range { start: first, step: None, end: second })
            }
        } else {
            // Explicit list: [a, b, c, ...]
            let mut list = vec![first];
            while self.try_consume(&Token::Comma) {
                list.push(self.parse_expr()?);
            }
            Ok(ForRange::List(list))
        }
    }

    fn parse_if_else(&mut self) -> Result<Statement, ParseError> {
        let start = self.current_span();
        self.expect(&Token::If)?;
        self.expect(&Token::LParen)?;
        let condition = self.parse_expr()?;
        self.expect(&Token::RParen)?;

        let then_body = if self.peek_is(&Token::LBrace) {
            self.parse_block()?
        } else {
            vec![self.parse_statement()?]
        };

        let else_body = if self.try_consume(&Token::Else) {
            if self.peek_is(&Token::LBrace) {
                self.parse_block()?
            } else {
                vec![self.parse_statement()?]
            }
        } else {
            Vec::new()
        };

        Ok(Statement::IfElse(IfElse {
            condition,
            then_body,
            else_body,
            span: Span { start: start.start, end: self.prev_span().end },
        }))
    }

    fn parse_ternary(&mut self) -> Result<Expr, ParseError> {
        let expr = self.parse_or()?;
        if self.try_consume(&Token::Question) {
            let then_expr = self.parse_expr()?;
            self.expect(&Token::Colon)?;
            let else_expr = self.parse_expr()?;
            Ok(Expr::Ternary {
                cond: Box::new(expr),
                then_expr: Box::new(then_expr),
                else_expr: Box::new(else_expr),
            })
        } else {
            Ok(expr)
        }
    }

    fn parse_or(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_and()?;
        while self.try_consume(&Token::Or) {
            let right = self.parse_and()?;
            left = Expr::BinaryOp {
                op: BinaryOp::Or,
                left: Box::new(left),
                right: Box::new(right),
            };
        }
        Ok(left)
    }

    fn parse_and(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_equality()?;
        while self.try_consume(&Token::And) {
            let right = self.parse_equality()?;
            left = Expr::BinaryOp {
                op: BinaryOp::And,
                left: Box::new(left),
                right: Box::new(right),
            };
        }
        Ok(left)
    }

    fn parse_equality(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_comparison()?;
        loop {
            let op = if self.try_consume(&Token::EqEq) {
                BinaryOp::Eq
            } else if self.try_consume(&Token::BangEq) {
                BinaryOp::Ne
            } else {
                break;
            };
            let right = self.parse_comparison()?;
            left = Expr::BinaryOp { op, left: Box::new(left), right: Box::new(right) };
        }
        Ok(left)
    }

    fn parse_comparison(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_additive()?;
        loop {
            let op = if self.try_consume(&Token::Lt) {
                BinaryOp::Lt
            } else if self.try_consume(&Token::Le) {
                BinaryOp::Le
            } else if self.try_consume(&Token::Gt) {
                BinaryOp::Gt
            } else if self.try_consume(&Token::Ge) {
                BinaryOp::Ge
            } else {
                break;
            };
            let right = self.parse_additive()?;
            left = Expr::BinaryOp { op, left: Box::new(left), right: Box::new(right) };
        }
        Ok(left)
    }

    fn parse_additive(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_multiplicative()?;
        loop {
            let op = if self.try_consume(&Token::Plus) {
                BinaryOp::Add
            } else if self.try_consume(&Token::Minus) {
                BinaryOp::Sub
            } else {
                break;
            };
            let right = self.parse_multiplicative()?;
            left = Expr::BinaryOp { op, left: Box::new(left), right: Box::new(right) };
        }
        Ok(left)
    }

    fn parse_multiplicative(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_unary()?;
        loop {
            let op = if self.try_consume(&Token::Star) {
                BinaryOp::Mul
            } else if self.try_consume(&Token::Slash) {
                BinaryOp::Div
            } else if self.try_consume(&Token::Percent) {
                BinaryOp::Mod
            } else {
                break;
            };
            let right = self.parse_unary()?;
            left = Expr::BinaryOp { op, left: Box::new(left), right: Box::new(right) };
        }
        Ok(left)
    }

    fn parse_unary(&mut self) -> Result<Expr, ParseError> {
        if self.try_consume(&Token::Minus) {
            let operand = self.parse_unary()?;
            return Ok(Expr::UnaryOp { op: UnaryOp::Neg, operand: Box::new(operand) });
        }
        if self.try_consume(&Token::Bang) {
            let operand = self.parse_unary()?;
            return Ok(Expr::UnaryOp { op: UnaryOp::Not, operand: Box::new(operand) });
        }
        self.parse_postfix()
    }

    fn parse_postfix(&mut self) -> Result<Expr, ParseError> {
        let mut expr = self.parse_primary()?;
        while self.peek_is(&Token::LBracket) {
            self.expect(&Token::LBracket)?;
            let index = self.parse_expr()?;
            self.expect(&Token::RBracket)?;
            expr = Expr::Index { array: Box::new(expr), index: Box::new(index) };
        }
        Ok(expr)
    }

    fn parse_primary(&mut self) -> Result<Expr, ParseError> {
        let tok = &self.tokens[self.pos];
        match &tok.token {
            Token::Number(n) => {
                let n = *n;
                self.pos += 1;
                Ok(Expr::Number(n))
            }
            Token::True => {
                self.pos += 1;
                Ok(Expr::Bool(true))
            }
            Token::False => {
                self.pos += 1;
                Ok(Expr::Bool(false))
            }
            Token::StringLit(s) => {
                let s = s.clone();
                self.pos += 1;
                Ok(Expr::String(s))
            }
            Token::Ident(_) => {
                let name = self.expect_ident()?;
                // Function call
                if self.peek_is(&Token::LParen) {
                    self.expect(&Token::LParen)?;
                    let args = self.parse_arguments()?;
                    self.expect(&Token::RParen)?;
                    return Ok(Expr::FnCall { name, args });
                }
                Ok(Expr::Ident(name))
            }
            Token::LBracket => {
                self.pos += 1;
                let mut elems = Vec::new();
                while !self.peek_is(&Token::RBracket) && !self.at_eof() {
                    elems.push(self.parse_expr()?);
                    if !self.try_consume(&Token::Comma) {
                        break;
                    }
                }
                self.expect(&Token::RBracket)?;
                Ok(Expr::Vector(elems))
            }
            Token::LParen => {
                self.pos += 1;
                let expr = self.parse_expr()?;
                self.expect(&Token::RParen)?;
                Ok(expr)
            }
            _ => Err(ParseError {
                span: tok.span,
                msg: format!("unexpected token: {:?}", tok.token),
            }),
        }
    }

    // Utility methods

    fn at_eof(&self) -> bool {
        self.pos >= self.tokens.len() || self.tokens[self.pos].token == Token::Eof
    }

    fn current_span(&self) -> Span {
        if self.pos < self.tokens.len() {
            self.tokens[self.pos].span
        } else {
            let end = self.tokens.last().map_or(0, |t| t.span.end);
            Span { start: end, end }
        }
    }

    fn prev_span(&self) -> Span {
        if self.pos > 0 {
            self.tokens[self.pos - 1].span
        } else {
            Span { start: 0, end: 0 }
        }
    }

    fn peek_is(&self, token: &Token) -> bool {
        self.pos < self.tokens.len() && std::mem::discriminant(&self.tokens[self.pos].token) == std::mem::discriminant(token)
    }

    fn peek_is_ident(&self) -> bool {
        matches!(self.tokens.get(self.pos), Some(SpannedToken { token: Token::Ident(_), .. }))
    }

    fn peek_ahead_is(&self, offset: usize, token: &Token) -> bool {
        let idx = self.pos + offset;
        idx < self.tokens.len()
            && std::mem::discriminant(&self.tokens[idx].token) == std::mem::discriminant(token)
    }

    fn try_consume(&mut self, token: &Token) -> bool {
        if self.peek_is(token) {
            self.pos += 1;
            true
        } else {
            false
        }
    }

    fn expect(&mut self, token: &Token) -> Result<(), ParseError> {
        if self.peek_is(token) {
            self.pos += 1;
            Ok(())
        } else {
            Err(ParseError {
                span: self.current_span(),
                msg: format!("expected {:?}, got {:?}", token, self.tokens.get(self.pos).map(|t| &t.token)),
            })
        }
    }

    fn expect_ident(&mut self) -> Result<String, ParseError> {
        if let Some(SpannedToken { token: Token::Ident(name), .. }) = self.tokens.get(self.pos) {
            let name = name.clone();
            self.pos += 1;
            Ok(name)
        } else {
            Err(ParseError {
                span: self.current_span(),
                msg: format!("expected identifier, got {:?}", self.tokens.get(self.pos).map(|t| &t.token)),
            })
        }
    }
}
