package p.studio.compiler.pbs;

import p.studio.compiler.models.IRReservedMetadata;
import p.studio.compiler.models.SourceKind;
import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.semantics.PbsBuiltinLayoutResolver;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.*;
import java.util.stream.Collectors;

public final class PbsReservedMetadataExtractor {
    private static final String ATTR_HOST = "Host";
    private static final String ATTR_CAPABILITY = "Capability";
    private static final String ATTR_BUILTIN_TYPE = "BuiltinType";
    private static final String ATTR_INTRINSIC_CALL = "IntrinsicCall";
    private static final String ATTR_BUILTIN_CONST = "BuiltinConst";
    private static final String ATTR_ASSET_LOWERING = "AssetLowering";

    public IRReservedMetadata extract(
            final PbsAst.File ast,
            final SourceKind sourceKind) {
        if (sourceKind != SourceKind.SDK_INTERFACE) {
            return IRReservedMetadata.empty();
        }

        final var hostMethodBindings = new ArrayList<IRReservedMetadata.HostMethodBinding>();
        final var builtinTypeSurfaces = new ArrayList<IRReservedMetadata.BuiltinTypeSurface>();
        final var builtinConstSurfaces = new ArrayList<IRReservedMetadata.BuiltinConstSurface>();
        final var inferredLayoutsByName = PbsBuiltinLayoutResolver.resolve(ast).layouts().stream()
                .collect(Collectors.toMap(
                        PbsBuiltinLayoutResolver.ResolvedBuiltinLayout::sourceTypeName,
                        layout -> layout));

        for (final var topDecl : ast.topDecls()) {
            if (topDecl instanceof PbsAst.HostDecl hostDecl) {
                extractHostBindings(hostDecl, hostMethodBindings);
                continue;
            }
            if (topDecl instanceof PbsAst.BuiltinTypeDecl builtinTypeDecl) {
                extractBuiltinTypeSurface(builtinTypeDecl, inferredLayoutsByName, builtinTypeSurfaces);
                continue;
            }
            if (topDecl instanceof PbsAst.ConstDecl constDecl) {
                extractBuiltinConstSurface(constDecl, builtinConstSurfaces);
            }
        }

        return new IRReservedMetadata(
                ReadOnlyList.wrap(hostMethodBindings),
                ReadOnlyList.wrap(builtinTypeSurfaces),
                ReadOnlyList.wrap(builtinConstSurfaces),
                requiredCapabilities(hostMethodBindings));
    }

    private void extractHostBindings(
            final PbsAst.HostDecl hostDecl,
            final List<IRReservedMetadata.HostMethodBinding> hostMethodBindings) {
        for (final var signature : hostDecl.signatures()) {
            final var hostAttribute = firstAttributeNamed(signature.attributes(), ATTR_HOST);
            if (hostAttribute.isEmpty()) {
                continue;
            }
            final var hostMetadata = hostAttribute.get();
            final var abiModule = stringArgument(hostMetadata, "module").orElse("");
            final var abiMethod = stringArgument(hostMetadata, "name").orElse(signature.name());
            final var abiVersion = longArgument(hostMetadata, "version").orElse(0L);
            final var capabilityAttribute = firstAttributeNamed(signature.attributes(), ATTR_CAPABILITY);
            final var capability = capabilityAttribute
                    .flatMap(attr -> stringArgument(attr, "name"))
                    .orElse("");
            final var assetLoweringAttribute = firstAttributeNamed(signature.attributes(), ATTR_ASSET_LOWERING);
            hostMethodBindings.add(new IRReservedMetadata.HostMethodBinding(
                    hostDecl.name(),
                    signature.name(),
                    abiModule,
                    abiMethod,
                    abiVersion,
                    PbsCallableSlotCounter.returnSlots(signature.returnKind(), signature.returnType()),
                    assetLoweringAttribute.isPresent()
                            ? safeToInteger(longArgument(assetLoweringAttribute.get(), "param").orElse(-1L))
                            : null,
                    capabilityAttribute.isPresent(),
                    capability,
                    signature.span()));
        }
    }

