// FILE: src/parser/core/primary.rs

use crate::parser::ast::Expr;
use crate::parser::lexer::{Token, TokenKind};
use super::utils::Parser;
use super::precedence::parse_expression;
use crate::parser::types::LambdaParam;

fn pos(parser: &Parser) -> (usize, usize) {
    parser.current().map(|t| (t.line, t.col)).unwrap_or((0, 0))
}

pub fn parse_primary(parser: &mut Parser) -> Result<Expr, super::driver::ParseError> {
    let (line, col) = pos(parser);
    let tk = parser.current().ok_or_else(|| super::driver::ParseError {
        msg: "EOF in parse_primary".into(),
        line,
        col,
    })?;

    match &tk.kind {
        TokenKind::Identifier(id) => {
            if peek_arrow(parser) {
                let param = if id == "$" { LambdaParam::Variadic } else { LambdaParam::Named(id.clone()) };
                parser.advance(); // Identifier
                parser.advance(); // =>
                if let Some(Token { kind: TokenKind::LBrace, .. }) = parser.current() {
                    let block = parse_block_expr(parser)?;
                    Ok(Expr::InlineLambdaAST {
                        params: vec![param],
                        body: Box::new(block),
                        line,
                        col,
                    })
                } else {
                    let body = parse_expression(parser)?;
                    Ok(Expr::InlineLambdaAST {
                        params: vec![param],
                        body: Box::new(body),
                        line,
                        col,
                    })
                }
            } else {
                let name = id.clone();
                parser.advance();
                Ok(Expr::Ident {
                    name,
                    line,
                    col,
                })
            }
        }
        TokenKind::Placeholder => {
            if peek_arrow(parser) {
                parser.advance(); // _
                parser.advance(); // =>
                if let Some(Token { kind: TokenKind::LBrace, .. }) = parser.current() {
                    let block = parse_block_expr(parser)?;
                    Ok(Expr::InlineLambdaAST {
                        params: vec![LambdaParam::Named("_".into())],
                        body: Box::new(block),
                        line,
                        col,
                    })
                } else {
                    let body = parse_expression(parser)?;
                    Ok(Expr::InlineLambdaAST {
                        params: vec![LambdaParam::Named("_".into())],
                        body: Box::new(body),
                        line,
                        col,
                    })
                }
            } else {
                parser.advance();
                Ok(Expr::Placeholder { line, col })
            }
        }
        TokenKind::Minus => {
            parser.advance();
            if let Some(Token { kind: TokenKind::Number(n), line: l, col: c }) = parser.current() {
                let v = -*n;
                parser.advance();
                Ok(Expr::Number { value: v, line: *l, col: *c })
            } else if let Some(Token { kind: TokenKind::FloatLiteral(f), line: l, col: c }) = parser.current() {
                let v = -*f;
                parser.advance();
                Ok(Expr::NumberFloat { value: v, line: *l, col: *c })
            } else {
                let sub = parse_primary(parser)?;
                Ok(Expr::UnaryMinus { expr: Box::new(sub), line, col })
            }
        }
        TokenKind::Bang => {
            parser.advance();
            let sub = parse_primary(parser)?;
            Ok(Expr::UnaryNot { expr: Box::new(sub), line, col })
        }
        TokenKind::Amp => {
            parser.advance();
            let sub = parse_primary(parser)?;
            Ok(Expr::Ref { expr: Box::new(sub), line, col })
        }
        TokenKind::Number(n) => {
            let v = *n;
            parser.advance();
            Ok(Expr::Number { value: v, line, col })
        }
        TokenKind::FloatLiteral(f) => {
            let v = *f;
            parser.advance();
            Ok(Expr::NumberFloat { value: v, line, col })
        }
        TokenKind::StringLiteral(s) => {
            let v = s.clone();
            parser.advance();
            Ok(Expr::StrLit { value: v, line, col })
        }
        TokenKind::BoolLiteral(b) => {
            let v = *b;
            parser.advance();
            Ok(Expr::Bool { value: v, line, col })
        }
        TokenKind::LBracket => parse_bracketed_array(parser),
        TokenKind::LParen => {
            parser.advance();
            let mut params = Vec::new();
            let start = parser.pos;

            while !parser.is_done() {
                if let Some(Token { kind: TokenKind::RParen, .. }) = parser.current() {
                    break;
                }

                match parser.current() {
                    Some(Token { kind: TokenKind::Amp, .. }) => {
                        parser.advance();
                        if let Some(Token { kind: TokenKind::Identifier(name), .. }) = parser.current() {
                            params.push(LambdaParam::Ref(name.clone()));
                            parser.advance();
                        } else {
                            return Err(super::driver::ParseError {
                                msg: "Expected identifier after '&' in param list".into(),
                                line: pos(parser).0,
                                col: pos(parser).1,
                            });
                        }
                    }
                    Some(Token { kind: TokenKind::Identifier(id), .. }) => {
                        if id == "$" {
                            params.push(LambdaParam::Variadic);
                        } else {
                            params.push(LambdaParam::Named(id.clone()));
                        }
                        parser.advance();
                    }
                    Some(Token { kind: TokenKind::Placeholder, .. }) => {
                        params.push(LambdaParam::Named("_".into()));
                        parser.advance();
                    }
                    _ => {
                        params.clear();
                        break;
                    }
                }

                if let Some(Token { kind: TokenKind::Comma, .. }) = parser.current() {
                    parser.advance();
                } else {
                    break;
                }
            }

            let mut is_lambda = false;
            if let Some(Token { kind: TokenKind::RParen, .. }) = parser.current() {
                parser.advance();
                if let Some(Token { kind: TokenKind::Arrow, .. }) = parser.current() {
                    is_lambda = true;
                    parser.advance();
                } else if !params.is_empty() {
                    parser.pos = start;
                    params.clear();
                }
            } else {
                parser.pos = start;
                params.clear();
            }

            if is_lambda {
                if let Some(Token { kind: TokenKind::LBrace, .. }) = parser.current() {
                    let block = parse_block_expr(parser)?;
                    Ok(Expr::InlineLambdaAST {
                        params,
                        body: Box::new(block),
                        line,
                        col,
                    })
                } else {
                    let body = parse_expression(parser)?;
                    Ok(Expr::InlineLambdaAST {
                        params,
                        body: Box::new(body),
                        line,
                        col,
                    })
                }
            } else {
                let inner = parse_expression(parser)?;
                parser.expect(&TokenKind::RParen)?;
                Ok(inner)
            }
        }
        other => Err(super::driver::ParseError {
            msg: format!("Unexpected token in primary: {:?}", other),
            line,
            col,
        }),
    }
}

