package p.studio.compiler.pbs.semantics;

import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.source.Span;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.diagnostics.RelatedSpan;
import p.studio.compiler.source.diagnostics.Severity;
import p.studio.compiler.source.identifiers.NameId;
import p.studio.compiler.source.tables.NameTable;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.HashMap;
import java.util.List;

final class PbsDeclarationRuleValidator {
    private final NameTable nameTable;
    private final DiagnosticSink diagnostics;

    PbsDeclarationRuleValidator(
            final NameTable nameTable,
            final DiagnosticSink diagnostics) {
        this.nameTable = nameTable;
        this.diagnostics = diagnostics;
    }

    void validateCallableHeader(
            final String callableName,
            final ReadOnlyList<PbsAst.Parameter> parameters,
            final PbsAst.ReturnKind returnKind,
            final PbsAst.TypeRef returnType,
            final PbsAst.TypeRef resultErrorType,
            final Span span,
            final boolean allowSelf) {
        validateParameters(parameters, "callable '%s'".formatted(callableName));
        validateReturnSurface(
                returnKind,
                returnType,
                resultErrorType,
                "callable '%s'".formatted(callableName),
                span,
                allowSelf);
        PbsTypeSurfaceSemanticsValidator.validateTypeSurfaceRecursive(
                parameters.stream().map(PbsAst.Parameter::typeRef).toList(),
                "callable '%s'".formatted(callableName),
                allowSelf,
                diagnostics);
    }

    void validateParameters(
            final ReadOnlyList<PbsAst.Parameter> parameters,
            final String ownerDescription) {
        final var seen = new HashMap<NameId, Span>();
        for (final var parameter : parameters) {
            final var parameterNameId = nameTable.register(parameter.name());
            final var first = seen.putIfAbsent(parameterNameId, parameter.span());
            if (first != null) {
                p.studio.compiler.source.diagnostics.Diagnostics.report(diagnostics, 
                        Severity.Error,
                        PbsSemanticsErrors.E_SEM_DUPLICATE_PARAMETER_NAME.name(),
                        "Duplicate parameter name '%s' in %s".formatted(parameter.name(), ownerDescription),
                        parameter.span(),
                        List.of(new RelatedSpan("First parameter is here", first)));
            }
        }
    }

    void validateTypeSurfaces(
            final List<PbsAst.TypeRef> typeRoots,
            final String ownerDescription,
            final boolean allowSelf) {
        PbsTypeSurfaceSemanticsValidator.validateTypeSurfaceRecursive(typeRoots, ownerDescription, allowSelf, diagnostics);
    }

    void validateErrorDeclaration(final PbsAst.ErrorDecl errorDecl) {
        final var seenCases = new HashMap<NameId, Span>();
        for (final var caseName : errorDecl.cases()) {
            final var caseNameId = nameTable.register(caseName);
            final var first = seenCases.putIfAbsent(caseNameId, errorDecl.span());
            if (first != null) {
                p.studio.compiler.source.diagnostics.Diagnostics.report(diagnostics, 
                        Severity.Error,
                        PbsSemanticsErrors.E_SEM_DUPLICATE_ERROR_CASE_LABEL.name(),
                        "Duplicate error case label '%s' in '%s'".formatted(caseName, errorDecl.name()),
                        errorDecl.span(),
                        List.of(new RelatedSpan("First case is here", first)));
            }
        }
    }