    private void extractBuiltinTypeSurface(
            final PbsAst.BuiltinTypeDecl builtinTypeDecl,
            final Map<String, PbsBuiltinLayoutResolver.ResolvedBuiltinLayout> inferredLayoutsByName,
            final List<IRReservedMetadata.BuiltinTypeSurface> builtinTypeSurfaces) {
        final var builtinTypeAttribute = firstAttributeNamed(builtinTypeDecl.attributes(), ATTR_BUILTIN_TYPE);
        if (builtinTypeAttribute.isEmpty()) {
            return;
        }

        final var builtinTypeMetadata = builtinTypeAttribute.get();
        final var canonicalTypeName = stringArgument(builtinTypeMetadata, "name").orElse(builtinTypeDecl.name());
        final var canonicalVersion = longArgument(builtinTypeMetadata, "version").orElse(0L);
        final var inferredLayout = inferredLayoutsByName.get(builtinTypeDecl.name());

        final var fields = new ArrayList<IRReservedMetadata.BuiltinFieldSurface>(builtinTypeDecl.fields().size());
        for (int i = 0; i < builtinTypeDecl.fields().size(); i++) {
            final var field = builtinTypeDecl.fields().get(i);
            final var slotStart = inferredLayout != null && i < inferredLayout.fields().size()
                    ? inferredLayout.fields().get(i).slotStart()
                    : i;
            fields.add(new IRReservedMetadata.BuiltinFieldSurface(
                    field.name(),
                    typeSurfaceKey(field.typeRef()),
                    slotStart,
                    field.span()));
        }

        final var intrinsics = new ArrayList<IRReservedMetadata.IntrinsicSurface>(builtinTypeDecl.signatures().size());
        for (final var signature : builtinTypeDecl.signatures()) {
            final var intrinsicAttribute = firstAttributeNamed(signature.attributes(), ATTR_INTRINSIC_CALL);
            if (intrinsicAttribute.isEmpty()) {
                continue;
            }
            final var intrinsicMetadata = intrinsicAttribute.get();
            final var intrinsicName = stringArgument(intrinsicMetadata, "name").orElse(signature.name());
            intrinsics.add(new IRReservedMetadata.IntrinsicSurface(
                    signature.name(),
                    canonicalIntrinsicName(canonicalTypeName, intrinsicName),
                    longArgument(intrinsicMetadata, "version").orElse(canonicalVersion),
                    signature.parameters().size(),
                    PbsCallableSlotCounter.returnSlots(signature.returnKind(), signature.returnType()),
                    signature.span()));
        }

        builtinTypeSurfaces.add(new IRReservedMetadata.BuiltinTypeSurface(
                builtinTypeDecl.name(),
                canonicalTypeName,
                canonicalVersion,
                ReadOnlyList.wrap(fields),
                ReadOnlyList.wrap(intrinsics),
                builtinTypeDecl.span()));
    }

    private void extractBuiltinConstSurface(
            final PbsAst.ConstDecl constDecl,
            final List<IRReservedMetadata.BuiltinConstSurface> builtinConstSurfaces) {
        final var builtinConstAttribute = firstAttributeNamed(constDecl.attributes(), ATTR_BUILTIN_CONST);
        if (builtinConstAttribute.isEmpty()) {
            return;
        }

        final var builtinConstMetadata = builtinConstAttribute.get();
        builtinConstSurfaces.add(new IRReservedMetadata.BuiltinConstSurface(
                constDecl.name(),
                stringArgument(builtinConstMetadata, "target").orElse(""),
                stringArgument(builtinConstMetadata, "name").orElse(constDecl.name()),
                longArgument(builtinConstMetadata, "version").orElse(0L),
                constDecl.span()));
    }

