package p.studio.compiler.pbs.lowering;

import p.studio.compiler.models.IRReservedMetadata;
import p.studio.compiler.messages.Addressable;
import p.studio.compiler.messages.FESurfaceContext;
import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.semantics.PbsSemanticsErrors;
import p.studio.compiler.source.diagnostics.DiagnosticPhase;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.NameId;
import p.studio.compiler.source.tables.NameTable;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static p.studio.compiler.pbs.lowering.PbsExecutableLoweringModels.*;

final class PbsExecutableMetadataIndexFactory {

    PbsExecutableMetadataIndex create(
            final PbsAst.File ast,
            final ReadOnlyList<PbsAst.TopDecl> supplementalTopDecls,
            final IRReservedMetadata reservedMetadata,
            final NameTable nameTable,
            final FESurfaceContext feSurfaceContext,
            final DiagnosticSink diagnostics) {
        final var hostByMethodName = indexHostBindingsByMethodName(reservedMetadata, nameTable);
        final var builtinCanonicalBySourceType = indexBuiltinCanonicalBySourceType(reservedMetadata);
        final var builtinConstOwnerByNameId = indexBuiltinConstOwners(reservedMetadata, nameTable);
        final var globalSlotByNameId = indexGlobalSlots(ast, nameTable);
        final var constDeclByNameId = indexConstDecls(ast, nameTable);
        final var addressableByAddress = indexAddressables(feSurfaceContext);
        final var builtinSignatureByOwnerAndMethod = indexBuiltinSignatures(
                ast,
                supplementalTopDecls,
                builtinCanonicalBySourceType);
        final var intrinsicByOwnerAndMethod = new HashMap<PbsIntrinsicOwnerMethodKey, List<IRReservedMetadata.IntrinsicSurface>>();
        final var intrinsicReturnOwnerByCanonical = new HashMap<PbsIntrinsicCanonicalKey, String>();
        indexIntrinsicMetadata(
                reservedMetadata,
                builtinCanonicalBySourceType,
                builtinSignatureByOwnerAndMethod,
                intrinsicByOwnerAndMethod,
                intrinsicReturnOwnerByCanonical,
                diagnostics);
        return new PbsExecutableMetadataIndex(
                hostByMethodName,
                builtinCanonicalBySourceType,
                builtinConstOwnerByNameId,
                globalSlotByNameId,
                constDeclByNameId,
                addressableByAddress,
                intrinsicByOwnerAndMethod,
                intrinsicReturnOwnerByCanonical);
    }

    private Map<NameId, List<IRReservedMetadata.HostMethodBinding>> indexHostBindingsByMethodName(
            final IRReservedMetadata reservedMetadata,
            final NameTable nameTable) {
        final var hostByMethodName = new HashMap<NameId, List<IRReservedMetadata.HostMethodBinding>>();
        for (final var hostBinding : reservedMetadata.hostMethodBindings()) {
            hostByMethodName
                    .computeIfAbsent(nameTable.register(hostBinding.sourceMethodName()), ignored -> new ArrayList<>())
                    .add(hostBinding);
        }
        return hostByMethodName;
    }

    private Map<String, String> indexBuiltinCanonicalBySourceType(final IRReservedMetadata reservedMetadata) {
        final var builtinCanonicalBySourceType = new HashMap<String, String>();
        for (final var builtinType : reservedMetadata.builtinTypeSurfaces()) {
            builtinCanonicalBySourceType.putIfAbsent(builtinType.sourceTypeName(), builtinType.canonicalTypeName());
        }
        return builtinCanonicalBySourceType;
    }

    private Map<NameId, String> indexBuiltinConstOwners(
            final IRReservedMetadata reservedMetadata,
            final NameTable nameTable) {
        final var builtinConstOwnerByNameId = new HashMap<NameId, String>();
        for (final var builtinConst : reservedMetadata.builtinConstSurfaces()) {
            if (builtinConst.sourceConstName().isBlank() || builtinConst.canonicalTarget().isBlank()) {
                continue;
            }
            builtinConstOwnerByNameId.put(nameTable.register(builtinConst.sourceConstName()), builtinConst.canonicalTarget());
        }
        return builtinConstOwnerByNameId;
    }

    private Map<NameId, Integer> indexGlobalSlots(
            final PbsAst.File ast,
            final NameTable nameTable) {
        final var globalSlotByNameId = new HashMap<NameId, Integer>();
        var nextSlot = 0;
        for (final var topDecl : ast.topDecls()) {
            if (!(topDecl instanceof PbsAst.GlobalDecl globalDecl)) {
                continue;
            }
            globalSlotByNameId.putIfAbsent(nameTable.register(globalDecl.name()), nextSlot++);
        }
        return globalSlotByNameId;
    }

    private Map<NameId, PbsAst.ConstDecl> indexConstDecls(
            final PbsAst.File ast,
            final NameTable nameTable) {
        final var constDeclByNameId = new HashMap<NameId, PbsAst.ConstDecl>();
        for (final var topDecl : ast.topDecls()) {
            if (!(topDecl instanceof PbsAst.ConstDecl constDecl)) {
                continue;
            }
            constDeclByNameId.putIfAbsent(nameTable.register(constDecl.name()), constDecl);
        }
        return constDeclByNameId;
    }