    void validateEnumDeclaration(final PbsAst.EnumDecl enumDecl) {
        final var seenLabels = new HashMap<NameId, Span>();
        final var seenIds = new HashMap<Long, Span>();
        for (final var enumCase : enumDecl.cases()) {
            final var labelId = nameTable.register(enumCase.name());
            final var firstLabel = seenLabels.putIfAbsent(labelId, enumCase.span());
            if (firstLabel != null) {
                p.studio.compiler.source.diagnostics.Diagnostics.report(diagnostics, 
                        Severity.Error,
                        PbsSemanticsErrors.E_SEM_DUPLICATE_ENUM_CASE_LABEL.name(),
                        "Duplicate enum case label '%s' in '%s'".formatted(enumCase.name(), enumDecl.name()),
                        enumCase.span(),
                        List.of(new RelatedSpan("First case is here", firstLabel)));
            }

            if (enumCase.explicitValue() != null) {
                final var firstId = seenIds.putIfAbsent(enumCase.explicitValue(), enumCase.span());
                if (firstId != null) {
                    p.studio.compiler.source.diagnostics.Diagnostics.report(diagnostics, 
                            Severity.Error,
                            PbsSemanticsErrors.E_SEM_DUPLICATE_ENUM_CASE_ID.name(),
                            "Duplicate enum case id '%s' in '%s'".formatted(enumCase.explicitValue(), enumDecl.name()),
                            enumCase.span(),
                            List.of(new RelatedSpan("First id is here", firstId)));
                }
            }
        }
    }

    void validateCallbackDeclaration(final PbsAst.CallbackDecl callbackDecl) {
        validateParameters(callbackDecl.parameters(), "callback '%s'".formatted(callbackDecl.name()));
        validateReturnSurface(
                callbackDecl.returnKind(),
                callbackDecl.returnType(),
                callbackDecl.resultErrorType(),
                "callback '%s'".formatted(callbackDecl.name()),
                callbackDecl.span(),
                false);
        PbsTypeSurfaceSemanticsValidator.validateTypeSurfaceRecursive(
                callbackDecl.parameters().stream().map(PbsAst.Parameter::typeRef).toList(),
                "callback '%s'".formatted(callbackDecl.name()),
                false,
                diagnostics);
    }

    void validateConstDeclaration(final PbsAst.ConstDecl constDecl) {
        validateConstDeclaration(constDecl, true);
    }

    void validateConstDeclaration(
            final PbsAst.ConstDecl constDecl,
            final boolean requireInitializer) {
        if (constDecl.explicitType() == null) {
            p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                    PbsSemanticsErrors.E_SEM_MISSING_CONST_TYPE_ANNOTATION.name(),
                    "Const declaration '%s' must include an explicit type annotation".formatted(constDecl.name()),
                    constDecl.span());
        } else {
            PbsTypeSurfaceSemanticsValidator.validateTypeSurfaceRecursive(
                    List.of(constDecl.explicitType()),
                    "const '%s'".formatted(constDecl.name()),
                    false,
                    diagnostics);
        }

