package p.studio.compiler.pbs.parser;

import org.junit.jupiter.api.Test;
import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.lexer.PbsLexer;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.FileId;

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

class PbsExprParserTest {

    @Test
    void shouldParseApplyElseAndPostfixChains() {
        final var source = """
                fn exprs(opt: int, fallback: int, x: int) -> int {
                  f apply g apply x;
                  opt else fallback else x;
                  obj.member().next!;
                  return x;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);
        final var ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Valid expression forms should parse cleanly");

        final var body = ast.functions().getFirst().body().statements();
        final var applyExpr = assertInstanceOf(PbsAst.ApplyExpr.class,
                assertInstanceOf(PbsAst.ExpressionStatement.class, body.get(0)).expression());
        assertInstanceOf(PbsAst.ApplyExpr.class, applyExpr.argument());

        final var elseExpr = assertInstanceOf(PbsAst.ElseExpr.class,
                assertInstanceOf(PbsAst.ExpressionStatement.class, body.get(1)).expression());
        assertInstanceOf(PbsAst.ElseExpr.class, elseExpr.fallbackExpression());

        assertInstanceOf(PbsAst.PropagateExpr.class,
                assertInstanceOf(PbsAst.ExpressionStatement.class, body.get(2)).expression());
    }

    @Test
    void shouldParseIfSwitchAndHandleExpressionsWithTailBlocks() {
        final var source = """
                fn flow(cond: bool, state: int) -> int {
                  let a: int = if cond { 1 } else { 2 };
                  let b: int = switch state { default: { 0 }, _: { 1 } };
                  let c: int = handle run apply () { Err.fail -> Other.fail, _ -> { 1 } };
                  return 0;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);
        final var ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Expression control-flow forms should parse cleanly");

        final var body = ast.functions().getFirst().body().statements();
        final var ifExpr = assertInstanceOf(PbsAst.IfExpr.class,
                assertInstanceOf(PbsAst.LetStatement.class, body.get(0)).initializer());
        assertInstanceOf(PbsAst.IntLiteralExpr.class, ifExpr.thenBlock().tailExpression());
        final var elseBlock = assertInstanceOf(PbsAst.BlockExpr.class, ifExpr.elseExpression()).block();
        assertInstanceOf(PbsAst.IntLiteralExpr.class, elseBlock.tailExpression());

        final var switchExpr = assertInstanceOf(PbsAst.SwitchExpr.class,
                assertInstanceOf(PbsAst.LetStatement.class, body.get(1)).initializer());
        assertEquals(2, switchExpr.arms().size());
        final var defaultPattern = assertInstanceOf(PbsAst.WildcardSwitchPattern.class,
                switchExpr.arms().getFirst().pattern());
        assertEquals(PbsAst.WildcardSwitchKind.DEFAULT, defaultPattern.kind());
        assertInstanceOf(PbsAst.IntLiteralExpr.class, switchExpr.arms().getFirst().block().tailExpression());
        final var underscorePattern = assertInstanceOf(PbsAst.WildcardSwitchPattern.class,
                switchExpr.arms().get(1).pattern());
        assertEquals(PbsAst.WildcardSwitchKind.UNDERSCORE, underscorePattern.kind());
        assertInstanceOf(PbsAst.IntLiteralExpr.class, switchExpr.arms().get(1).block().tailExpression());

        final var handle = assertInstanceOf(PbsAst.HandleExpr.class,
                assertInstanceOf(PbsAst.LetStatement.class, body.get(2)).initializer());
        assertEquals(2, handle.arms().size());
        assertInstanceOf(PbsAst.IntLiteralExpr.class, handle.arms().get(1).block().tailExpression());
    }

