#include "parser.h"
#include <format>

namespace camellya {

Parser::Parser(std::vector<Token> tokens) : tokens_(std::move(tokens)) {}

Program Parser::parse() {
    std::vector<StmtPtr> statements;
    
    while (!is_at_end()) {
        try {
            statements.push_back(declaration());
        } catch (const ParseError& error) {
            synchronize();
        }
    }
    
    return Program(std::move(statements));
}

Token Parser::peek() const {
    return tokens_[current_];
}

Token Parser::previous() const {
    return tokens_[current_ - 1];
}

bool Parser::is_at_end() const {
    return peek().type == TokenType::END_OF_FILE;
}

Token Parser::advance() {
    if (!is_at_end()) current_++;
    return previous();
}

bool Parser::check(TokenType type) const {
    if (is_at_end()) return false;
    return peek().type == type;
}

bool Parser::match(std::initializer_list<TokenType> types) {
    for (TokenType type : types) {
        if (check(type)) {
            advance();
            return true;
        }
    }
    return false;
}

Token Parser::consume(TokenType type, const std::string& message) {
    if (check(type)) return advance();
    throw error(peek(), message);
}

ParseError Parser::error(const Token& token, const std::string& message) {
    std::string error_msg = std::format("Line {}: Error at '{}': {}", 
                                        token.line, token.lexeme, message);
    return ParseError(error_msg);
}

void Parser::synchronize() {
    advance();
    
    while (!is_at_end()) {
        if (previous().type == TokenType::SEMICOLON) return;
        
        switch (peek().type) {
            case TokenType::CLASS:
            case TokenType::FUNC:
            case TokenType::VAR:
            case TokenType::FOR:
            case TokenType::IF:
            case TokenType::WHILE:
            case TokenType::RETURN:
                return;
            default:
                break;
        }
        
        advance();
    }
}

StmtPtr Parser::declaration() {
    if (match({TokenType::CLASS})) {
        return class_declaration();
    }
    if (match({TokenType::FUNC})) {
        return function_declaration();
    }
    if (check(TokenType::NUMBER) || check(TokenType::STRING) || 
        check(TokenType::BOOL) || check(TokenType::LIST) || 
        check(TokenType::MAP) || check(TokenType::IDENTIFIER)) {
        // Type name followed by identifier
        Token lookahead = tokens_[current_];
        if (current_ + 1 < tokens_.size() && tokens_[current_ + 1].type == TokenType::IDENTIFIER) {
            return var_declaration();
        }
    }
    
    return statement();
}

StmtPtr Parser::class_declaration() {
    Token name = consume(TokenType::IDENTIFIER, "Expected class name.");
    consume(TokenType::LEFT_BRACE, "Expected '{' before class body.");
    
    std::vector<StmtPtr> members;
    
    while (!check(TokenType::RIGHT_BRACE) && !is_at_end()) {
        // Parse field or method
        if (check(TokenType::FUNC)) {
            advance();
            members.push_back(function_declaration());
        } else {
            members.push_back(var_declaration());
            consume(TokenType::SEMICOLON, "Expected ';' after field declaration.");
        }
    }
    
    consume(TokenType::RIGHT_BRACE, "Expected '}' after class body.");
    
    return std::make_unique<ClassDecl>(name.lexeme, std::move(members));
}

StmtPtr Parser::function_declaration() {
    Token name = consume(TokenType::IDENTIFIER, "Expected function name.");
    consume(TokenType::LEFT_PAREN, "Expected '(' after function name.");
    
    std::vector<std::pair<std::string, std::string>> parameters;
    
    if (!check(TokenType::RIGHT_PAREN)) {
        do {
            Token type_token = advance();
            std::string type_name = type_token.lexeme;
            Token param_name = consume(TokenType::IDENTIFIER, "Expected parameter name.");
            parameters.emplace_back(type_name, param_name.lexeme);
        } while (match({TokenType::COMMA}));
    }
    
    consume(TokenType::RIGHT_PAREN, "Expected ')' after parameters.");
    
    std::string return_type = "nil";
    if (match({TokenType::ARROW})) {
        Token type_token = advance();
        return_type = type_token.lexeme;
    }
    
    consume(TokenType::LEFT_BRACE, "Expected '{' before function body.");
    StmtPtr body = block_statement();
    
    auto body_block = std::shared_ptr<BlockStmt>(static_cast<BlockStmt*>(body.release()));
    
    return std::make_unique<FunctionDecl>(name.lexeme, std::move(parameters), 
                                          return_type, std::move(body_block));
}

StmtPtr Parser::var_declaration() {
    Token type_token = advance();
    std::string type_name = type_token.lexeme;
    Token name = consume(TokenType::IDENTIFIER, "Expected variable name.");
    
    ExprPtr initializer = nullptr;
    if (match({TokenType::EQUAL})) {
        initializer = expression();
    }
    
    return std::make_unique<VarDecl>(type_name, name.lexeme, std::move(initializer));
}

StmtPtr Parser::statement() {
    if (match({TokenType::IF})) return if_statement();
    if (match({TokenType::WHILE})) return while_statement();
    if (match({TokenType::FOR})) return for_statement();
    if (match({TokenType::RETURN})) return return_statement();
    if (match({TokenType::LEFT_BRACE})) return block_statement();
    
    return expression_statement();
}

StmtPtr Parser::if_statement() {
    consume(TokenType::LEFT_PAREN, "Expected '(' after 'if'.");
    ExprPtr condition = expression();
    consume(TokenType::RIGHT_PAREN, "Expected ')' after if condition.");
    
    StmtPtr then_branch = statement();
    StmtPtr else_branch = nullptr;
    
    if (match({TokenType::ELSE})) {
        else_branch = statement();
    }
    
    return std::make_unique<IfStmt>(std::move(condition), std::move(then_branch), 
                                    std::move(else_branch));
}

StmtPtr Parser::while_statement() {
    consume(TokenType::LEFT_PAREN, "Expected '(' after 'while'.");
    ExprPtr condition = expression();
    consume(TokenType::RIGHT_PAREN, "Expected ')' after while condition.");
    
    StmtPtr body = statement();
    
    return std::make_unique<WhileStmt>(std::move(condition), std::move(body));
}

StmtPtr Parser::for_statement() {
    consume(TokenType::LEFT_PAREN, "Expected '(' after 'for'.");
    
    StmtPtr initializer = nullptr;
    if (!match({TokenType::SEMICOLON})) {
        initializer = declaration();
        consume(TokenType::SEMICOLON, "Expected ';' after for initializer.");
    }
    
    ExprPtr condition = nullptr;
    if (!check(TokenType::SEMICOLON)) {
        condition = expression();
    }
    consume(TokenType::SEMICOLON, "Expected ';' after for condition.");
    
    ExprPtr increment = nullptr;
    if (!check(TokenType::RIGHT_PAREN)) {
        increment = expression();
    }
    consume(TokenType::RIGHT_PAREN, "Expected ')' after for clauses.");
    
    StmtPtr body = statement();
    
    return std::make_unique<ForStmt>(std::move(initializer), std::move(condition),
                                     std::move(increment), std::move(body));
}

StmtPtr Parser::return_statement() {
    ExprPtr value = nullptr;
    
    if (!check(TokenType::SEMICOLON)) {
        value = expression();
    }
    
    consume(TokenType::SEMICOLON, "Expected ';' after return value.");
    return std::make_unique<ReturnStmt>(std::move(value));
}

StmtPtr Parser::block_statement() {
    std::vector<StmtPtr> statements;
    
    while (!check(TokenType::RIGHT_BRACE) && !is_at_end()) {
        auto stmt = declaration();
        // If declaration returned a VarDecl (not a class/function/statement), consume semicolon
        if (dynamic_cast<VarDecl*>(stmt.get())) {
            consume(TokenType::SEMICOLON, "Expected ';' after variable declaration.");
        }
        statements.push_back(std::move(stmt));
    }
    
    consume(TokenType::RIGHT_BRACE, "Expected '}' after block.");
    return std::make_unique<BlockStmt>(std::move(statements));
}

StmtPtr Parser::expression_statement() {
    ExprPtr expr = expression();
    consume(TokenType::SEMICOLON, "Expected ';' after expression.");
    return std::make_unique<ExprStmt>(std::move(expr));
}

ExprPtr Parser::expression() {
    return assignment();
}

ExprPtr Parser::assignment() {
    ExprPtr expr = logical_or();
    
    if (match({TokenType::EQUAL})) {
        Token equals = previous();
        ExprPtr value = assignment();
        
        if (auto* var_expr = dynamic_cast<VariableExpr*>(expr.get())) {
            return std::make_unique<AssignExpr>(var_expr->name, std::move(value));
        } else if (auto* get_expr = dynamic_cast<GetExpr*>(expr.get())) {
            return std::make_unique<SetExpr>(std::move(get_expr->object), 
                                             get_expr->name, std::move(value));
        } else if (auto* index_expr = dynamic_cast<IndexExpr*>(expr.get())) {
            return std::make_unique<IndexSetExpr>(std::move(index_expr->object),
                                                  std::move(index_expr->index),
                                                  std::move(value));
        }
        
        throw error(equals, "Invalid assignment target.");
    }
    
    return expr;
}

ExprPtr Parser::logical_or() {
    ExprPtr expr = logical_and();
    
    while (match({TokenType::OR})) {
        Token op = previous();
        ExprPtr right = logical_and();
        expr = std::make_unique<BinaryExpr>(std::move(expr), op.lexeme, std::move(right));
    }
    
    return expr;
}

ExprPtr Parser::logical_and() {
    ExprPtr expr = equality();
    
    while (match({TokenType::AND})) {
        Token op = previous();
        ExprPtr right = equality();
        expr = std::make_unique<BinaryExpr>(std::move(expr), op.lexeme, std::move(right));
    }
    
    return expr;
}

ExprPtr Parser::equality() {
    ExprPtr expr = comparison();
    
    while (match({TokenType::EQUAL_EQUAL, TokenType::BANG_EQUAL})) {
        Token op = previous();
        ExprPtr right = comparison();
        expr = std::make_unique<BinaryExpr>(std::move(expr), op.lexeme, std::move(right));
    }
    
    return expr;
}

ExprPtr Parser::comparison() {
    ExprPtr expr = term();
    
    while (match({TokenType::GREATER, TokenType::GREATER_EQUAL, 
                  TokenType::LESS, TokenType::LESS_EQUAL})) {
        Token op = previous();
        ExprPtr right = term();
        expr = std::make_unique<BinaryExpr>(std::move(expr), op.lexeme, std::move(right));
    }
    
    return expr;
}

ExprPtr Parser::term() {
    ExprPtr expr = factor();
    
    while (match({TokenType::MINUS, TokenType::PLUS})) {
        Token op = previous();
        ExprPtr right = factor();
        expr = std::make_unique<BinaryExpr>(std::move(expr), op.lexeme, std::move(right));
    }
    
    return expr;
}

ExprPtr Parser::factor() {
    ExprPtr expr = unary();
    
    while (match({TokenType::SLASH, TokenType::STAR, TokenType::PERCENT})) {
        Token op = previous();
        ExprPtr right = unary();
        expr = std::make_unique<BinaryExpr>(std::move(expr), op.lexeme, std::move(right));
    }
    
    return expr;
}

ExprPtr Parser::unary() {
    if (match({TokenType::BANG, TokenType::MINUS})) {
        Token op = previous();
        ExprPtr right = unary();
        return std::make_unique<UnaryExpr>(op.lexeme, std::move(right));
    }
    
    return call();
}

ExprPtr Parser::call() {
    ExprPtr expr = primary();
    
    while (true) {
        if (match({TokenType::LEFT_PAREN})) {
            expr = finish_call(std::move(expr));
        } else if (match({TokenType::DOT})) {
            Token name = consume(TokenType::IDENTIFIER, "Expected property name after '.'.");
            expr = std::make_unique<GetExpr>(std::move(expr), name.lexeme);
        } else if (match({TokenType::LEFT_BRACKET})) {
            ExprPtr index = expression();
            consume(TokenType::RIGHT_BRACKET, "Expected ']' after index.");
            expr = std::make_unique<IndexExpr>(std::move(expr), std::move(index));
        } else {
            break;
        }
    }
    
    return expr;
}

ExprPtr Parser::primary() {
    if (match({TokenType::TRUE})) {
        return std::make_unique<LiteralExpr>(true);
    }
    if (match({TokenType::FALSE})) {
        return std::make_unique<LiteralExpr>(false);
    }
    if (match({TokenType::NIL})) {
        return std::make_unique<LiteralExpr>();
    }
    
    if (match({TokenType::NUMBER_LITERAL})) {
        return std::make_unique<LiteralExpr>(std::get<double>(previous().literal));
    }
    
    if (match({TokenType::STRING_LITERAL})) {
        return std::make_unique<LiteralExpr>(std::get<std::string>(previous().literal));
    }
    
    if (match({TokenType::THIS})) {
        return std::make_unique<VariableExpr>("this");
    }
    
    if (match({TokenType::IDENTIFIER})) {
        return std::make_unique<VariableExpr>(previous().lexeme);
    }
    
    if (match({TokenType::LEFT_PAREN})) {
        ExprPtr expr = expression();
        consume(TokenType::RIGHT_PAREN, "Expected ')' after expression.");
        return expr;
    }
    
    if (match({TokenType::LEFT_BRACKET})) {
        std::vector<ExprPtr> elements;
        
        if (!check(TokenType::RIGHT_BRACKET)) {
            do {
                elements.push_back(expression());
            } while (match({TokenType::COMMA}));
        }
        
        consume(TokenType::RIGHT_BRACKET, "Expected ']' after list elements.");
        return std::make_unique<ListExpr>(std::move(elements));
    }
    
    if (match({TokenType::LEFT_BRACE})) {
        std::vector<std::pair<ExprPtr, ExprPtr>> pairs;
        
        if (!check(TokenType::RIGHT_BRACE)) {
            do {
                ExprPtr key = expression();
                consume(TokenType::COLON, "Expected ':' after map key.");
                ExprPtr value = expression();
                pairs.emplace_back(std::move(key), std::move(value));
            } while (match({TokenType::COMMA}));
        }
        
        consume(TokenType::RIGHT_BRACE, "Expected '}' after map pairs.");
        return std::make_unique<MapExpr>(std::move(pairs));
    }
    
    throw error(peek(), "Expected expression.");
}

ExprPtr Parser::finish_call(ExprPtr callee) {
    std::vector<ExprPtr> arguments;
    
    if (!check(TokenType::RIGHT_PAREN)) {
        do {
            arguments.push_back(expression());
        } while (match({TokenType::COMMA}));
    }
    
    consume(TokenType::RIGHT_PAREN, "Expected ')' after arguments.");
    
    return std::make_unique<CallExpr>(std::move(callee), std::move(arguments));
}

} // namespace camellya