use alloc::{string::String, vec::Vec};

use crate::{
    SourceId, Span,
    ast::{Ast, BinaryOp, CompareOp, Expr, ExprId, Field, Literal, MatchArm, Param},
    diagnostic::{Diagnostic, Result},
    lexer::{Lexer, Token, TokenKind},
};

#[derive(Debug, Clone)]
pub struct ParseOutput {
    pub ast: Ast,
    pub root: ExprId,
    pub source_form: SourceForm,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SourceForm {
    Expr,
    Fields,
}

pub fn parse_source(source: &str) -> Result<ParseOutput> {
    parse_source_with_source_id(SourceId(0), source)
}

pub fn parse_source_with_source_id(source_id: SourceId, source: &str) -> Result<ParseOutput> {
    let tokens = Lexer::with_source_id(source_id, source).tokenize()?;
    Parser::new(tokens).parse()
}

pub struct Parser {
    tokens: Vec<Token>,
    pos: usize,
    ast: Ast,
}

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

    pub fn parse(mut self) -> Result<ParseOutput> {
        let (root, source_form) = if self.starts_field() {
            let fields = self.parse_fields_until_eof()?;
            let span = fields
                .first()
                .map(|f| {
                    fields
                        .iter()
                        .fold(f.span, |acc, field| acc.join(field.span))
                })
                .unwrap_or_else(|| self.peek().span);
            (
                self.ast.push(Expr::Object(fields), span),
                SourceForm::Fields,
            )
        } else {
            let expr = self.parse_expr(0)?;
            self.expect_eof()?;
            (expr, SourceForm::Expr)
        };
        Ok(ParseOutput {
            ast: self.ast,
            root,
            source_form,
        })
    }

    fn parse_expr(&mut self, min_bp: u8) -> Result<ExprId> {
        let mut lhs = self.parse_prefix()?;

        loop {
            if self.at_eof() || self.is_expr_stop() {
                break;
            }

            // postfix: path reference
            if self.consume_kind(&TokenKind::Dot).is_some() {
                let (field, field_span) = self.expect_ident()?;
                let span = self.ast.span(lhs).join(field_span);
                lhs = self.ast.push(Expr::Path { base: lhs, field }, span);
                continue;
            }

            // postfix: call
            if self.consume_kind(&TokenKind::LParen).is_some() {
                let mut args = Vec::new();
                if self.consume_kind(&TokenKind::RParen).is_none() {
                    loop {
                        args.push(self.parse_expr(0)?);
                        if self.consume_kind(&TokenKind::Comma).is_some() {
                            continue;
                        }
                        self.expect_kind(&TokenKind::RParen, "expected ')' after call arguments")?;
                        break;
                    }
                }
                let span = self.ast.span(lhs).join(self.previous_span());
                lhs = self.ast.push(Expr::Call { callee: lhs, args }, span);
                continue;
            }

            let Some((kind, l_bp, r_bp)) = self.peek_infix() else {
                break;
            };
            if l_bp < min_bp {
                break;
            }
            let op_span = self.advance().span;
            let rhs = self.parse_expr(r_bp)?;
            let span = self.ast.span(lhs).join(self.ast.span(rhs)).join(op_span);
            lhs = match kind {
                InfixKind::And => self.ast.push(
                    Expr::Binary {
                        op: BinaryOp::And,
                        lhs,
                        rhs,
                    },
                    span,
                ),
                InfixKind::Patch => self.ast.push(
                    Expr::Binary {
                        op: BinaryOp::Patch,
                        lhs,
                        rhs,
                    },
                    span,
                ),
                InfixKind::Default => self.ast.push(
                    Expr::Default {
                        base: lhs,
                        fallback: rhs,
                    },
                    span,
                ),
            };
        }

        Ok(lhs)
    }

