package io.gitlab.jfronny.muscript.compiler;

import io.gitlab.jfronny.muscript.ast.*;
import io.gitlab.jfronny.muscript.ast.bool.*;
import io.gitlab.jfronny.muscript.ast.compare.Equal;
import io.gitlab.jfronny.muscript.ast.compare.Greater;
import io.gitlab.jfronny.muscript.ast.conditional.UnresolvedConditional;
import io.gitlab.jfronny.muscript.ast.dynamic.*;
import io.gitlab.jfronny.muscript.ast.dynamic.assign.DynamicAssign;
import io.gitlab.jfronny.muscript.ast.literal.DynamicLiteral;
import io.gitlab.jfronny.muscript.ast.math.*;
import io.gitlab.jfronny.muscript.ast.string.Concatenate;
import io.gitlab.jfronny.muscript.data.Script;
import io.gitlab.jfronny.muscript.data.dynamic.additional.DFinal;
import io.gitlab.jfronny.muscript.error.*;
import org.jetbrains.annotations.Nullable;

import java.util.*;

public class Parser {
    private final Lexer lexer;

    private final TokenData previous = new TokenData();
    private final TokenData current = new TokenData();

    public static Expr<?> parse(String source) {
        return new Parser(new Lexer(source)).parse().optimize();
    }

    public static Script parseScript(String source) {
        return new Parser(new Lexer(source)).parseScript().optimize();
    }

    public Parser(Lexer lexer) {
        this.lexer = lexer;
    }

    /**
     * Generate a single expression.
     * Must exhaust the source!
     *
     * @return the resulting expression
     */
    public Expr<?> parse() {
        advance();
        Expr<?> expr = expression();
        if (!isAtEnd())
            throw new ParseException(LocationalError.create(lexer.source, lexer.start, lexer.current - 1, "Unexpected element after end of expression"));
        return expr;
    }

    /**
     * Generate a script instance.
     * Multiple instructions will be executed in sequence and the result of the last one will be returned.
     *
     * @return the resulting expression
     */
    public Script parseScript() {
        advance();
        List<Expr<?>> expressions = new LinkedList<>();
        while (!isAtEnd()) {
            expressions.add(expression());
            match(Token.Semicolon); // Consume semicolon if present
        }
        if (expressions.isEmpty()) throw new ParseException(LocationalError.create(lexer.source, lexer.start, lexer.current - 1, "Missing any elements in closure"));
        return new Script(expressions);
    }

    // Expressions
    private Expr<?> expression() {
        try {
            return conditional();
        } catch (RuntimeException e) {
            if (e instanceof ParseException) throw e;
            else if (e instanceof LocationalException le) {
                throw new ParseException(le.asPrintable(lexer.source), le.getCause());
            } else throw error(e.getMessage());
        }
    }

    private Expr<?> conditional() {
        Expr<?> expr = and();

        if (match(Token.QuestionMark)) {
            int start = previous.start;
            int end = previous.current - 1;
            Expr<?> trueExpr = expression();
            consume(Token.Colon, "Expected ':' after first part of condition");
            Expr<?> falseExpr = expression();
            expr = new UnresolvedConditional(start, end, asBool(expr), trueExpr, falseExpr);
        }

        return expr;
    }

    private Expr<?> and() {
        Expr<?> expr = or();

        while (match(Token.And)) {
            int start = previous.start;
            int end = previous.current - 1;
            Expr<?> right = or();
            expr = new And(start, end, asBool(expr), asBool(right));
        }

        return expr;
    }

    private Expr<?> or() {
        Expr<?> expr = equality();

        while (match(Token.Or)) {
            int start = previous.start;
            int end = previous.current - 1;
            Expr<?> right = equality();
            expr = new Or(start, end, asBool(expr), asBool(right));
        }

        return expr;
    }

    private Expr<?> equality() {
        Expr<?> expr = concat();

        while (match(Token.EqualEqual, Token.BangEqual)) {
            Token op = previous.token;
            int start = previous.start;
            int end = previous.current - 1;
            Expr<?> right = concat();
            BoolExpr e = new Equal(start, end, expr, right);
            if (op == Token.BangEqual) e = new Not(start, end, e);
            expr = e;
        }

        return expr;
    }

    private Expr<?> concat() {
        Expr<?> expr = comparison();

        while (match(Token.Concat)) {
            int start = previous.start;
            int end = previous.current - 1;
            Expr<?> right = comparison();
            expr = new Concatenate(start, end, asString(expr), asString(right));
        }

        return expr;
    }

    private Expr<?> comparison() {
        Expr<?> expr = term();

        while (match(Token.Greater, Token.GreaterEqual, Token.Less, Token.LessEqual)) {
            Token op = previous.token;
            int start = previous.start;
            int end = previous.current - 1;
            NumberExpr right = asNumber(term());
            expr = switch (op) {
                case Greater -> new Greater(start, end, asNumber(expr), right);
                case GreaterEqual -> new Not(start, end, new Greater(start, end, right, asNumber(expr)));
                case Less -> new Greater(start, end, right, asNumber(expr));
                case LessEqual -> new Not(start, end, new Greater(start, end, asNumber(expr), right));
                default -> throw new IllegalStateException();
            };
        }

        return expr;
    }

