package p.studio.compiler.pbs.semantics;

import lombok.experimental.UtilityClass;
import p.studio.compiler.messages.FESurfaceContext;
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.Diagnostics;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.*;

@UtilityClass
final class PbsFlowSemanticSupport {
    record ExprResult(
            TypeView type,
            List<CallableSymbol> callables,
            boolean methodTarget) {
        static ExprResult type(final TypeView type) {
            return new ExprResult(type, List.of(), false);
        }
    
        static ExprResult callables(final List<CallableSymbol> callables, final boolean methodTarget) {
            return new ExprResult(TypeView.unknown(), List.copyOf(callables), methodTarget);
        }
    }
    
    enum ExprUse {
        VALUE,
        CALL_TARGET,
        APPLY_TARGET
    }
    
    enum Kind {
        UNKNOWN,
        UNIT,
        INT,
        FLOAT,
        BOOL,
        STR,
        STRUCT,
        SERVICE,
        CONTRACT,
        CALLBACK,
        ENUM,
        ERROR,
        OPTIONAL,
        RESULT,
        TUPLE,
        TYPE_REF,
        ADDRESSABLE,
        ASSET_NAMESPACE
    }
    
    record TupleField(
            String label,
            TypeView type) {
    }
    
    record TypeView(
            Kind kind,
            String name,
            List<TupleField> tupleFields,
            TypeView inner,
            TypeView errorType,
            List<TypeView> callbackInputs,
            TypeView callbackOutput) {
        static TypeView unknown() {
            return new TypeView(Kind.UNKNOWN, null, List.of(), null, null, List.of(), null);
        }
    
        static TypeView unit() {
            return new TypeView(Kind.UNIT, "unit", List.of(), null, null, List.of(), null);
        }
    
        static TypeView intType() {
            return new TypeView(Kind.INT, "int", List.of(), null, null, List.of(), null);
        }
    
        static TypeView floatType() {
            return new TypeView(Kind.FLOAT, "float", List.of(), null, null, List.of(), null);
        }
    
        static TypeView bool() {
            return new TypeView(Kind.BOOL, "bool", List.of(), null, null, List.of(), null);
        }
    
        static TypeView str() {
            return new TypeView(Kind.STR, "str", List.of(), null, null, List.of(), null);
        }
    
        static TypeView struct(final String name) {
            return new TypeView(Kind.STRUCT, name, List.of(), null, null, List.of(), null);
        }
    
        static TypeView service(final String name) {
            return new TypeView(Kind.SERVICE, name, List.of(), null, null, List.of(), null);
        }
    
        static TypeView contract(final String name) {
            return new TypeView(Kind.CONTRACT, name, List.of(), null, null, List.of(), null);
        }
    
        static TypeView callback(
                final String name,
                final List<TypeView> inputTypes,
                final TypeView outputType) {
            return new TypeView(Kind.CALLBACK, name, List.of(), null, null, List.copyOf(inputTypes), outputType);
        }
    
        static TypeView enumType(final String name) {
            return new TypeView(Kind.ENUM, name, List.of(), null, null, List.of(), null);
        }
    
        static TypeView error(final String name) {
            return new TypeView(Kind.ERROR, name, List.of(), null, null, List.of(), null);
        }
    
        static TypeView optional(final TypeView inner) {
            return new TypeView(Kind.OPTIONAL, "optional", List.of(), inner, null, List.of(), null);
        }
    
        static TypeView result(final TypeView error, final TypeView payload) {
            return new TypeView(Kind.RESULT, "result", List.of(), payload, error, List.of(), null);
        }
    
        static TypeView tuple(final List<TupleField> fields) {
            return new TypeView(Kind.TUPLE, "tuple", List.copyOf(fields), null, null, List.of(), null);
        }
    
        static TypeView typeRef(final String name) {
            return new TypeView(Kind.TYPE_REF, name, List.of(), null, null, List.of(), null);
        }

        static TypeView addressable(final String address) {
            return new TypeView(Kind.ADDRESSABLE, address, List.of(), null, null, List.of(), null);
        }

        static TypeView assetNamespace(final String path) {
            return new TypeView(Kind.ASSET_NAMESPACE, path, List.of(), null, null, List.of(), null);
        }
    }
    
