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.utilities.structures.ReadOnlyList;

import java.util.*;

public final class PbsBuiltinLayoutResolver {
    private static final Set<String> BUILTIN_SCALAR_TYPES = Set.of("int", "float", "bool", "str");

    private enum ResolveState {
        UNVISITED,
        RESOLVING,
        RESOLVED,
        FAILED
    }

    private enum WidthStatus {
        OK,
        CYCLE,
        FAILED
    }

    private final Map<String, PbsAst.BuiltinTypeDecl> builtinDeclByName = new LinkedHashMap<>();
    private final Map<String, ResolveState> stateByBuiltinName = new HashMap<>();
    private final Map<String, ResolvedBuiltinLayout> resolvedLayouts = new HashMap<>();
    private final Set<String> reportedCycles = new HashSet<>();
    private final DiagnosticSink diagnostics;

    private PbsBuiltinLayoutResolver(
            final PbsAst.File ast,
            final DiagnosticSink diagnostics) {
        this.diagnostics = diagnostics;
        for (final var topDecl : ast.topDecls()) {
            if (!(topDecl instanceof PbsAst.BuiltinTypeDecl builtinTypeDecl)) {
                continue;
            }
            builtinDeclByName.putIfAbsent(builtinTypeDecl.name(), builtinTypeDecl);
        }
        for (final var builtinTypeName : builtinDeclByName.keySet()) {
            stateByBuiltinName.put(builtinTypeName, ResolveState.UNVISITED);
        }
    }

    public static ResolvedBuiltinLayoutTable resolve(final PbsAst.File ast) {
        return new PbsBuiltinLayoutResolver(ast, null).resolveAll();
    }

    public static ResolvedBuiltinLayoutTable resolve(
            final PbsAst.File ast,
            final DiagnosticSink diagnostics) {
        return new PbsBuiltinLayoutResolver(ast, diagnostics).resolveAll();
    }

    private ResolvedBuiltinLayoutTable resolveAll() {
        for (final var builtinTypeName : builtinDeclByName.keySet()) {
            resolveBuiltinLayout(builtinTypeName);
        }
        return new ResolvedBuiltinLayoutTable(ReadOnlyList.wrap(new ArrayList<>(resolvedLayouts.values())));
    }

    private WidthResult resolveBuiltinLayout(final String builtinTypeName) {
        final var state = stateByBuiltinName.getOrDefault(builtinTypeName, ResolveState.UNVISITED);
        if (state == ResolveState.RESOLVED) {
            return WidthResult.ok(resolvedLayouts.get(builtinTypeName).totalWidth());
        }
        if (state == ResolveState.FAILED) {
            return WidthResult.failed();
        }
        if (state == ResolveState.RESOLVING) {
            return WidthResult.cycle();
        }

        final var declaration = builtinDeclByName.get(builtinTypeName);
        if (declaration == null) {
            return WidthResult.failed();
        }

        stateByBuiltinName.put(builtinTypeName, ResolveState.RESOLVING);

        long slotCursor = 0L;
        var failed = false;
        final var fieldLayouts = new ArrayList<ResolvedBuiltinFieldLayout>(declaration.fields().size());
        for (final var field : declaration.fields()) {
            final var width = resolveFieldWidth(builtinTypeName, field);
            if (width.status() != WidthStatus.OK) {
                failed = true;
                continue;
            }
            fieldLayouts.add(new ResolvedBuiltinFieldLayout(
                    field.name(),
                    slotCursor,
                    width.width(),
                    field.span()));
            slotCursor += width.width();
        }

        if (failed) {
            stateByBuiltinName.put(builtinTypeName, ResolveState.FAILED);
            return WidthResult.failed();
        }

        final var layout = new ResolvedBuiltinLayout(
                declaration.name(),
                slotCursor,
                ReadOnlyList.wrap(fieldLayouts));
        resolvedLayouts.put(builtinTypeName, layout);
        stateByBuiltinName.put(builtinTypeName, ResolveState.RESOLVED);
        return WidthResult.ok(slotCursor);
    }

    private WidthResult resolveFieldWidth(
            final String ownerBuiltinName,
            final PbsAst.BuiltinFieldDecl field) {
        final var typeRef = normalizeFieldType(field.typeRef());
        if (typeRef == null || typeRef.kind() != PbsAst.TypeRefKind.SIMPLE) {
            reportTypeNotAdmissible(ownerBuiltinName, field, field.typeRef());
            return WidthResult.failed();
        }

        if (BUILTIN_SCALAR_TYPES.contains(typeRef.name())) {
            return WidthResult.ok(1L);
        }

        final var nestedBuiltin = builtinDeclByName.get(typeRef.name());
        if (nestedBuiltin == null) {
            reportTypeNotAdmissible(ownerBuiltinName, field, field.typeRef());
            return WidthResult.failed();
        }

        final var nestedWidth = resolveBuiltinLayout(typeRef.name());
        if (nestedWidth.status() == WidthStatus.OK) {
            return nestedWidth;
        }
        if (nestedWidth.status() == WidthStatus.CYCLE) {
            reportCycleIfNeeded(typeRef.name(), ownerBuiltinName, field.name(), field.span());
            return WidthResult.cycle();
        }

        reportUnresolvableWidth(ownerBuiltinName, field, typeRef.name());
        return WidthResult.failed();
    }