    private Expr<?> term() {
        Expr<?> expr = factor();

        while (match(Token.Plus, Token.Minus)) {
            Token op = previous.token;
            int start = previous.start;
            int end = previous.current - 1;
            NumberExpr right = asNumber(factor());
            expr = switch (op) {
                case Plus -> new Plus(start, end, asNumber(expr), right);
                case Minus -> new Minus(start, end, asNumber(expr), right);
                default -> throw new IllegalStateException();
            };
        }

        return expr;
    }

    private Expr<?> factor() {
        Expr<?> expr = exp();

        while (match(Token.Star, Token.Slash, Token.Percentage)) {
            Token op = previous.token;
            int start = previous.start;
            int end = previous.current - 1;
            NumberExpr right = asNumber(exp());
            expr = switch (op) {
                case Star -> new Multiply(start, end, asNumber(expr), right);
                case Slash -> new Divide(start, end, asNumber(expr), right);
                case Percentage -> new Modulo(start, end, asNumber(expr), right);
                default -> throw new IllegalStateException();
            };
        }

        return expr;
    }

    private Expr<?> exp() {
        Expr<?> expr = unary();

        while (match(Token.UpArrow)) {
            int start = previous.start;
            int end = previous.current - 1;
            NumberExpr right = asNumber(unary());
            expr = new Power(start, end, asNumber(expr), right);
        }

        return expr;
    }

    private Expr<?> unary() {
        if (match(Token.Bang, Token.Minus)) {
            Token op = previous.token;
            int start = previous.start;
            int end = previous.current - 1;
            Expr<?> right = unary();
            return switch (op) {
                case Bang -> new Not(start, end, asBool(right));
                case Minus -> new Invert(start, end, asNumber(right));
                default -> throw new IllegalStateException();
            };
        }

        return call();
    }

    private Expr<?> call() {
        Expr<?> expr = primary();

        while (match(Token.LeftParen, Token.Dot, Token.LeftBracket, Token.DoubleColon)) {
            int start = previous.start;
            int end = previous.current - 1;
            expr = switch (previous.token) {
                case LeftParen -> finishCall(start, end, expr);
                case Dot -> {
                    TokenData name = consume(Token.Identifier, "Expected field name after '.'");
                    yield new Get(start, end, asDynamic(expr), Expr.literal(name.start, name.current - 1, name.lexeme));
                }
                case DoubleColon -> {
                    DynamicExpr callable;
                    if (match(Token.Identifier)) {
                        callable = new Variable(previous.start, previous.current - 1, previous.lexeme);
                    } else if (match(Token.LeftParen)) {
                        callable = expression().asDynamicExpr();
                        consume(Token.RightParen, "Expected ')' after expression");
                    } else throw error("Bind operator requires right side to be a literal identifier or to be wrapped in parentheses.");
                    yield new Bind(start, end, callable, expr.asDynamicExpr());
                }
                case LeftBracket -> {
                    expr = new Get(start, end, asDynamic(expr), expression());
                    consume(Token.RightBracket, "Expected closing bracket");
                    yield expr;
                }
                default -> throw new IllegalStateException();
            };
        }

        return expr;
    }

    private Expr<?> finishCall(int start, int end, Expr<?> callee) {
        List<Call.Arg> args = new ArrayList<>(2);

        if (!check(Token.RightParen)) {
            do {
                args.add(new Call.Arg(asDynamic(expression()), match(Token.Ellipsis)));
            } while (match(Token.Comma));
        }

        consume(Token.RightParen, "Expected ')' after function arguments");
        return new Call(start, end, asDynamic(callee), args);
    }

    private Expr<?> primary() {
        if (match(Token.Null)) return Expr.literalNull(previous.start, previous.current - 1);
        if (match(Token.String)) return Expr.literal(previous.start, previous.current - 1, previous.lexeme);
        if (match(Token.True, Token.False)) return Expr.literal(previous.start, previous.current - 1, previous.lexeme.equals("true"));
        if (match(Token.Number)) return Expr.literal(previous.start, previous.current - 1, Double.parseDouble(previous.lexeme));
        if (match(Token.Identifier)) {
            int start = previous.start;
            int end = previous.current - 1;
            String name = previous.lexeme;
            if (match(Token.Assign)) return new DynamicAssign(start, end, name, expression().asDynamicExpr());
            else return new Variable(start, end, name);
        }

        if (match(Token.LeftParen)) {
            Expr<?> expr = expression();
            consume(Token.RightParen, "Expected ')' after expression");
            return expr;
        }

        if (match(Token.LeftBrace)) {
            int start = previous.start;
            if (match(Token.Arrow)) return finishClosure(start, null, false);
            if (match(Token.RightBrace)) return new DynamicLiteral<>(start, previous.start, DFinal.of(Map.of()));
            consume(Token.Identifier, "Expected arrow or identifier as first element in closure or object");
            String first = previous.lexeme;
            if (check(Token.Arrow)) return finishClosure(start, first, false);
            if (match(Token.Ellipsis)) return finishClosure(start, first, true);
            if (check(Token.Comma)) return finishClosure(start, first, false);
            if (match(Token.Assign)) {
                return finishObject(start, first, expression().asDynamicExpr());
            }
            throw error("Unexpected");
        }

        throw error("Expected expression.");
    }