    fn parse_prefix(&mut self) -> Result<ExprId> {
        let token = self.advance().clone();
        match token.kind {
            TokenKind::String(value) => Ok(self
                .ast
                .push(Expr::Literal(Literal::String(value)), token.span)),
            TokenKind::Int(value) => Ok(self
                .ast
                .push(Expr::Literal(Literal::Int(value)), token.span)),
            TokenKind::Float(value) => Ok(self
                .ast
                .push(Expr::Literal(Literal::Float(value)), token.span)),
            TokenKind::True => Ok(self
                .ast
                .push(Expr::Literal(Literal::Bool(true)), token.span)),
            TokenKind::False => Ok(self
                .ast
                .push(Expr::Literal(Literal::Bool(false)), token.span)),
            TokenKind::Ident(name) => Ok(self.ast.push(Expr::Ident(name), token.span)),
            TokenKind::Regex(pattern) => {
                Ok(self.ast.push(Expr::RegexConstraint(pattern), token.span))
            }
            TokenKind::Underscore => Ok(self.ast.push(Expr::Wildcard, token.span)),
            TokenKind::LBrace => self.parse_object_after_lbrace(token.span),
            TokenKind::LBracket => self.parse_array_after_lbracket(token.span),
            TokenKind::LParen => self.parse_group_or_function(token.span),
            TokenKind::Let => self.parse_let(token.span),
            TokenKind::Match => self.parse_match(token.span),
            TokenKind::Import => self.parse_import(token.span),
            TokenKind::Gt | TokenKind::Gte | TokenKind::Lt | TokenKind::Lte => {
                let op = match token.kind {
                    TokenKind::Gt => CompareOp::Gt,
                    TokenKind::Gte => CompareOp::Gte,
                    TokenKind::Lt => CompareOp::Lt,
                    TokenKind::Lte => CompareOp::Lte,
                    _ => unreachable!(),
                };
                let value = self.parse_expr(8)?;
                let span = token.span.join(self.ast.span(value));
                Ok(self.ast.push(Expr::CompareConstraint { op, value }, span))
            }
            _ => Err(Diagnostic::syntax(token.span, "expected expression")),
        }
    }

    fn parse_object_after_lbrace(&mut self, start_span: Span) -> Result<ExprId> {
        let mut fields = Vec::new();
        if self.consume_kind(&TokenKind::RBrace).is_some() {
            return Ok(self
                .ast
                .push(Expr::Object(fields), start_span.join(self.previous_span())));
        }
        loop {
            fields.push(self.parse_field()?);
            if self.consume_kind(&TokenKind::Semicolon).is_some() {
                if self.consume_kind(&TokenKind::RBrace).is_some() {
                    break;
                }
                continue;
            }
            self.expect_kind(&TokenKind::RBrace, "expected ';' or '}' after object field")?;
            break;
        }
        let span = start_span.join(self.previous_span());
        Ok(self.ast.push(Expr::Object(fields), span))
    }

    fn parse_array_after_lbracket(&mut self, start_span: Span) -> Result<ExprId> {
        let mut items = Vec::new();
        if self.consume_kind(&TokenKind::RBracket).is_some() {
            return Ok(self
                .ast
                .push(Expr::Array(items), start_span.join(self.previous_span())));
        }
        loop {
            items.push(self.parse_expr(0)?);
            if self.consume_kind(&TokenKind::Comma).is_some() {
                if self.consume_kind(&TokenKind::RBracket).is_some() {
                    break;
                }
                continue;
            }
            self.expect_kind(&TokenKind::RBracket, "expected ',' or ']' after array item")?;
            break;
        }
        let span = start_span.join(self.previous_span());
        Ok(self.ast.push(Expr::Array(items), span))
    }

    fn parse_group_or_function(&mut self, start_span: Span) -> Result<ExprId> {
        if self.looks_like_params() {
            let params = self.parse_params_after_lparen()?;
            self.expect_kind(&TokenKind::Arrow, "expected '=>' after function parameters")?;
            let body = self.parse_expr(0)?;
            let span = start_span.join(self.ast.span(body));
            return Ok(self.ast.push(Expr::Function { params, body }, span));
        }

        let expr = self.parse_expr(0)?;
        self.expect_kind(&TokenKind::RParen, "expected ')' after expression")?;
        Ok(expr)
    }

    fn parse_params_after_lparen(&mut self) -> Result<Vec<Param>> {
        let mut params = Vec::new();
        if self.consume_kind(&TokenKind::RParen).is_some() {
            return Ok(params);
        }
        loop {
            let (name, name_span) = self.expect_ident()?;
            let constraint = if self.consume_kind(&TokenKind::Colon).is_some() {
                Some(self.parse_expr(0)?)
            } else {
                None
            };
            let span = constraint
                .map(|id| name_span.join(self.ast.span(id)))
                .unwrap_or(name_span);
            params.push(Param {
                name,
                constraint,
                span,
            });
            if self.consume_kind(&TokenKind::Comma).is_some() {
                continue;
            }
            self.expect_kind(&TokenKind::RParen, "expected ',' or ')' after parameter")?;
            break;
        }
        Ok(params)
    }