    private Optional<PbsAst.Attribute> firstAttributeNamed(
            final ReadOnlyList<PbsAst.Attribute> attributes,
            final String attributeName) {
        for (final var attribute : attributes) {
            if (attributeName.equals(attribute.name())) {
                return Optional.of(attribute);
            }
        }
        return Optional.empty();
    }

    private Optional<String> stringArgument(
            final PbsAst.Attribute attribute,
            final String argumentName) {
        final var argument = findArgument(attribute, argumentName);
        if (argument.isEmpty()) {
            return Optional.empty();
        }
        if (argument.get().value() instanceof PbsAst.AttributeStringValue value) {
            return Optional.of(normalizeStringLiteralValue(value.value()));
        }
        return Optional.empty();
    }

    private OptionalLong longArgument(
            final PbsAst.Attribute attribute,
            final String argumentName) {
        final var argument = findArgument(attribute, argumentName);
        if (argument.isEmpty()) {
            return OptionalLong.empty();
        }
        if (argument.get().value() instanceof PbsAst.AttributeIntValue value) {
            return OptionalLong.of(value.value());
        }
        return OptionalLong.empty();
    }

    private Optional<PbsAst.AttributeArgument> findArgument(
            final PbsAst.Attribute attribute,
            final String argumentName) {
        for (final var argument : attribute.arguments()) {
            if (argumentName.equals(argument.name())) {
                return Optional.of(argument);
            }
        }
        return Optional.empty();
    }

    private String normalizeStringLiteralValue(final String rawValue) {
        if (rawValue == null) {
            return "";
        }
        if (rawValue.length() >= 2 && rawValue.startsWith("\"") && rawValue.endsWith("\"")) {
            return rawValue.substring(1, rawValue.length() - 1);
        }
        return rawValue;
    }

    private String canonicalIntrinsicName(
            final String canonicalTypeName,
            final String intrinsicName) {
        final var normalizedTypeName = canonicalTypeName == null ? "" : canonicalTypeName.trim();
        final var normalizedIntrinsicName = intrinsicName == null ? "" : intrinsicName.trim();
        if (normalizedIntrinsicName.isBlank()) {
            return normalizedTypeName;
        }
        if (normalizedIntrinsicName.contains(".")) {
            return normalizedIntrinsicName;
        }
        if (normalizedTypeName.isBlank()) {
            return normalizedIntrinsicName;
        }
        return normalizedTypeName + "." + normalizedIntrinsicName;
    }

    public static String typeSurfaceKey(final PbsAst.TypeRef typeRef) {
        if (typeRef == null) {
            return "unit";
        }
        return switch (typeRef.kind()) {
            case SIMPLE -> "simple:" + typeRef.name();
            case SELF -> "self";
            case OPTIONAL -> "optional(" + typeSurfaceKey(typeRef.inner()) + ")";
            case UNIT -> "unit";
            case GROUP -> "group(" + typeSurfaceKey(typeRef.inner()) + ")";
            case NAMED_TUPLE -> "tuple(" + typeRef.fields().stream()
                    .map(field -> typeSurfaceKey(field.typeRef()))
                    .reduce((left, right) -> left + "," + right)
                    .orElse("") + ")";
            case ERROR -> "error";
        };
    }

    private ReadOnlyList<String> requiredCapabilities(
            final List<IRReservedMetadata.HostMethodBinding> hostMethodBindings) {
        final var required = new LinkedHashSet<String>();
        for (final var hostBinding : hostMethodBindings) {
            if (hostBinding.requiredCapability() == null || hostBinding.requiredCapability().isBlank()) {
                continue;
            }
            required.add(hostBinding.requiredCapability().trim().toLowerCase());
        }
        return ReadOnlyList.wrap(required.stream().toList());
    }

    private Integer safeToInteger(final long value) {
        if (value < Integer.MIN_VALUE || value > Integer.MAX_VALUE) {
            return null;
        }
        return (int) value;
    }
}
