use crate::ast::{
    Block, Expr, FieldDecl, FunctionDecl, Item, Module, Param, Stmt, StructDecl, TypeRef, UseDecl,
};
use crate::token::{Keyword, Span, Token, TokenKind};

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

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

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

    pub fn parse_module(&mut self) -> Result<Module, ParseError> {
        let mut items = Vec::new();
        while !self.is_at_end() {
            self.skip_newlines();
            if self.is_at_end() {
                break;
            }
            items.push(self.parse_item()?);
            self.skip_newlines();
        }
        Ok(Module { items })
    }

    fn parse_item(&mut self) -> Result<Item, ParseError> {
        if self.matches_keyword(Keyword::Use) {
            return self.parse_use().map(Item::Use);
        }

        let is_public = self.matches_keyword(Keyword::Pub);
        let is_async = self.matches_keyword(Keyword::Async);

        if self.matches_keyword(Keyword::Fn) {
            return self.parse_function(is_public, is_async).map(Item::Function);
        }

        if self.matches_keyword(Keyword::Struct) {
            return self.parse_struct(is_public).map(Item::Struct);
        }

        Err(self.error_here("expected module item"))
    }

    fn parse_use(&mut self) -> Result<UseDecl, ParseError> {
        let start = self.previous_span();
        let mut path = Vec::new();
        path.push(self.expect_identifier()?);
        while self.matches(TokenKind::Dot) {
            path.push(self.expect_identifier()?);
        }
        Ok(UseDecl { path, span: start })
    }

    fn parse_function(
        &mut self,
        is_public: bool,
        is_async: bool,
    ) -> Result<FunctionDecl, ParseError> {
        let start = self.previous_span();
        let name = self.expect_identifier()?;
        self.expect(TokenKind::LeftParen, "expected '(' after function name")?;
        let mut params = Vec::new();
        if !self.check(&TokenKind::RightParen) {
            loop {
                let param_name = self.expect_identifier()?;
                self.expect(TokenKind::Colon, "expected ':' after parameter name")?;
                let ty = self.parse_type()?;
                params.push(Param {
                    name: param_name,
                    ty,
                    span: start,
                });
                if !self.matches(TokenKind::Comma) {
                    break;
                }
            }
        }
        self.expect(TokenKind::RightParen, "expected ')' after parameters")?;

        let return_type = if self.matches(TokenKind::Arrow) {
            Some(self.parse_type()?)
        } else {
            None
        };

        self.expect(TokenKind::Colon, "expected ':' before function body")?;
        self.skip_newlines();
        let body = self.parse_block()?;

        Ok(FunctionDecl {
            is_public,
            is_async,
            name,
            params,
            return_type,
            body,
            span: start,
        })
    }

    fn parse_struct(&mut self, is_public: bool) -> Result<StructDecl, ParseError> {
        let start = self.previous_span();
        let name = self.expect_identifier()?;
        self.expect(TokenKind::Colon, "expected ':' after struct name")?;
        self.skip_newlines();
        self.expect(TokenKind::Indent, "expected indented struct body")?;

        let mut fields = Vec::new();
        while !self.check(&TokenKind::Dedent) && !self.is_at_end() {
            self.skip_newlines();
            if self.check(&TokenKind::Dedent) {
                break;
            }
            let field_name = self.expect_identifier()?;
            self.expect(TokenKind::Colon, "expected ':' after field name")?;
            let ty = self.parse_type()?;
            fields.push(FieldDecl {
                name: field_name,
                ty,
                span: start,
            });
            self.skip_newlines();
        }

        self.expect(TokenKind::Dedent, "expected end of struct body")?;

        Ok(StructDecl {
            is_public,
            name,
            fields,
            span: start,
        })
    }

    fn parse_block(&mut self) -> Result<Block, ParseError> {
        self.expect(TokenKind::Indent, "expected indented block")?;
        let mut statements = Vec::new();

        while !self.check(&TokenKind::Dedent) && !self.is_at_end() {
            self.skip_newlines();
            if self.check(&TokenKind::Dedent) {
                break;
            }
            statements.push(self.parse_statement()?);
            self.skip_newlines();
        }

        self.expect(TokenKind::Dedent, "expected end of block")?;
        Ok(Block { statements })
    }

    fn parse_statement(&mut self) -> Result<Stmt, ParseError> {
        if self.matches_keyword(Keyword::Let) {
            return self.parse_let(false);
        }
        if self.matches_keyword(Keyword::Var) {
            return self.parse_let(true);
        }
        if self.matches_keyword(Keyword::Return) {
            let span = self.previous_span();
            if self.check(&TokenKind::Newline) || self.check(&TokenKind::Dedent) {
                return Ok(Stmt::Return(None, span));
            }
            let expr = self.parse_expression()?;
            return Ok(Stmt::Return(Some(expr), span));
        }
        Ok(Stmt::Expr(self.parse_expression()?))
    }

    fn parse_let(&mut self, mutable: bool) -> Result<Stmt, ParseError> {
        let span = self.previous_span();
        let name = self.expect_identifier()?;
        let ty = if self.matches(TokenKind::Colon) {
            Some(self.parse_type()?)
        } else {
            None
        };
        self.expect(TokenKind::Equal, "expected '=' in variable declaration")?;
        let value = self.parse_expression()?;
        Ok(Stmt::Let {
            mutable,
            name,
            ty,
            value,
            span,
        })
    }

    fn parse_expression(&mut self) -> Result<Expr, ParseError> {
        let mut expr = self.parse_primary()?;
        while self.matches(TokenKind::LeftParen) {
            let mut args = Vec::new();
            if !self.check(&TokenKind::RightParen) {
                loop {
                    args.push(self.parse_expression()?);
                    if !self.matches(TokenKind::Comma) {
                        break;
                    }
                }
            }
            let span = self.previous_span();
            self.expect(TokenKind::RightParen, "expected ')' after arguments")?;
            expr = Expr::Call {
                callee: Box::new(expr),
                args,
                span,
            };
        }
        Ok(expr)
    }

    fn parse_primary(&mut self) -> Result<Expr, ParseError> {
        let token = self.advance().clone();
        match token.kind {
            TokenKind::Identifier(value) => Ok(Expr::Identifier(value, token.span)),
            TokenKind::Integer(value) => {
                let parsed = value.parse::<i64>().unwrap_or_default();
                Ok(Expr::Integer(parsed, token.span))
            }
            TokenKind::String(value) => Ok(Expr::String(value, token.span)),
            _ => Err(ParseError {
                message: "expected expression".to_string(),
                span: token.span,
            }),
        }
    }

    fn parse_type(&mut self) -> Result<TypeRef, ParseError> {
        let token = self.advance().clone();
        match token.kind {
            TokenKind::Identifier(name) => Ok(TypeRef {
                name,
                span: token.span,
            }),
            _ => Err(ParseError {
                message: "expected type name".to_string(),
                span: token.span,
            }),
        }
    }

    fn expect_identifier(&mut self) -> Result<String, ParseError> {
        let token = self.advance().clone();
        match token.kind {
            TokenKind::Identifier(name) => Ok(name),
            _ => Err(ParseError {
                message: "expected identifier".to_string(),
                span: token.span,
            }),
        }
    }

    fn expect(&mut self, kind: TokenKind, message: &str) -> Result<(), ParseError> {
        if self.matches(kind) {
            Ok(())
        } else {
            Err(self.error_here(message))
        }
    }

    fn matches_keyword(&mut self, keyword: Keyword) -> bool {
        if matches!(self.peek().kind, TokenKind::Keyword(found) if found == keyword) {
            self.advance();
            return true;
        }
        false
    }

    fn matches(&mut self, kind: TokenKind) -> bool {
        if self.check(&kind) {
            self.advance();
            return true;
        }
        false
    }

    fn check(&self, kind: &TokenKind) -> bool {
        if self.is_at_end() {
            return matches!(kind, TokenKind::Eof);
        }
        same_variant(&self.peek().kind, kind)
    }

    fn skip_newlines(&mut self) {
        while self.matches(TokenKind::Newline) {}
    }

    fn is_at_end(&self) -> bool {
        matches!(self.peek().kind, TokenKind::Eof)
    }

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

    fn advance(&mut self) -> &Token {
        if !self.is_at_end() {
            self.current += 1;
        }
        &self.tokens[self.current.saturating_sub(1)]
    }

    fn previous_span(&self) -> Span {
        if self.current == 0 {
            Span::default()
        } else {
            self.tokens[self.current - 1].span
        }
    }

    fn error_here(&self, message: &str) -> ParseError {
        ParseError {
            message: message.to_string(),
            span: self.peek().span,
        }
    }
}

fn same_variant(left: &TokenKind, right: &TokenKind) -> bool {
    std::mem::discriminant(left) == std::mem::discriminant(right)
}