    fn parse_let(&mut self, start_span: Span) -> Result<ExprId> {
        let mut bindings = Vec::new();
        while !self.check_kind(&TokenKind::In) && !self.at_eof() {
            bindings.push(self.parse_field()?);
            self.expect_kind(&TokenKind::Semicolon, "expected ';' after let binding")?;
        }
        self.expect_kind(&TokenKind::In, "expected 'in' after let bindings")?;
        let body = self.parse_expr(0)?;
        let span = start_span.join(self.ast.span(body));
        Ok(self.ast.push(Expr::Let { bindings, body }, span))
    }

    fn parse_match(&mut self, start_span: Span) -> Result<ExprId> {
        let scrutinee = self.parse_expr(0)?;
        self.expect_kind(&TokenKind::LBrace, "expected '{' after match scrutinee")?;
        let mut arms = Vec::new();
        if self.consume_kind(&TokenKind::RBrace).is_none() {
            loop {
                let pattern = self.parse_expr(0)?;
                self.expect_kind(&TokenKind::Colon, "expected ':' after match pattern")?;
                let body = self.parse_expr(0)?;
                let span = self.ast.span(pattern).join(self.ast.span(body));
                arms.push(MatchArm {
                    pattern,
                    body,
                    span,
                });
                if self.consume_kind(&TokenKind::Semicolon).is_some() {
                    if self.consume_kind(&TokenKind::RBrace).is_some() {
                        break;
                    }
                    continue;
                }
                self.expect_kind(&TokenKind::RBrace, "expected ';' or '}' after match arm")?;
                break;
            }
        }
        let span = start_span.join(self.previous_span());
        Ok(self.ast.push(Expr::Match { scrutinee, arms }, span))
    }

    fn parse_import(&mut self, start_span: Span) -> Result<ExprId> {
        let token = self.advance().clone();
        let path = match token.kind {
            TokenKind::String(path) | TokenKind::Ident(path) => path,
            _ => return Err(Diagnostic::syntax(token.span, "expected import path")),
        };
        Ok(self
            .ast
            .push(Expr::Import(path), start_span.join(token.span)))
    }

    fn parse_fields_until_eof(&mut self) -> Result<Vec<Field>> {
        let mut fields = Vec::new();
        while !self.at_eof() {
            fields.push(self.parse_field()?);
            self.consume_kind(&TokenKind::Semicolon);
        }
        Ok(fields)
    }

    fn parse_field(&mut self) -> Result<Field> {
        let (first, first_span) = self.expect_ident()?;
        let mut path = Vec::new();
        path.push(first);
        let mut span = first_span;
        while self.consume_kind(&TokenKind::Dot).is_some() {
            let (name, name_span) = self.expect_ident()?;
            span = span.join(name_span);
            path.push(name);
        }
        self.expect_kind(&TokenKind::Equal, "expected '=' after field name")?;
        let value = self.parse_expr(0)?;
        span = span.join(self.ast.span(value));
        Ok(Field { path, value, span })
    }

    fn starts_field(&self) -> bool {
        matches!(self.peek_kind(), TokenKind::Ident(_))
            && matches!(self.peek_kind_n(1), TokenKind::Equal | TokenKind::Dot)
    }