    @Test
    void shouldParseIfSwitchAndHandleExpressionsWithReturnBlocks() {
        final var source = """
                fn flow(cond: bool, state: int) -> int {
                  let a: int = if cond { return 1; } else { return 2; };
                  let b: int = switch state { default: { return 0; } };
                  let c: int = handle run apply () { Err.fail -> Other.fail, _ -> { return 1; } };
                  return 0;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);
        final var ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Expression control-flow forms with return blocks should parse cleanly");

        final var body = ast.functions().getFirst().body().statements();
        final var ifExpr = assertInstanceOf(PbsAst.IfExpr.class,
                assertInstanceOf(PbsAst.LetStatement.class, body.get(0)).initializer());
        assertEquals(1, ifExpr.thenBlock().statements().size());
        assertInstanceOf(PbsAst.ReturnStatement.class, ifExpr.thenBlock().statements().getFirst());
        assertNull(ifExpr.thenBlock().tailExpression());
        final var elseBlock = assertInstanceOf(PbsAst.BlockExpr.class, ifExpr.elseExpression()).block();
        assertEquals(1, elseBlock.statements().size());
        assertInstanceOf(PbsAst.ReturnStatement.class, elseBlock.statements().getFirst());
        assertNull(elseBlock.tailExpression());

        final var switchExpr = assertInstanceOf(PbsAst.SwitchExpr.class,
                assertInstanceOf(PbsAst.LetStatement.class, body.get(1)).initializer());
        assertEquals(1, switchExpr.arms().getFirst().block().statements().size());
        assertInstanceOf(PbsAst.ReturnStatement.class, switchExpr.arms().getFirst().block().statements().getFirst());
        assertNull(switchExpr.arms().getFirst().block().tailExpression());

        final var handle = assertInstanceOf(PbsAst.HandleExpr.class,
                assertInstanceOf(PbsAst.LetStatement.class, body.get(2)).initializer());
        assertEquals(2, handle.arms().size());
        assertEquals(1, handle.arms().get(1).block().statements().size());
        assertInstanceOf(PbsAst.ReturnStatement.class, handle.arms().get(1).block().statements().getFirst());
        assertNull(handle.arms().get(1).block().tailExpression());
    }

    @Test
    void shouldParseConstructionAndResultSurfaces() {
        final var source = """
                fn build(ctx: int, v: int) -> int {
                  some(v);
                  none;
                  ok(v);
                  err(Game.Fail);
                  new Vec(v);
                  new Vec.zero(v);
                  bind(ctx, handler);
                  return v;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);
        final var ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Constructor and optional/result expression surfaces should parse cleanly");

        final var body = ast.functions().getFirst().body().statements();
        assertInstanceOf(PbsAst.SomeExpr.class,
                assertInstanceOf(PbsAst.ExpressionStatement.class, body.get(0)).expression());
        assertInstanceOf(PbsAst.NoneExpr.class,
                assertInstanceOf(PbsAst.ExpressionStatement.class, body.get(1)).expression());
        assertInstanceOf(PbsAst.OkExpr.class,
                assertInstanceOf(PbsAst.ExpressionStatement.class, body.get(2)).expression());
        assertInstanceOf(PbsAst.ErrExpr.class,
                assertInstanceOf(PbsAst.ExpressionStatement.class, body.get(3)).expression());
        assertInstanceOf(PbsAst.NewExpr.class,
                assertInstanceOf(PbsAst.ExpressionStatement.class, body.get(4)).expression());
        assertInstanceOf(PbsAst.NewExpr.class,
                assertInstanceOf(PbsAst.ExpressionStatement.class, body.get(5)).expression());
        assertInstanceOf(PbsAst.BindExpr.class,
                assertInstanceOf(PbsAst.ExpressionStatement.class, body.get(6)).expression());
    }

    @Test
    void shouldRejectInvalidExpressionFormsDeterministically() {
        final var source = """
                fn bad(a: int, b: int, c: int) -> int {
                  a < b < c;
                  a == b == c;
                  value?;
                  (a: 1);
                  (a: 1, 2);
                  (1, 2, 3, 4, 5, 6, 7);
                  if a { 1; };
                  return 0;
                }
                """;

        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);
        PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.hasErrors(), "Invalid expression forms should produce diagnostics");
        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(ParseErrors.E_PARSE_NON_ASSOC.name())));
        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(ParseErrors.E_PARSE_INVALID_PROPAGATE_OPERATOR.name())));
        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(ParseErrors.E_PARSE_INVALID_TUPLE_LITERAL.name())));
    }
}