pub fn peek_arrow(parser: &Parser) -> bool {
    parser.pos + 1 < parser.tokens.len()
        && matches!(parser.tokens[parser.pos + 1].kind, TokenKind::Arrow)
}

pub fn parse_bracketed_array(parser: &mut Parser) -> Result<Expr, super::driver::ParseError> {
    let (line, col) = pos(parser);
    parser.expect(&TokenKind::LBracket)?;
    let mut items = Vec::new();
    let mut keyed = Vec::new();
    let mut is_keyed = false;
    let mut first = true;
    while !parser.is_done() {
        if let Some(Token { kind: TokenKind::RBracket, .. }) = parser.current() {
            parser.advance();
            return if is_keyed {
                Ok(Expr::KeyedArray {
                    pairs: keyed,
                    line,
                    col,
                })
            } else {
                Ok(Expr::BracketedArray {
                    items,
                    line,
                    col,
                })
            };
        }
        if !first {
            if let Some(Token { kind: TokenKind::Comma, .. }) = parser.current() {
                parser.advance();
            }
        }
        first = false;

        let save = parser.pos;
        if let Some(Token { kind: TokenKind::Identifier(id), .. }) = parser.current() {
            let possible = id.clone();
            parser.advance();
            if let Some(Token { kind: TokenKind::Colon, .. }) = parser.current() {
                parser.advance();
                let val = parse_expression(parser)?;
                keyed.push((possible, val));
                is_keyed = true;
                continue;
            } else {
                parser.pos = save;
            }
        }
        let v = parse_expression(parser)?;
        if is_keyed {
            return Err(super::driver::ParseError {
                msg: "Mix of keyed and non-keyed elements".into(),
                line,
                col,
            });
        }
        items.push(v);
    }
    Err(super::driver::ParseError {
        msg: "Unclosed '['".into(),
        line,
        col,
    })
}

pub fn parse_block_expr(parser: &mut Parser) -> Result<Expr, super::driver::ParseError> {
    let (line, col) = pos(parser);
    parser.expect(&TokenKind::LBrace)?;
    let mut stmts = Vec::new();
    while !parser.is_done() {
        while let Some(Token { kind: TokenKind::Semi, .. }) = parser.current() {
            parser.advance();
        }
        if let Some(Token { kind: TokenKind::RBrace, .. }) = parser.current() {
            parser.advance();
            break;
        }
        stmts.push(super::stmt::parse_statement(parser)?);
    }
    Ok(Expr::Block {
        stmts,
        line,
        col,
    })
}