    record CallableSymbol(
            String name,
            List<TypeView> inputTypes,
            TypeView outputType,
            Span span) {
    }
    
    record CallbackSignature(
            List<TypeView> inputTypes,
            TypeView outputType) {
    }
    
    record StructInfo(
            Map<String, StructFieldInfo> fields,
            Map<String, List<CallableSymbol>> methods) {
    }

    record StructFieldInfo(
            TypeView type,
            boolean isPublic,
            boolean isMutable) {
    }
    
    record ServiceInfo(
            Map<String, List<CallableSymbol>> methods) {
    }
    
    record ContractInfo(
            Map<String, List<CallableSymbol>> methods) {
    }
    
    static final class Model {
        final Map<String, List<CallableSymbol>> topLevelCallables = new HashMap<>();
        final Map<String, StructInfo> structs = new HashMap<>();
        final Map<String, ServiceInfo> services = new HashMap<>();
        final Map<String, ContractInfo> contracts = new HashMap<>();
        final Map<String, CallbackSignature> callbacks = new HashMap<>();
        final Map<String, Set<String>> enums = new HashMap<>();
        final Map<String, Set<String>> errors = new HashMap<>();
        final Map<String, TypeView> constTypes = new HashMap<>();
        final Map<String, TypeView> globalTypes = new HashMap<>();
        final Map<String, TypeView> serviceSingletons = new HashMap<>();
        final Set<String> knownStructNames = new HashSet<>();
        final Set<String> knownServiceNames = new HashSet<>();
        final Set<String> knownContractNames = new HashSet<>();
        final Set<String> knownEnumNames = new HashSet<>();
        final Set<String> knownErrorNames = new HashSet<>();
        final Set<String> knownCallbackNames = new HashSet<>();
        final Map<String, Integer> assetIdByAddress = new HashMap<>();
        final Set<String> assetNamespacePrefixes = new HashSet<>();
    
        static Model from(final PbsAst.File ast) {
            return from(ast, ReadOnlyList.empty(), FESurfaceContext.empty(), DiagnosticSink.empty());
        }

        static Model from(
                final PbsAst.File ast,
                final ReadOnlyList<PbsAst.TopDecl> supplementalTopDecls,
                final FESurfaceContext feSurfaceContext,
                final DiagnosticSink diagnostics) {
            final var model = new Model();
            for (final var topDecl : ast.topDecls()) {
                model.registerKnownTopDecl(topDecl);
            }
            for (final var topDecl : supplementalTopDecls) {
                model.registerKnownTopDecl(topDecl);
            }
            for (final var topDecl : ast.topDecls()) {
                model.ingestTopDecl(topDecl);
            }
            for (final var topDecl : supplementalTopDecls) {
                model.ingestTopDecl(topDecl);
            }
            model.ingestAssetSurface(feSurfaceContext, diagnostics);
            return model;
        }

        private void registerKnownTopDecl(final PbsAst.TopDecl topDecl) {
            if (topDecl instanceof PbsAst.StructDecl structDecl) {
                knownStructNames.add(structDecl.name());
                return;
            }
            if (topDecl instanceof PbsAst.BuiltinTypeDecl builtinTypeDecl) {
                knownStructNames.add(builtinTypeDecl.name());
                return;
            }
            if (topDecl instanceof PbsAst.ServiceDecl serviceDecl) {
                knownServiceNames.add(serviceDecl.name());
                return;
            }
            if (topDecl instanceof PbsAst.HostDecl hostDecl) {
                knownServiceNames.add(hostDecl.name());
                return;
            }
            if (topDecl instanceof PbsAst.ContractDecl contractDecl) {
                knownContractNames.add(contractDecl.name());
                return;
            }
            if (topDecl instanceof PbsAst.EnumDecl enumDecl) {
                knownEnumNames.add(enumDecl.name());
                return;
            }
            if (topDecl instanceof PbsAst.ErrorDecl errorDecl) {
                knownErrorNames.add(errorDecl.name());
                return;
            }
            if (topDecl instanceof PbsAst.CallbackDecl callbackDecl) {
                knownCallbackNames.add(callbackDecl.name());
            }
        }

