package p.studio.compiler.pbs.semantics;

import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.Model;
import p.studio.compiler.pbs.semantics.PbsFlowSemanticSupport.TypeView;
import p.studio.compiler.source.Span;
import p.studio.compiler.source.diagnostics.DiagnosticSink;

import java.util.*;

final class PbsConstSemanticsValidator {
    private final PbsFlowTypeOps typeOps = new PbsFlowTypeOps();

    private record ConstAnalysis(
            TypeView type,
            boolean constant,
            Span invalidSpan,
            Set<String> dependencies) {
    }

    void validate(final PbsAst.File ast, final DiagnosticSink diagnostics) {
        final var model = Model.from(ast);
        final var constByName = collectConstDeclarations(ast);
        final var dependencyGraph = new LinkedHashMap<String, Set<String>>();

        for (final var entry : constByName.entrySet()) {
            final var constDecl = entry.getValue();
            if (constDecl.explicitType() == null || constDecl.initializer() == null) {
                continue;
            }

            final var declaredType = typeOps.typeFromTypeRef(constDecl.explicitType(), model, null);
            final var analysis = analyzeConstExpression(constDecl.initializer(), constByName, model, diagnostics);
            dependencyGraph.put(entry.getKey(), analysis.dependencies());

            if (!analysis.constant()) {
                final var invalidSpan = analysis.invalidSpan() == null ? constDecl.initializer().span() : analysis.invalidSpan();
                p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_CONST_NON_CONSTANT_INITIALIZER.name(),
                        "Const initializer for '%s' is not a constant expression".formatted(constDecl.name()),
                        invalidSpan);
                continue;
            }

            if (!typeOps.compatible(analysis.type(), declaredType)) {
                p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_CONST_INITIALIZER_TYPE_MISMATCH.name(),
                        "Const initializer type is incompatible with declared type for '%s'".formatted(constDecl.name()),
                        constDecl.initializer().span());
            }
        }

        validateConstCycles(constByName, dependencyGraph, diagnostics);
    }

    private LinkedHashMap<String, PbsAst.ConstDecl> collectConstDeclarations(final PbsAst.File ast) {
        final var constByName = new LinkedHashMap<String, PbsAst.ConstDecl>();
        for (final var topDecl : ast.topDecls()) {
            if (topDecl instanceof PbsAst.ConstDecl constDecl) {
                constByName.putIfAbsent(constDecl.name(), constDecl);
            }
        }
        return constByName;
    }

    private ConstAnalysis analyzeConstExpression(
            final PbsAst.Expression expression,
            final Map<String, PbsAst.ConstDecl> constByName,
            final Model model,
            final DiagnosticSink diagnostics) {
        if (expression instanceof PbsAst.IntLiteralExpr || expression instanceof PbsAst.BoundedLiteralExpr) {
            return new ConstAnalysis(TypeView.intType(), true, null, Set.of());
        }
        if (expression instanceof PbsAst.FloatLiteralExpr) {
            return new ConstAnalysis(TypeView.floatType(), true, null, Set.of());
        }
        if (expression instanceof PbsAst.StringLiteralExpr) {
            return new ConstAnalysis(TypeView.str(), true, null, Set.of());
        }
        if (expression instanceof PbsAst.BoolLiteralExpr) {
            return new ConstAnalysis(TypeView.bool(), true, null, Set.of());
        }
        if (expression instanceof PbsAst.GroupExpr groupExpr) {
            return analyzeConstExpression(groupExpr.expression(), constByName, model, diagnostics);
        }
        if (expression instanceof PbsAst.IdentifierExpr identifierExpr) {
            final var depName = identifierExpr.name();
            final var depDecl = constByName.get(depName);
            if (depDecl == null || depDecl.explicitType() == null || depDecl.initializer() == null) {
                p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics,
                        PbsSemanticsErrors.E_SEM_CONST_UNRESOLVED_REFERENCE.name(),
                        "Const initializer references unresolved const '%s'".formatted(depName),
                        identifierExpr.span());
                return new ConstAnalysis(TypeView.unknown(), false, identifierExpr.span(), Set.of());
            }
            return new ConstAnalysis(
                    typeOps.typeFromTypeRef(depDecl.explicitType(), model, null),
                    true,
                    null,
                    Set.of(depName));
        }
        if (expression instanceof PbsAst.MemberExpr memberExpr) {
            if (memberExpr.receiver() instanceof PbsAst.IdentifierExpr enumIdentifier) {
                final var cases = model.enums.get(enumIdentifier.name());
                if (cases != null && cases.contains(memberExpr.memberName())) {
                    return new ConstAnalysis(TypeView.enumType(enumIdentifier.name()), true, null, Set.of());
                }
            }
            return new ConstAnalysis(TypeView.unknown(), false, memberExpr.span(), Set.of());
        }
        if (expression instanceof PbsAst.UnaryExpr unaryExpr) {
            final var operand = analyzeConstExpression(unaryExpr.expression(), constByName, model, diagnostics);
            if (!operand.constant()) {
                return operand;
            }
            if ("-".equals(unaryExpr.operator()) && isNumeric(operand.type())) {
                return operand;
            }
            if (("!".equals(unaryExpr.operator()) || "not".equals(unaryExpr.operator())) && typeOps.isBool(operand.type())) {
                return new ConstAnalysis(TypeView.bool(), true, null, operand.dependencies());
            }
            return new ConstAnalysis(TypeView.unknown(), false, unaryExpr.span(), operand.dependencies());
        }
        if (expression instanceof PbsAst.BinaryExpr binaryExpr) {
            final var left = analyzeConstExpression(binaryExpr.left(), constByName, model, diagnostics);
            final var right = analyzeConstExpression(binaryExpr.right(), constByName, model, diagnostics);
            final var deps = mergeDependencies(left.dependencies(), right.dependencies());
            if (!left.constant()) {
                return new ConstAnalysis(TypeView.unknown(), false, left.invalidSpan(), deps);
            }
            if (!right.constant()) {
                return new ConstAnalysis(TypeView.unknown(), false, right.invalidSpan(), deps);
            }

            final var operator = binaryExpr.operator();
            if (isArithmeticOperator(operator)) {
                if (!isNumeric(left.type()) || !isNumeric(right.type())) {
                    return new ConstAnalysis(TypeView.unknown(), false, binaryExpr.span(), deps);
                }
                return new ConstAnalysis(typeOps.inferBinaryResult(operator, left.type(), right.type()), true, null, deps);
            }
            if (isComparisonOperator(operator)) {
                if (!typeOps.compatible(left.type(), right.type())) {
                    return new ConstAnalysis(TypeView.unknown(), false, binaryExpr.span(), deps);
                }
                return new ConstAnalysis(TypeView.bool(), true, null, deps);
            }
            if (isBooleanOperator(operator)) {
                if (!typeOps.isBool(left.type()) || !typeOps.isBool(right.type())) {
                    return new ConstAnalysis(TypeView.unknown(), false, binaryExpr.span(), deps);
                }
                return new ConstAnalysis(TypeView.bool(), true, null, deps);
            }
            return new ConstAnalysis(TypeView.unknown(), false, binaryExpr.span(), deps);
        }

        return new ConstAnalysis(TypeView.unknown(), false, expression.span(), Set.of());
    }

    private void validateConstCycles(
            final Map<String, PbsAst.ConstDecl> constByName,
            final Map<String, Set<String>> dependencyGraph,
            final DiagnosticSink diagnostics) {
        final var indegree = new HashMap<String, Integer>();
        final var reverse = new HashMap<String, Set<String>>();
        for (final var node : dependencyGraph.keySet()) {
            indegree.put(node, 0);
            reverse.put(node, new HashSet<>());
        }

        for (final var entry : dependencyGraph.entrySet()) {
            for (final var dep : entry.getValue()) {
                if (!dependencyGraph.containsKey(dep)) {
                    continue;
                }
                indegree.put(entry.getKey(), indegree.get(entry.getKey()) + 1);
                reverse.get(dep).add(entry.getKey());
            }
        }

        final var ready = new PriorityQueue<String>();
        for (final var entry : indegree.entrySet()) {
            if (entry.getValue() == 0) {
                ready.add(entry.getKey());
            }
        }

        var visited = 0;
        while (!ready.isEmpty()) {
            final var current = ready.remove();
            visited++;
            for (final var dependent : reverse.getOrDefault(current, Set.of())) {
                final var next = indegree.get(dependent) - 1;
                indegree.put(dependent, next);
                if (next == 0) {
                    ready.add(dependent);
                }
            }
        }

        if (visited == dependencyGraph.size()) {
            return;
        }

        final var cyclic = new ArrayList<String>();
        for (final var entry : indegree.entrySet()) {
            if (entry.getValue() > 0) {
                cyclic.add(entry.getKey());
            }
        }
        cyclic.sort(String::compareTo);

        for (final var constName : cyclic) {
            final var constDecl = constByName.get(constName);
            if (constDecl == null) {
                continue;
            }
            p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_CONST_CYCLIC_DEPENDENCY.name(),
                    "Cyclic const dependency detected involving '%s'".formatted(constName),
                    constDecl.span());
        }
    }

    private Set<String> mergeDependencies(final Set<String> left, final Set<String> right) {
        final var merged = new HashSet<String>();
        merged.addAll(left);
        merged.addAll(right);
        return merged;
    }

    private boolean isNumeric(final TypeView type) {
        return typeOps.isInt(type) || typeOps.isFloat(type);
    }

    private boolean isArithmeticOperator(final String operator) {
        return "+".equals(operator)
                || "-".equals(operator)
                || "*".equals(operator)
                || "/".equals(operator)
                || "%".equals(operator);
    }

    private boolean isComparisonOperator(final String operator) {
        return "==".equals(operator)
                || "!=".equals(operator)
                || "<".equals(operator)
                || "<=".equals(operator)
                || ">".equals(operator)
                || ">=".equals(operator);
    }

    private boolean isBooleanOperator(final String operator) {
        return "and".equals(operator)
                || "or".equals(operator)
                || "&&".equals(operator)
                || "||".equals(operator);
    }
}