        if (requireInitializer && constDecl.initializer() == null) {
            p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                    PbsSemanticsErrors.E_SEM_MISSING_CONST_INITIALIZER.name(),
                    "Non-builtin const declaration '%s' must include an initializer".formatted(constDecl.name()),
                    constDecl.span());
        }
    }

    void validateGlobalDeclaration(final PbsAst.GlobalDecl globalDecl) {
        if (globalDecl.explicitType() == null) {
            p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_MISSING_GLOBAL_TYPE_ANNOTATION.name(),
                    "Global declaration '%s' must include an explicit type annotation".formatted(globalDecl.name()),
                    globalDecl.span());
        } else {
            PbsTypeSurfaceSemanticsValidator.validateTypeSurfaceRecursive(
                    List.of(globalDecl.explicitType()),
                    "global '%s'".formatted(globalDecl.name()),
                    false,
                    diagnostics);
        }

        if (globalDecl.initializer() == null) {
            p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_MISSING_GLOBAL_INITIALIZER.name(),
                    "Global declaration '%s' must include an initializer".formatted(globalDecl.name()),
                    globalDecl.span());
            return;
        }

        final var invalidSpan = firstUnsupportedGlobalInitializerSpan(globalDecl.initializer());
        if (invalidSpan != null) {
            p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics,
                    PbsSemanticsErrors.E_SEM_GLOBAL_UNSUPPORTED_INITIALIZER.name(),
                    "Global initializer for '%s' uses an unsupported form".formatted(globalDecl.name()),
                    invalidSpan);
        }
    }

    private Span firstUnsupportedGlobalInitializerSpan(final PbsAst.Expression expression) {
        if (expression == null) {
            return null;
        }
        if (expression instanceof PbsAst.IntLiteralExpr
                || expression instanceof PbsAst.FloatLiteralExpr
                || expression instanceof PbsAst.BoundedLiteralExpr
                || expression instanceof PbsAst.StringLiteralExpr
                || expression instanceof PbsAst.BoolLiteralExpr
                || expression instanceof PbsAst.IdentifierExpr
                || expression instanceof PbsAst.ThisExpr
                || expression instanceof PbsAst.UnitExpr) {
            return null;
        }
        if (expression instanceof PbsAst.GroupExpr groupExpr) {
            return firstUnsupportedGlobalInitializerSpan(groupExpr.expression());
        }
        if (expression instanceof PbsAst.UnaryExpr unaryExpr) {
            return firstUnsupportedGlobalInitializerSpan(unaryExpr.expression());
        }
        if (expression instanceof PbsAst.BinaryExpr binaryExpr) {
            final var left = firstUnsupportedGlobalInitializerSpan(binaryExpr.left());
            if (left != null) {
                return left;
            }
            return firstUnsupportedGlobalInitializerSpan(binaryExpr.right());
        }
        if (expression instanceof PbsAst.MemberExpr memberExpr) {
            return firstUnsupportedGlobalInitializerSpan(memberExpr.receiver());
        }
        if (expression instanceof PbsAst.NewExpr newExpr) {
            for (final var argument : newExpr.arguments()) {
                final var invalid = firstUnsupportedGlobalInitializerSpan(argument);
                if (invalid != null) {
                    return invalid;
                }
            }
            return null;
        }
        return expression.span();
    }

    private void validateReturnSurface(
            final PbsAst.ReturnKind returnKind,
            final PbsAst.TypeRef returnType,
            final PbsAst.TypeRef resultErrorType,
            final String ownerDescription,
            final Span span,
            final boolean allowSelf) {
        if (returnKind == PbsAst.ReturnKind.RESULT && PbsTypeSurfaceSemanticsValidator.isOptionalType(returnType)) {
            p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, 
                    PbsSemanticsErrors.E_SEM_INVALID_MIXED_OPTIONAL_RESULT_RETURN.name(),
                    "Return surface for %s cannot combine 'result' with top-level 'optional' payload".formatted(ownerDescription),
                    span);
        }

        if (returnType != null && returnType.kind() == PbsAst.TypeRefKind.NAMED_TUPLE) {
            validateDuplicateOutputLabels(returnType, ownerDescription);
        }

        if (returnType != null) {
            PbsTypeSurfaceSemanticsValidator.validateTypeSurfaceRecursive(List.of(returnType), ownerDescription, allowSelf, diagnostics);
        }
        if (resultErrorType != null) {
            PbsTypeSurfaceSemanticsValidator.validateTypeSurfaceRecursive(List.of(resultErrorType), ownerDescription, allowSelf, diagnostics);
        }
    }

    private void validateDuplicateOutputLabels(
            final PbsAst.TypeRef namedTupleReturnType,
            final String ownerDescription) {
        final var seen = new HashMap<NameId, Span>();
        for (final var field : namedTupleReturnType.fields()) {
            final var labelId = nameTable.register(field.label());
            final var first = seen.putIfAbsent(labelId, field.span());
            if (first != null) {
                p.studio.compiler.source.diagnostics.Diagnostics.report(diagnostics, 
                        Severity.Error,
                        PbsSemanticsErrors.E_SEM_DUPLICATE_RETURN_LABEL.name(),
                        "Duplicate return label '%s' in %s".formatted(field.label(), ownerDescription),
                        field.span(),
                        List.of(new RelatedSpan("First label is here", first)));
            }
        }
    }
}