        private void ingestTopDecl(final PbsAst.TopDecl topDecl) {
            if (topDecl instanceof PbsAst.StructDecl structDecl) {
                final var fields = new HashMap<String, StructFieldInfo>();
                for (final var field : structDecl.fields()) {
                    fields.put(
                            field.name(),
                            new StructFieldInfo(
                                    typeFrom(field.typeRef()),
                                    field.isPublic(),
                                    field.isMutable()));
                }
                final var methods = new HashMap<String, List<CallableSymbol>>();
                for (final var method : structDecl.methods()) {
                    methods.computeIfAbsent(method.name(), ignored -> new ArrayList<>())
                            .add(callableFrom(
                                    method.name(),
                                    method.parameters(),
                                    method.returnKind(),
                                    method.returnType(),
                                    method.resultErrorType(),
                                    method.span()));
                }
                structs.put(structDecl.name(), new StructInfo(fields, methods));
                return;
            }
            if (topDecl instanceof PbsAst.BuiltinTypeDecl builtinTypeDecl) {
                final var fields = new HashMap<String, StructFieldInfo>();
                for (final var field : builtinTypeDecl.fields()) {
                    fields.put(
                            field.name(),
                            new StructFieldInfo(
                                    typeFrom(field.typeRef()),
                                    true,
                                    false));
                }
                final var methods = new HashMap<String, List<CallableSymbol>>();
                for (final var signature : builtinTypeDecl.signatures()) {
                    methods.computeIfAbsent(signature.name(), ignored -> new ArrayList<>())
                            .add(callableFrom(
                                    signature.name(),
                                    signature.parameters(),
                                    signature.returnKind(),
                                    signature.returnType(),
                                    signature.resultErrorType(),
                                    signature.span()));
                }
                structs.put(builtinTypeDecl.name(), new StructInfo(fields, methods));
                return;
            }
            if (topDecl instanceof PbsAst.ServiceDecl serviceDecl) {
                final var methods = new HashMap<String, List<CallableSymbol>>();
                for (final var method : serviceDecl.methods()) {
                    methods.computeIfAbsent(method.name(), ignored -> new ArrayList<>())
                            .add(callableFrom(
                                    method.name(),
                                    method.parameters(),
                                    method.returnKind(),
                                    method.returnType(),
                                    method.resultErrorType(),
                                    method.span()));
                }
                services.put(serviceDecl.name(), new ServiceInfo(methods));
                serviceSingletons.put(serviceDecl.name(), TypeView.service(serviceDecl.name()));
                return;
            }
            if (topDecl instanceof PbsAst.HostDecl hostDecl) {
                final var methods = new HashMap<String, List<CallableSymbol>>();
                for (final var signature : hostDecl.signatures()) {
                    methods.computeIfAbsent(signature.name(), ignored -> new ArrayList<>())
                            .add(callableFrom(
                                    signature.name(),
                                    signature.parameters(),
                                    signature.returnKind(),
                                    signature.returnType(),
                                    signature.resultErrorType(),
                                    signature.span()));
                }
                // Host owners are value singletons with callable members, same access shape as services.
                services.put(hostDecl.name(), new ServiceInfo(methods));
                serviceSingletons.put(hostDecl.name(), TypeView.service(hostDecl.name()));
                return;
            }
            if (topDecl instanceof PbsAst.ContractDecl contractDecl) {
                final var methods = new HashMap<String, List<CallableSymbol>>();
                for (final var signature : contractDecl.signatures()) {
                    methods.computeIfAbsent(signature.name(), ignored -> new ArrayList<>())
                            .add(callableFrom(
                                    signature.name(),
                                    signature.parameters(),
                                    signature.returnKind(),
                                    signature.returnType(),
                                    signature.resultErrorType(),
                                    signature.span()));
                }
                contracts.put(contractDecl.name(), new ContractInfo(methods));
                return;
            }
            if (topDecl instanceof PbsAst.FunctionDecl functionDecl) {
                topLevelCallables.computeIfAbsent(functionDecl.name(), ignored -> new ArrayList<>())
                        .add(callableFrom(
                                functionDecl.name(),
                                functionDecl.parameters(),
                                functionDecl.returnKind(),
                                functionDecl.returnType(),
                                functionDecl.resultErrorType(),
                                functionDecl.span()));
                return;
            }
            if (topDecl instanceof PbsAst.CallbackDecl(
                    String name, ReadOnlyList<PbsAst.Parameter> parameters, PbsAst.ReturnKind returnKind,
                    PbsAst.TypeRef returnType, PbsAst.TypeRef resultErrorType, Span span
            )) {
                final var symbol = callableFrom(
                        name,
                        parameters,
                        returnKind,
                        returnType,
                        resultErrorType,
                        span);
                callbacks.put(name, new CallbackSignature(symbol.inputTypes(), symbol.outputType()));
                return;
            }
            if (topDecl instanceof PbsAst.EnumDecl enumDecl) {
                final var cases = new HashSet<String>();
                for (final var enumCase : enumDecl.cases()) {
                    cases.add(enumCase.name());
                }
                enums.put(enumDecl.name(), cases);
                return;
            }
            if (topDecl instanceof PbsAst.ErrorDecl errorDecl) {
                errors.put(errorDecl.name(), new HashSet<>(errorDecl.cases().asList()));
                return;
            }
            if (topDecl instanceof PbsAst.ConstDecl constDecl && constDecl.explicitType() != null) {
                constTypes.put(constDecl.name(), typeFrom(constDecl.explicitType()));
                return;
            }
            if (topDecl instanceof PbsAst.GlobalDecl globalDecl && globalDecl.explicitType() != null) {
                globalTypes.put(globalDecl.name(), typeFrom(globalDecl.explicitType()));
            }
        }