    fn looks_like_params(&self) -> bool {
        match self.peek_kind() {
            TokenKind::RParen => matches!(self.peek_kind_n(1), TokenKind::Arrow),
            TokenKind::Ident(_) => {
                let mut i = self.pos;
                loop {
                    if !matches!(self.kind_at(i), TokenKind::Ident(_)) {
                        return false;
                    }
                    i += 1;
                    if matches!(self.kind_at(i), TokenKind::Colon) {
                        // Skip a simple constraint expression approximately until comma/rparen.
                        i += 1;
                        let mut depth = 0usize;
                        while !matches!(self.kind_at(i), TokenKind::Eof) {
                            match self.kind_at(i) {
                                TokenKind::LParen | TokenKind::LBrace | TokenKind::LBracket => {
                                    depth += 1
                                }
                                TokenKind::RParen if depth == 0 => break,
                                TokenKind::RParen | TokenKind::RBrace | TokenKind::RBracket => {
                                    depth = depth.saturating_sub(1)
                                }
                                TokenKind::Comma if depth == 0 => break,
                                _ => {}
                            }
                            i += 1;
                        }
                    }
                    if matches!(self.kind_at(i), TokenKind::Comma) {
                        i += 1;
                        continue;
                    }
                    if matches!(self.kind_at(i), TokenKind::RParen) {
                        return matches!(self.kind_at(i + 1), TokenKind::Arrow);
                    }
                    return false;
                }
            }
            _ => false,
        }
    }

    fn peek_infix(&self) -> Option<(InfixKind, u8, u8)> {
        match self.peek_kind() {
            TokenKind::Default => Some((InfixKind::Default, 1, 2)),
            TokenKind::SlashSlash => Some((InfixKind::Patch, 3, 4)),
            TokenKind::Amp => Some((InfixKind::And, 5, 6)),
            _ => None,
        }
    }

    fn is_expr_stop(&self) -> bool {
        matches!(
            self.peek_kind(),
            TokenKind::Semicolon
                | TokenKind::Comma
                | TokenKind::RParen
                | TokenKind::RBracket
                | TokenKind::RBrace
                | TokenKind::Colon
                | TokenKind::In
                | TokenKind::Arrow
        )
    }

    fn expect_ident(&mut self) -> Result<(String, Span)> {
        let token = self.advance().clone();
        match token.kind {
            TokenKind::Ident(name) => Ok((name, token.span)),
            _ => Err(Diagnostic::syntax(token.span, "expected identifier")),
        }
    }

    fn expect_kind(&mut self, expected: &TokenKind, message: &'static str) -> Result<Span> {
        if let Some(span) = self.consume_kind(expected) {
            Ok(span)
        } else {
            Err(Diagnostic::syntax(self.peek().span, message))
        }
    }

    fn expect_eof(&mut self) -> Result<()> {
        if self.at_eof() {
            Ok(())
        } else {
            Err(Diagnostic::syntax(self.peek().span, "expected end of file"))
        }
    }

    fn consume_kind(&mut self, expected: &TokenKind) -> Option<Span> {
        if core::mem::discriminant(self.peek_kind()) == core::mem::discriminant(expected) {
            Some(self.advance().span)
        } else {
            None
        }
    }

    fn check_kind(&self, expected: &TokenKind) -> bool {
        core::mem::discriminant(self.peek_kind()) == core::mem::discriminant(expected)
    }

    fn at_eof(&self) -> bool {
        matches!(self.peek_kind(), TokenKind::Eof)
    }

    fn advance(&mut self) -> &Token {
        let index = self.pos;
        if !self.at_eof() {
            self.pos += 1;
        }
        &self.tokens[index]
    }

    fn previous_span(&self) -> Span {
        self.tokens[self.pos.saturating_sub(1)].span
    }

    fn peek(&self) -> &Token {
        &self.tokens[self.pos]
    }

    fn peek_kind(&self) -> &TokenKind {
        &self.peek().kind
    }

    fn peek_kind_n(&self, n: usize) -> &TokenKind {
        self.kind_at(self.pos + n)
    }

    fn kind_at(&self, index: usize) -> &TokenKind {
        self.tokens
            .get(index)
            .map(|token| &token.kind)
            .unwrap_or(&TokenKind::Eof)
    }
}

#[derive(Debug, Clone, Copy)]
enum InfixKind {
    And,
    Patch,
    Default,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::Expr;

    #[test]
    fn parses_top_level_fields_as_object() {
        let parsed = parse_source("port = Int & >= 1 default 8080;").unwrap();
        assert!(matches!(parsed.ast.get(parsed.root).expr, Expr::Object(_)));
    }

    #[test]
    fn parses_object_dot_field() {
        let parsed = parse_source("{ feature.enable = false; }").unwrap();
        let Expr::Object(fields) = &parsed.ast.get(parsed.root).expr else {
            panic!()
        };
        assert_eq!(fields[0].path, ["feature", "enable"]);
    }
}