package p.studio.compiler.pbs.semantics;

import org.junit.jupiter.api.Test;
import p.studio.compiler.pbs.PbsFrontendCompiler;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.FileId;

import static org.junit.jupiter.api.Assertions.*;

class PbsSemanticsFallthroughTest {

    @Test
    void shouldRejectPossibleFallthroughForPlainNonUnitAndResultCallables() {
        final var source = """
                declare error Err { Fail; }

                fn plainIfBad(flag: bool) -> int {
                  if flag { return 1; }
                }

                fn plainLoopBad(flag: bool) -> int {
                  while flag {
                    return 1;
                  }
                }

                fn resultBad(flag: bool) -> result<Err> int {
                  if flag { return ok(1); }
                }
                """;
        final var diagnostics = DiagnosticSink.empty();

        new PbsFrontendCompiler().compileFile(new FileId(0), source, diagnostics);

        final var plainFallthroughCount = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsSemanticsErrors.E_SEM_POSSIBLE_FALLTHROUGH_NON_UNIT.name()))
                .count();
        final var resultFallthroughCount = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsSemanticsErrors.E_SEM_POSSIBLE_FALLTHROUGH_RESULT.name()))
                .count();

        assertEquals(2, plainFallthroughCount);
        assertEquals(1, resultFallthroughCount);
    }

    @Test
    void shouldAllowFallthroughForUnitAndOptionalCallables() {
        final var source = """
                fn unitOk(flag: bool) -> void {
                  if flag { return; }
                }

                fn optionalOk(flag: bool, v: int) -> optional int {
                  if flag { return some(v); }
                }
                """;
        final var diagnostics = DiagnosticSink.empty();

        new PbsFrontendCompiler().compileFile(new FileId(0), source, diagnostics);

        assertFalse(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(PbsSemanticsErrors.E_SEM_POSSIBLE_FALLTHROUGH_NON_UNIT.name())));
        assertFalse(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(PbsSemanticsErrors.E_SEM_POSSIBLE_FALLTHROUGH_RESULT.name())));
    }

    @Test
    void shouldAcceptNonFallthroughWhenAllPathsReturn() {
        final var source = """
                declare enum Mode(Idle, Run);
                declare error Err { Fail; }

                fn plainSwitchWildcardGood(mode: Mode) -> int {
                  switch mode {
                    Mode.Idle: { return 1; },
                    _: { return 2; },
                  };
                }

                fn plainSwitchEnumExhaustiveGood(mode: Mode) -> int {
                  switch mode {
                    Mode.Idle: { return 1; },
                    Mode.Run: { return 2; },
                  };
                }

                fn plainNestedIfGood(flag: bool, mode: Mode) -> int {
                  if flag {
                    switch mode {
                      Mode.Idle: { return 1; },
                      _: { return 2; },
                    };
                  } else {
                    return 3;
                  }
                }

                fn resultGood(flag: bool) -> result<Err> int {
                  if flag { return ok(1); }
                  return err(Err.Fail);
                }
                """;
        final var diagnostics = DiagnosticSink.empty();

        new PbsFrontendCompiler().compileFile(new FileId(0), source, diagnostics);

        assertFalse(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(PbsSemanticsErrors.E_SEM_POSSIBLE_FALLTHROUGH_NON_UNIT.name())));
        assertFalse(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(PbsSemanticsErrors.E_SEM_POSSIBLE_FALLTHROUGH_RESULT.name())));
    }

    @Test
    void shouldTreatTerminalHandleArmsAsCallableCompletion() {
        final var source = """
                declare enum Mode(Idle, Run);
                declare error InErr { A; B; }
                declare error OutErr { Retry; Abort; }

                fn source(v: int) -> result<InErr> int {
                  if v > 0 { return err(InErr.A); }
                  return err(InErr.B);
                }

                fn handleShortRemapTerminal(v: int) -> result<OutErr> int {
                  handle source(v) {
                    InErr.A -> OutErr.Retry,
                    InErr.B -> OutErr.Abort,
                  };
                }

                fn handleBlockErrTerminal(v: int) -> result<OutErr> int {
                  handle source(v) {
                    InErr.A -> { err(OutErr.Retry) },
                    InErr.B -> { err(OutErr.Abort) },
                  };
                }

                fn handleTerminalInIf(flag: bool, v: int) -> result<OutErr> int {
                  if flag {
                    handle source(v) {
                      InErr.A -> OutErr.Retry,
                      InErr.B -> OutErr.Abort,
                    };
                  } else {
                    return err(OutErr.Abort);
                  }
                }

                fn handleTerminalInSwitch(mode: Mode, v: int) -> result<OutErr> int {
                  switch mode {
                    Mode.Idle: {
                      handle source(v) {
                        InErr.A -> { err(OutErr.Retry) },
                        InErr.B -> { err(OutErr.Abort) },
                      };
                    },
                    _: { return err(OutErr.Abort); },
                  };
                }
                """;
        final var diagnostics = DiagnosticSink.empty();

        new PbsFrontendCompiler().compileFile(new FileId(0), source, diagnostics);

        assertFalse(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(PbsSemanticsErrors.E_SEM_POSSIBLE_FALLTHROUGH_RESULT.name())));
    }

    @Test
    void shouldKeepFallthroughWhenHandleArmRecoversWithOk() {
        final var source = """
                declare error InErr { A; B; }
                declare error OutErr { Retry; Abort; }

                fn source(v: int) -> result<InErr> int {
                  if v > 0 { return err(InErr.A); }
                  return err(InErr.B);
                }

                fn handleRecoveringTail(v: int) -> result<OutErr> int {
                  handle source(v) {
                    InErr.A -> { ok(1) },
                    InErr.B -> { err(OutErr.Abort) },
                  };
                }

                fn handleRecoveringInIf(flag: bool, v: int) -> result<OutErr> int {
                  if flag {
                    handle source(v) {
                      InErr.A -> { ok(1) },
                      InErr.B -> { err(OutErr.Abort) },
                    };
                  } else {
                    return err(OutErr.Abort);
                  }
                }
                """;
        final var diagnostics = DiagnosticSink.empty();

        new PbsFrontendCompiler().compileFile(new FileId(0), source, diagnostics);

        final var resultFallthroughCount = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsSemanticsErrors.E_SEM_POSSIBLE_FALLTHROUGH_RESULT.name()))
                .count();
        assertTrue(resultFallthroughCount >= 2);
    }
}