    private PbsAst.TypeRef normalizeFieldType(final PbsAst.TypeRef typeRef) {
        var current = typeRef;
        while (current != null && current.kind() == PbsAst.TypeRefKind.GROUP) {
            current = current.inner();
        }
        return current;
    }

    private void reportTypeNotAdmissible(
            final String ownerBuiltinName,
            final PbsAst.BuiltinFieldDecl field,
            final PbsAst.TypeRef typeRef) {
        if (diagnostics == null) {
            return;
        }
        p.studio.compiler.source.diagnostics.Diagnostics.error(
                diagnostics,
                PbsSemanticsErrors.E_SEM_BUILTIN_LAYOUT_FIELD_TYPE_NOT_ADMISSIBLE.name(),
                "Builtin field '%s.%s' uses non-admissible layout type '%s'"
                        .formatted(ownerBuiltinName, field.name(), renderTypeSurface(typeRef)),
                field.span());
    }

    private void reportCycleIfNeeded(
            final String targetBuiltinName,
            final String ownerBuiltinName,
            final String ownerFieldName,
            final Span span) {
        if (diagnostics == null) {
            return;
        }
        final String cycleKey = ownerBuiltinName == null
                ? "cycle:" + targetBuiltinName
                : "cycle:" + ownerBuiltinName + "." + ownerFieldName + "->" + targetBuiltinName;
        if (!reportedCycles.add(cycleKey)) {
            return;
        }
        final var primarySpan = span != null
                ? span
                : Optional.ofNullable(builtinDeclByName.get(targetBuiltinName))
                .map(PbsAst.BuiltinTypeDecl::span)
                .orElse(Span.none());
        p.studio.compiler.source.diagnostics.Diagnostics.error(
                diagnostics,
                PbsSemanticsErrors.E_SEM_BUILTIN_LAYOUT_CYCLIC_COMPOSITION.name(),
                ownerBuiltinName == null
                        ? "Builtin layout cycle detected for '%s'".formatted(targetBuiltinName)
                        : "Builtin field '%s.%s' introduces cyclic builtin layout dependency on '%s'"
                        .formatted(ownerBuiltinName, ownerFieldName, targetBuiltinName),
                primarySpan);
    }

    private void reportUnresolvableWidth(
            final String ownerBuiltinName,
            final PbsAst.BuiltinFieldDecl field,
            final String nestedBuiltinName) {
        if (diagnostics == null) {
            return;
        }
        p.studio.compiler.source.diagnostics.Diagnostics.error(
                diagnostics,
                PbsSemanticsErrors.E_SEM_BUILTIN_LAYOUT_UNRESOLVABLE_WIDTH.name(),
                "Builtin field '%s.%s' cannot resolve flattened layout width for '%s'"
                        .formatted(ownerBuiltinName, field.name(), nestedBuiltinName),
                field.span());
    }

    private String renderTypeSurface(final PbsAst.TypeRef typeRef) {
        if (typeRef == null) {
            return "void";
        }
        return switch (typeRef.kind()) {
            case SIMPLE -> typeRef.name();
            case SELF -> "Self";
            case OPTIONAL -> "optional %s".formatted(renderTypeSurface(typeRef.inner()));
            case UNIT -> "void";
            case GROUP -> "(%s)".formatted(renderTypeSurface(typeRef.inner()));
            case NAMED_TUPLE -> "(%s)".formatted(typeRef.fields().stream()
                    .map(field -> "%s: %s".formatted(field.label(), renderTypeSurface(field.typeRef())))
                    .reduce((left, right) -> left + ", " + right)
                    .orElse(""));
            case ERROR -> "<error>";
        };
    }

    private record WidthResult(
            WidthStatus status,
            long width) {
        static WidthResult ok(final long width) {
            return new WidthResult(WidthStatus.OK, width);
        }

        static WidthResult cycle() {
            return new WidthResult(WidthStatus.CYCLE, 0L);
        }

        static WidthResult failed() {
            return new WidthResult(WidthStatus.FAILED, 0L);
        }
    }

    public record ResolvedBuiltinLayoutTable(
            ReadOnlyList<ResolvedBuiltinLayout> layouts) {
        public Optional<ResolvedBuiltinLayout> findByName(final String builtinTypeName) {
            for (final var layout : layouts) {
                if (layout.sourceTypeName().equals(builtinTypeName)) {
                    return Optional.of(layout);
                }
            }
            return Optional.empty();
        }
    }

    public record ResolvedBuiltinLayout(
            String sourceTypeName,
            long totalWidth,
            ReadOnlyList<ResolvedBuiltinFieldLayout> fields) {
    }

    public record ResolvedBuiltinFieldLayout(
            String fieldName,
            long slotStart,
            long slotWidth,
            Span span) {
    }
}