        private void ingestAssetSurface(
                final FESurfaceContext feSurfaceContext,
                final DiagnosticSink diagnostics) {
            final var effectiveSurfaceContext = feSurfaceContext == null ? FESurfaceContext.empty() : feSurfaceContext;
            for (final var addressable : effectiveSurfaceContext.assets()) {
                if (addressable == null || addressable.address().isBlank()) {
                    continue;
                }
                assetIdByAddress.put(addressable.address(), addressable.assetId());
                registerAssetNamespaces(addressable.address());
            }
            for (final var address : assetIdByAddress.keySet()) {
                for (final var prefix : namespacePrefixesOf(address)) {
                    if (assetIdByAddress.containsKey(prefix)) {
                        Diagnostics.error(
                                diagnostics,
                                PbsSemanticsErrors.E_SEM_INVALID_MEMBER_ACCESS.name(),
                                "Asset surface collision: '%s' cannot be both terminal and namespace prefix".formatted(prefix),
                                astSyntheticSpan());
                    }
                }
            }
        }

        boolean hasAssetNamespaceRoot() {
            return assetNamespacePrefixes.contains("assets") || assetIdByAddress.containsKey("assets");
        }

        boolean isAssetNamespace(final String path) {
            return assetNamespacePrefixes.contains(path);
        }

        Integer resolveAssetId(final String address) {
            return assetIdByAddress.get(address);
        }

        private void registerAssetNamespaces(final String address) {
            for (final var prefix : namespacePrefixesOf(address)) {
                assetNamespacePrefixes.add(prefix);
            }
        }

        private List<String> namespacePrefixesOf(final String address) {
            final var segments = address.split("\\.");
            final var prefixes = new ArrayList<String>();
            if (segments.length <= 1) {
                return prefixes;
            }
            final var builder = new StringBuilder();
            for (int i = 0; i < segments.length - 1; i++) {
                if (segments[i] == null || segments[i].isBlank()) {
                    break;
                }
                if (builder.length() > 0) {
                    builder.append('.');
                }
                builder.append(segments[i]);
                prefixes.add(builder.toString());
            }
            return prefixes;
        }

        private Span astSyntheticSpan() {
            return Span.none();
        }
    
        private CallableSymbol callableFrom(
                final String name,
                final ReadOnlyList<PbsAst.Parameter> parameters,
                final PbsAst.ReturnKind returnKind,
                final PbsAst.TypeRef returnType,
                final PbsAst.TypeRef resultErrorType,
                final Span span) {
            final var input = new ArrayList<TypeView>(parameters.size());
            for (final var parameter : parameters) {
                input.add(typeFrom(parameter.typeRef()));
            }
            return new CallableSymbol(name, input, callableReturn(returnKind, returnType, resultErrorType), span);
        }
    