    private Expr<?> finishClosure(int start, @Nullable String firstArg, boolean firstVariadic) {
        List<String> boundArgs = new LinkedList<>();
        boolean variadic = false;
        if (firstArg != null) {
            boundArgs.add(firstArg);
            if (firstVariadic) {
                consume(Token.Arrow, "Variadic argument MUST be the last argument");
                variadic = true;
            } else {
                while (!match(Token.Arrow)) {
                    consume(Token.Comma, "Closure parameters MUST be comma-seperated");
                    consume(Token.Identifier, "Closure arguments MUST be identifiers");
                    boundArgs.add(previous.lexeme);
                    if (match(Token.Ellipsis)) {
                        variadic = true;
                        consume(Token.Arrow, "Variadic argument MUST be the last argument");
                        break;
                    }
                }
            }
        }
        List<Expr<?>> expressions = new LinkedList<>();
        while (!match(Token.RightBrace)) {
            expressions.add(expression());
            match(Token.Semicolon); // Consume semicolon if present
        }
        int end = previous.start;
        return new Closure(start, end, boundArgs, expressions, variadic);
    }

    private Expr<?> finishObject(int start, @Nullable String firstArg, @Nullable DynamicExpr firstValue) {
        Map<String, DynamicExpr> content = new LinkedHashMap<>();
        content.put(firstArg, firstValue);
        while (match(Token.Comma)) {
            consume(Token.Identifier, "Object element MUST start with an identifier");
            String name = previous.lexeme;
            consume(Token.Assign, "Object element name and value MUST be seperated with '='");
            content.put(name, expression().asDynamicExpr());
        }
        consume(Token.RightBrace, "Expected end of object");
        return new ObjectLiteral(start, previous.start, content);
    }

    // Type conversion
    private BoolExpr asBool(Expr<?> expression) {
        try {
            return expression.asBoolExpr();
        } catch (TypeMismatchException e) {
            throw error(e.getMessage(), expression);
        }
    }

    private NumberExpr asNumber(Expr<?> expression) {
        try {
            return expression.asNumberExpr();
        } catch (TypeMismatchException e) {
            throw error(e.getMessage(), expression);
        }
    }

    private StringExpr asString(Expr<?> expression) {
        try {
            return expression.asStringExpr();
        } catch (TypeMismatchException e) {
            throw error(e.getMessage(), expression);
        }
    }

    private DynamicExpr asDynamic(Expr<?> expression) {
        try {
            return expression.asDynamicExpr();
        } catch (TypeMismatchException e) {
            throw error(e.getMessage(), expression);
        }
    }

    // Helpers
    private ParseException error(String message) {
        return new ParseException(LocationalError.create(lexer.source, current.current - 1, message));
    }

    private ParseException error(String message, Expr<?> expr) {
        return new ParseException(LocationalError.create(lexer.source, expr.chStart, expr.chEnd, message));
    }

    private TokenData consume(Token token, String message) {
        if (check(token)) return advance();
        throw error(message + " but got " + current.token);
    }

    private boolean match(Token... tokens) {
        for (Token token : tokens) {
            if (check(token)) {
                advance();
                return true;
            }
        }

        return false;
    }

    private boolean check(Token token) {
        if (isAtEnd()) return false;
        return current.token == token;
    }

    private TokenData advance() {
        previous.set(current);

        lexer.next();
        current.set(lexer.token, lexer.lexeme, lexer.start, lexer.current, lexer.ch);

        if (current.token == Token.Error) {
            throw error(current.lexeme);
        }

        return previous;
    }

    private boolean isAtEnd() {
        return current.token == Token.EOF;
    }

    // Token data
    private static class TokenData {
        public Token token;
        public String lexeme;
        public int start, current;
        public char ch;

        public void set(Token token, String lexeme, int start, int current, char ch) {
            this.token = token;
            this.lexeme = lexeme;
            this.start = start;
            this.current = current;
            this.ch = ch;
        }

        public void set(TokenData data) {
            set(data.token, data.lexeme, data.start, data.current, data.ch);
        }

        @Override
        public String toString() {
            return String.format("%s '%s'", token, lexeme);
        }
    }

    // Parse Exception
    public static class ParseException extends RuntimeException {
        public final LocationalError error;

        public ParseException(LocationalError error) {
            super(error.toString());
            this.error = error;
        }

        public ParseException(LocationalError error, Throwable cause) {
            super(error.toString(), cause);
            this.error = error;
        }
    }
}