package p.studio.compiler.pbs;

import org.junit.jupiter.api.Test;
import p.studio.compiler.messages.HostAdmissionContext;
import p.studio.compiler.models.SourceKind;
import p.studio.compiler.pbs.lexer.LexErrors;
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.FileId;
import p.studio.utilities.structures.ReadOnlyList;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

class PbsFrontendCompilerTest {

    @Test
    void shouldLowerFunctionsToIr() {
        final var source = """
                fn a() -> int {
                  return 1;
                }

                fn b(x: int) -> int {
                  return x;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(0), source, diagnostics);
        final var functions = fileBackend.functions();

        assertTrue(diagnostics.isEmpty(), "Valid program should not report diagnostics");
        assertEquals(2, functions.size());
        assertEquals("a", functions.get(0).name());
        assertEquals(0, functions.get(0).parameterCount());
        assertEquals("b", functions.get(1).name());
        assertEquals(1, functions.get(1).parameterCount());
    }

    @Test
    void shouldReportDuplicateFunctionNames() {
        final var source = """
                fn a() -> int { return 1; }
                fn a() -> int { return 2; }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(0), source, diagnostics);

        assertTrue(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(PbsSemanticsErrors.E_SEM_DUPLICATE_CALLABLE_SHAPE.name())));
        assertEquals(0, fileBackend.functions().size());
    }

    @Test
    void shouldAllowOverloadedFunctionNamesWhenShapeDiffers() {
        final var source = """
                fn a(x: int) -> int { return x; }
                fn a(x: float) -> int { return 0; }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(0), source, diagnostics);

        assertTrue(diagnostics.stream().noneMatch(d ->
                d.getCode().equals(PbsSemanticsErrors.E_SEM_DUPLICATE_CALLABLE_SHAPE.name())));
        assertEquals(2, fileBackend.functions().size());
    }

    @Test
    void shouldLowerExecutableFunctionsWithCallsiteCategories() {
        final var source = """
                fn b(x: int) -> int {
                  return x;
                }

                fn a() -> int {
                  return b(1);
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(42), source, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Valid program should not report diagnostics");
        assertEquals(2, fileBackend.executableFunctions().size());
        final var executableA = fileBackend.executableFunctions().stream()
                .filter(fn -> fn.callableName().equals("a"))
                .findFirst()
                .orElseThrow();
        assertTrue(executableA.instructions().stream().anyMatch(i ->
                i.kind() == p.studio.compiler.models.IRBackendExecutableFunction.InstructionKind.CALL_FUNC));
    }

    @Test
    void shouldClassifyCallableCallEvenWhenNameLooksHostLike() {
        final var source = """
                fn draw_pixel(x: int, y: int) -> int {
                  return x;
                }

                fn main() -> int {
                  return draw_pixel(1, 2);
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(43), source, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Valid program should not report diagnostics");
        final var executableMain = fileBackend.executableFunctions().stream()
                .filter(fn -> fn.callableName().equals("main"))
                .findFirst()
                .orElseThrow();
        final var callInstruction = executableMain.instructions().stream()
                .filter(i -> i.kind() == p.studio.compiler.models.IRBackendExecutableFunction.InstructionKind.CALL_FUNC)
                .findFirst()
                .orElseThrow();
        assertEquals("draw_pixel", callInstruction.calleeCallableName());
    }

    @Test
    void shouldNotInferCallableCategoryFromMemberCallText() {
        final var source = """
                fn draw_pixel(x: int, y: int) -> int {
                  return x;
                }

                fn main() -> int {
                  return gfx.draw_pixel(1, 2);
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(44), source, diagnostics);

        assertTrue(fileBackend.executableFunctions().stream()
                .filter(fn -> fn.callableName().equals("main"))
                .findFirst()
                .map(fn -> fn.instructions().stream().noneMatch(i ->
                        i.kind() == p.studio.compiler.models.IRBackendExecutableFunction.InstructionKind.CALL_FUNC
                                || i.kind() == p.studio.compiler.models.IRBackendExecutableFunction.InstructionKind.CALL_HOST
                                || i.kind() == p.studio.compiler.models.IRBackendExecutableFunction.InstructionKind.CALL_INTRINSIC))
                .orElse(true));
    }

    @Test
    void shouldRejectAmbiguousCallableIdentityInsteadOfGuessingByText() {
        final var source = """
                fn draw_pixel(x: int) -> int {
                  return x;
                }

                fn draw_pixel(x: float) -> int {
                  return 0;
                }

                fn main() -> int {
                  return draw_pixel(1);
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(45), source, diagnostics);

        final var loweringDiagnostic = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsSemanticsErrors.E_SEM_EXEC_LOWERING_UNRESOLVED_CALLEE.name()))
                .findFirst()
                .orElseThrow();
        assertEquals(DiagnosticPhase.STATIC_SEMANTICS, loweringDiagnostic.getPhase());
        assertTrue(loweringDiagnostic.getMessage().contains("ambiguous callable identity"));
        assertEquals(0, fileBackend.executableFunctions().size());
    }

    @Test
    void shouldLowerExecutableCallsitesInDeterministicSourceOrder() {
        final var source = """
                fn b() -> int { return 1; }
                fn c() -> int { return 2; }
                fn a() -> int {
                  b();
                  c();
                  return 0;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(99), source, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Valid program should not report diagnostics");
        final var executableA = fileBackend.executableFunctions().stream()
                .filter(fn -> fn.callableName().equals("a"))
                .findFirst()
                .orElseThrow();
        assertEquals("b", executableA.instructions().get(0).calleeCallableName());
        assertEquals("c", executableA.instructions().get(1).calleeCallableName());
    }

    @Test
    void shouldStopLoweringReachablePathAfterReturn() {
        final var source = """
                fn a() -> int { return 1; }
                fn b() -> int { return 2; }
                fn main() -> int {
                  return a();
                  b();
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(100), source, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Valid program should not report diagnostics");
        final var executableMain = fileBackend.executableFunctions().stream()
                .filter(fn -> fn.callableName().equals("main"))
                .findFirst()
                .orElseThrow();
        assertEquals(2, executableMain.instructions().size());
        assertEquals("a", executableMain.instructions().get(0).calleeCallableName());
        assertEquals(p.studio.compiler.models.IRBackendExecutableFunction.InstructionKind.RET, executableMain.instructions().get(1).kind());
    }

    @Test
    void shouldEmitControlFlowInstructionsForIfAndWhile() {
        final var source = """
                fn cond() -> bool { return true; }
                fn t() -> int { return 1; }
                fn f() -> int { return 2; }
                fn loopBody() -> void { return; }
                fn main() -> int {
                  while cond() {
                    break;
                  }
                  if cond() {
                    return t();
                  } else {
                    return f();
                  }
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(101), source, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Valid program should not report diagnostics");
        final var executableMain = fileBackend.executableFunctions().stream()
                .filter(fn -> fn.callableName().equals("main"))
                .findFirst()
                .orElseThrow();
        assertTrue(executableMain.instructions().stream().anyMatch(i ->
                i.kind() == p.studio.compiler.models.IRBackendExecutableFunction.InstructionKind.LABEL));
        assertTrue(executableMain.instructions().stream().anyMatch(i ->
                i.kind() == p.studio.compiler.models.IRBackendExecutableFunction.InstructionKind.JMP));
        assertTrue(executableMain.instructions().stream().anyMatch(i ->
                i.kind() == p.studio.compiler.models.IRBackendExecutableFunction.InstructionKind.JMP_IF_FALSE));
    }

    @Test
    void shouldNotLowerWhenSyntaxErrorsExist() {
        final var source = """
                $
                fn run() -> int { return 1; }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(0), source, diagnostics);

        assertTrue(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(LexErrors.E_LEX_INVALID_CHAR.name())));
        assertEquals(0, fileBackend.functions().size());
    }

    @Test
    void shouldNotLowerWhenStaticSemanticsErrorsExist() {
        final var source = """
                fn run(x: int) -> int { return x; }
                fn run(y: int) -> int { return y; }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(0), source, diagnostics);

        assertTrue(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(PbsSemanticsErrors.E_SEM_DUPLICATE_CALLABLE_SHAPE.name())));
        assertEquals(0, fileBackend.functions().size());
    }

    @Test
    void shouldExtractReservedMetadataForInterfaceModuleSurfaces() {
        final var source = """
                [BuiltinType(name = "Color", version = 1)]
                declare builtin type Color(
                  pub raw: int
                ) {
                }
                declare host Gfx {
                  [Host(module = "gfx", name = "draw_pixel", version = 1)]
                  [Capability(name = "gfx")]
                  fn draw_pixel(x: int, y: int, color: Color) -> void;
                }
                [BuiltinConst(target = "Color", name = "white", version = 1)]
                declare const WHITE: Color;
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(0), source, diagnostics, SourceKind.SDK_INTERFACE);

        assertTrue(diagnostics.isEmpty());
        assertEquals(0, fileBackend.functions().size());
        assertEquals(1, fileBackend.reservedMetadata().hostMethodBindings().size());
        assertEquals(1, fileBackend.reservedMetadata().builtinTypeSurfaces().size());
        assertEquals(1, fileBackend.reservedMetadata().builtinConstSurfaces().size());
        assertEquals(1, fileBackend.reservedMetadata().requiredCapabilities().size());
        assertEquals("gfx", fileBackend.reservedMetadata().requiredCapabilities().getFirst());
        assertEquals("Color", fileBackend.reservedMetadata().builtinTypeSurfaces().getFirst().canonicalTypeName());
        assertEquals(1, fileBackend.reservedMetadata().builtinTypeSurfaces().getFirst().fields().size());
    }

    @Test
    void shouldRejectHostBindingWithoutCapabilityAtHostAdmission() {
        final var source = """
                declare host Gfx {
                  [Host(module = "gfx", name = "draw_pixel", version = 1)]
                  fn draw_pixel(x: int, y: int) -> void;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(7), source, diagnostics, SourceKind.SDK_INTERFACE);

        final var hostDiagnostic = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsHostAdmissionErrors.E_HOST_MISSING_REQUIRED_CAPABILITY.name()))
                .findFirst()
                .orElseThrow();
        assertEquals(DiagnosticPhase.HOST_ADMISSION, hostDiagnostic.getPhase());
        assertEquals(hostDiagnostic.getCode(), hostDiagnostic.getTemplateId());
        assertEquals(0, fileBackend.functions().size());
    }

    @Test
    void shouldRejectHostBindingCapabilityNotDeclaredInStrictContext() {
        final var source = """
                declare host Gfx {
                  [Host(module = "gfx", name = "draw_pixel", version = 1)]
                  [Capability(name = "gfx")]
                  fn draw_pixel(x: int, y: int) -> void;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        compiler.compileFile(
                new FileId(8),
                source,
                diagnostics,
                SourceKind.SDK_INTERFACE,
                HostAdmissionContext.strictWithDeclaredCapabilities(ReadOnlyList.wrap(java.util.List.of("input"))));

        final var hostDiagnostic = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsHostAdmissionErrors.E_HOST_CAPABILITY_NOT_DECLARED.name()))
                .findFirst()
                .orElseThrow();
        assertEquals(DiagnosticPhase.HOST_ADMISSION, hostDiagnostic.getPhase());
    }

    @Test
    void shouldRejectUnknownHostCapabilityAtHostAdmission() {
        final var source = """
                declare host Gfx {
                  [Host(module = "gfx", name = "draw_pixel", version = 1)]
                  [Capability(name = "video")]
                  fn draw_pixel(x: int, y: int) -> void;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        compiler.compileFile(new FileId(9), source, diagnostics, SourceKind.SDK_INTERFACE);

        final var hostDiagnostic = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsHostAdmissionErrors.E_HOST_UNKNOWN_CAPABILITY.name()))
                .findFirst()
                .orElseThrow();
        assertEquals(DiagnosticPhase.HOST_ADMISSION, hostDiagnostic.getPhase());
    }

    @Test
    void shouldRejectMalformedHostCapabilityMetadataAtHostAdmission() {
        final var source = """
                declare host Gfx {
                  [Host(module = "gfx", name = "draw_pixel", version = 1)]
                  [Capability(version = 1)]
                  fn draw_pixel(x: int, y: int) -> void;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        compiler.compileFile(new FileId(10), source, diagnostics, SourceKind.SDK_INTERFACE);

        final var hostDiagnostic = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsHostAdmissionErrors.E_HOST_MALFORMED_CAPABILITY_METADATA.name()))
                .findFirst()
                .orElseThrow();
        assertEquals(DiagnosticPhase.HOST_ADMISSION, hostDiagnostic.getPhase());
    }

    @Test
    void shouldInferBuiltinFieldSlotsFromDeclarationOrder() {
        final var source = """
                [BuiltinType(name = "Color", version = 1)]
                declare builtin type Color(
                  pub r: int,
                  pub g: int
                ) {
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(3), source, diagnostics, SourceKind.SDK_INTERFACE);

        assertTrue(diagnostics.isEmpty());
        final var fields = fileBackend.reservedMetadata().builtinTypeSurfaces().getFirst().fields();
        assertEquals(2, fields.size());
        assertEquals(0L, fields.get(0).slotStart());
        assertEquals(1L, fields.get(1).slotStart());
        assertEquals(0, fileBackend.functions().size());
    }

    @Test
    void shouldInferBuiltinFieldSlotsUsingFlattenedNestedWidth() {
        final var source = """
                [BuiltinType(name = "Color", version = 1)]
                declare builtin type Color(
                  pub raw: int
                ) {
                }

                [BuiltinType(name = "Pixel", version = 1)]
                declare builtin type Pixel(
                  pub x: int,
                  pub y: int,
                  pub color: Color
                ) {
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var compiler = new PbsFrontendCompiler();
        final var fileBackend = compiler.compileFile(new FileId(5), source, diagnostics, SourceKind.SDK_INTERFACE);

        assertTrue(diagnostics.isEmpty());
        final var pixel = fileBackend.reservedMetadata().builtinTypeSurfaces().stream()
                .filter(t -> t.sourceTypeName().equals("Pixel"))
                .findFirst()
                .orElseThrow();
        assertEquals(3, pixel.fields().size());
        assertEquals(0L, pixel.fields().get(0).slotStart());
        assertEquals(1L, pixel.fields().get(1).slotStart());
        assertEquals(2L, pixel.fields().get(2).slotStart());
    }
}