        private TypeView callableReturn(
                final PbsAst.ReturnKind returnKind,
                final PbsAst.TypeRef returnType,
                final PbsAst.TypeRef resultErrorType) {
            return switch (returnKind) {
                case INFERRED_UNIT, EXPLICIT_UNIT -> TypeView.unit();
                case PLAIN -> collapse(typeFrom(returnType));
                case RESULT -> TypeView.result(typeFrom(resultErrorType), collapse(typeFrom(returnType)));
            };
        }
    
        private TypeView collapse(final TypeView type) {
            if (type.kind() == Kind.TUPLE && type.tupleFields().size() == 1) {
                return type.tupleFields().getFirst().type();
            }
            return type;
        }
    
        private TypeView typeFrom(final PbsAst.TypeRef typeRef) {
            if (typeRef == null) {
                return TypeView.unit();
            }
            return switch (typeRef.kind()) {
                case UNIT -> TypeView.unit();
                case SELF, ERROR -> TypeView.unknown();
                case SIMPLE -> {
                    if ("int".equals(typeRef.name())) {
                        yield TypeView.intType();
                    }
                    if ("float".equals(typeRef.name())) {
                        yield TypeView.floatType();
                    }
                    if ("bool".equals(typeRef.name())) {
                        yield TypeView.bool();
                    }
                    if ("str".equals(typeRef.name())) {
                        yield TypeView.str();
                    }
                    if ("Addressable".equals(typeRef.name())) {
                        yield TypeView.addressable(typeRef.name());
                    }
                    if (structs.containsKey(typeRef.name()) || knownStructNames.contains(typeRef.name())) {
                        yield TypeView.struct(typeRef.name());
                    }
                    if (services.containsKey(typeRef.name()) || knownServiceNames.contains(typeRef.name())) {
                        yield TypeView.service(typeRef.name());
                    }
                    if (contracts.containsKey(typeRef.name()) || knownContractNames.contains(typeRef.name())) {
                        yield TypeView.contract(typeRef.name());
                    }
                    if (enums.containsKey(typeRef.name()) || knownEnumNames.contains(typeRef.name())) {
                        yield TypeView.enumType(typeRef.name());
                    }
                    if (errors.containsKey(typeRef.name()) || knownErrorNames.contains(typeRef.name())) {
                        yield TypeView.error(typeRef.name());
                    }
                    final var callbackSignature = callbacks.get(typeRef.name());
                    if (callbackSignature != null) {
                        yield TypeView.callback(typeRef.name(), callbackSignature.inputTypes(), callbackSignature.outputType());
                    }
                    if (knownCallbackNames.contains(typeRef.name())) {
                        yield TypeView.callback(typeRef.name(), List.of(), TypeView.unknown());
                    }
                    yield TypeView.unknown();
                }
                case OPTIONAL -> TypeView.optional(typeFrom(typeRef.inner()));
                case GROUP -> typeFrom(typeRef.inner());
                case NAMED_TUPLE -> {
                    final var fields = new ArrayList<TupleField>(typeRef.fields().size());
                    for (final var field : typeRef.fields()) {
                        fields.add(new TupleField(field.label(), typeFrom(field.typeRef())));
                    }
                    yield TypeView.tuple(fields);
                }
            };
        }
    }
    
    static final class Scope {
        private final Map<String, LocalSymbol> names = new HashMap<>();

        record LocalSymbol(
                TypeView type,
                boolean mutable) {
        }
    
        Scope copy() {
            final var scope = new Scope();
            scope.names.putAll(names);
            return scope;
        }
    
        void bind(final String name, final TypeView type) {
            bind(name, type, true);
        }

        void bind(final String name, final TypeView type, final boolean mutable) {
            names.put(name, new LocalSymbol(type, mutable));
        }

        TypeView resolve(final String name) {
            final var symbol = names.get(name);
            return symbol == null ? null : symbol.type();
        }

        boolean isMutable(final String name) {
            final var symbol = names.get(name);
            return symbol != null && symbol.mutable();
        }
    }
}