    private Map<String, Addressable> indexAddressables(final FESurfaceContext feSurfaceContext) {
        final var addressableByAddress = new HashMap<String, Addressable>();
        final var effectiveSurfaceContext = feSurfaceContext == null ? FESurfaceContext.empty() : feSurfaceContext;
        for (final var addressable : effectiveSurfaceContext.assets()) {
            if (addressable == null || addressable.address().isBlank()) {
                continue;
            }
            addressableByAddress.putIfAbsent(addressable.address(), addressable);
        }
        return addressableByAddress;
    }

    private Map<PbsIntrinsicOwnerMethodKey, PbsAst.FunctionSignature> indexBuiltinSignatures(
            final PbsAst.File ast,
            final ReadOnlyList<PbsAst.TopDecl> supplementalTopDecls,
            final Map<String, String> builtinCanonicalBySourceType) {
        final var builtinSignatureByOwnerAndMethod = new HashMap<PbsIntrinsicOwnerMethodKey, PbsAst.FunctionSignature>();
        final var topDeclsForIntrinsicInference = new ArrayList<PbsAst.TopDecl>(ast.topDecls().size() + supplementalTopDecls.size());
        topDeclsForIntrinsicInference.addAll(ast.topDecls().asList());
        topDeclsForIntrinsicInference.addAll(supplementalTopDecls.asList());
        for (final var topDecl : topDeclsForIntrinsicInference) {
            if (!(topDecl instanceof PbsAst.BuiltinTypeDecl builtinTypeDecl)) {
                continue;
            }
            final var ownerCanonical = builtinCanonicalBySourceType.getOrDefault(builtinTypeDecl.name(), "");
            if (ownerCanonical.isBlank()) {
                continue;
            }
            for (final var signature : builtinTypeDecl.signatures()) {
                builtinSignatureByOwnerAndMethod.putIfAbsent(
                        new PbsIntrinsicOwnerMethodKey(ownerCanonical, signature.name()),
                        signature);
            }
        }
        return builtinSignatureByOwnerAndMethod;
    }

    private void indexIntrinsicMetadata(
            final IRReservedMetadata reservedMetadata,
            final Map<String, String> builtinCanonicalBySourceType,
            final Map<PbsIntrinsicOwnerMethodKey, PbsAst.FunctionSignature> builtinSignatureByOwnerAndMethod,
            final Map<PbsIntrinsicOwnerMethodKey, List<IRReservedMetadata.IntrinsicSurface>> intrinsicByOwnerAndMethod,
            final Map<PbsIntrinsicCanonicalKey, String> intrinsicReturnOwnerByCanonical,
            final DiagnosticSink diagnostics) {
        for (final var builtinType : reservedMetadata.builtinTypeSurfaces()) {
            final var ownerCanonical = builtinType.canonicalTypeName();
            for (final var intrinsicSurface : builtinType.intrinsics()) {
                intrinsicByOwnerAndMethod
                        .computeIfAbsent(
                                new PbsIntrinsicOwnerMethodKey(ownerCanonical, intrinsicSurface.sourceMethodName()),
                                ignored -> new ArrayList<>())
                        .add(intrinsicSurface);

                final var signature = builtinSignatureByOwnerAndMethod.get(
                        new PbsIntrinsicOwnerMethodKey(ownerCanonical, intrinsicSurface.sourceMethodName()));
                if (signature == null) {
                    continue;
                }
                final var returnOwner = inferIntrinsicReturnOwner(
                        signature,
                        ownerCanonical,
                        builtinCanonicalBySourceType);
                if (returnOwner.isBlank()) {
                    continue;
                }
                final var key = new PbsIntrinsicCanonicalKey(
                        intrinsicSurface.canonicalName(),
                        intrinsicSurface.canonicalVersion());
                final var previous = intrinsicReturnOwnerByCanonical.putIfAbsent(key, returnOwner);
                if (previous != null && !previous.equals(returnOwner)) {
                    diagnostics.error(
                            DiagnosticPhase.STATIC_SEMANTICS,
                            PbsSemanticsErrors.E_SEM_EXEC_LOWERING_UNRESOLVED_CALLEE.name(),
                            PbsSemanticsErrors.E_SEM_EXEC_LOWERING_UNRESOLVED_CALLEE.name(),
                            Map.of(),
                            "inconsistent intrinsic return owner metadata for %s@%d: %s vs %s".formatted(
                                    intrinsicSurface.canonicalName(),
                                    intrinsicSurface.canonicalVersion(),
                                    previous,
                                    returnOwner),
                            intrinsicSurface.span());
                }
            }
        }
    }

    private String inferIntrinsicReturnOwner(
            final PbsAst.FunctionSignature signature,
            final String intrinsicOwnerCanonical,
            final Map<String, String> builtinCanonicalBySourceType) {
        if (signature.returnKind() == PbsAst.ReturnKind.INFERRED_UNIT
                || signature.returnKind() == PbsAst.ReturnKind.EXPLICIT_UNIT) {
            return "";
        }
        final var returnType = unwrapGroup(signature.returnType());
        if (returnType == null) {
            return "";
        }
        return switch (returnType.kind()) {
            case SELF -> intrinsicOwnerCanonical == null ? "" : intrinsicOwnerCanonical;
            case SIMPLE -> builtinCanonicalBySourceType.getOrDefault(returnType.name(), "");
            default -> "";
        };
    }

    private PbsAst.TypeRef unwrapGroup(final PbsAst.TypeRef typeRef) {
        if (typeRef == null) {
            return null;
        }
        if (typeRef.kind() != PbsAst.TypeRefKind.GROUP) {
            return typeRef;
        }
        return unwrapGroup(typeRef.inner());
    }
}
