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 java.nio.charset.StandardCharsets;

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

class PbsParserTest {

    @Test
    void shouldParseSingleFunctionWithArrowReturn() {
        final var source = """
            fn sum(a: int, b: int) -> int
            {
              return a + b;
            }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);
        final var tokens = PbsLexer.lex(source, fileId, diagnostics);

        final PbsAst.File ast = PbsParser.parse(tokens, fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Parser should not report diagnostics for valid function");
        assertEquals(0, ast.imports().size());
        assertEquals(1, ast.topDecls().size());
        assertEquals(1, ast.functions().size());

        final var fn = ast.functions().getFirst();
        assertEquals("sum", fn.name());
        assertEquals(2, fn.parameters().size());
        assertEquals(PbsAst.ReturnKind.PLAIN, fn.returnKind());
        assertEquals("int", fn.returnType().name());
        assertNull(fn.resultErrorType());
        assertEquals(1, fn.body().statements().size());
        assertInstanceOf(PbsAst.ReturnStatement.class, fn.body().statements().getFirst());

        assertEquals(fileId, fn.span().getFileId());
        assertTrue(fn.span().getEnd() > fn.span().getStart());
    }

    @Test
    void shouldParseImportsAndTopDeclOrdering() {
        final var source = """
            import { Vec2 as V2 } from @core:math;

            fn first() -> int { return 1; }
            declare struct Point(x: int, y: int);
            declare callback Tick(dt: int) -> int;
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);
        final var tokens = PbsLexer.lex(source, fileId, diagnostics);

        final PbsAst.File ast = PbsParser.parse(tokens, fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Parser should accept valid declarations in this slice");
        assertEquals(1, ast.imports().size());
        assertEquals(3, ast.topDecls().size());
        assertInstanceOf(PbsAst.FunctionDecl.class, ast.topDecls().get(0));
        assertInstanceOf(PbsAst.StructDecl.class, ast.topDecls().get(1));
        assertInstanceOf(PbsAst.CallbackDecl.class, ast.topDecls().get(2));

        final var importDecl = ast.imports().getFirst();
        assertEquals("core", importDecl.moduleRef().project());
        assertEquals("math", importDecl.moduleRef().pathSegments().getFirst());
    }

    @Test
    void shouldParseDeclarationFamiliesWithShape() {
        final var source = """
            fn f() -> int { return 1; }
            [Init]
            fn init() -> void { return; }
            [Frame]
            fn frame() -> void { return; }
            declare struct S(pub mut x: int, y: optional int) { fn run() -> void { return; } ctor make(x: int) { return; } }
            declare contract C { fn run(x: int) -> int; }
            declare service Game { fn tick(x: int) -> int { return x; } }
            declare error Err { Fail; Crash; }
            declare enum Mode(Idle = 0, Run = 1);
            declare callback TickCb(x: int) -> result<Err> int;
            declare global STATE: int = 10;
            declare const LIMIT: int = 10;
            implements C for S using s { fn run(x: int) -> int { return x; } }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);
        final var tokens = PbsLexer.lex(source, fileId, diagnostics);

        final PbsAst.File ast = PbsParser.parse(tokens, fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Parser should represent declaration families with expected shapes");
        assertEquals(12, ast.topDecls().size());

        final var initDecl = assertInstanceOf(PbsAst.FunctionDecl.class, ast.topDecls().get(1));
        assertEquals(PbsAst.LifecycleMarker.INIT, initDecl.lifecycleMarker());

        final var frameDecl = assertInstanceOf(PbsAst.FunctionDecl.class, ast.topDecls().get(2));
        assertEquals(PbsAst.LifecycleMarker.FRAME, frameDecl.lifecycleMarker());

        final var structDecl = assertInstanceOf(PbsAst.StructDecl.class, ast.topDecls().get(3));
        assertEquals(2, structDecl.fields().size());
        assertTrue(structDecl.hasBody());
        assertEquals(1, structDecl.methods().size());
        assertEquals(1, structDecl.ctors().size());

        final var contractDecl = assertInstanceOf(PbsAst.ContractDecl.class, ast.topDecls().get(4));
        assertEquals(1, contractDecl.signatures().size());

        final var serviceDecl = assertInstanceOf(PbsAst.ServiceDecl.class, ast.topDecls().get(5));
        assertEquals(1, serviceDecl.methods().size());

        final var errorDecl = assertInstanceOf(PbsAst.ErrorDecl.class, ast.topDecls().get(6));
        assertEquals(2, errorDecl.cases().size());

        final var enumDecl = assertInstanceOf(PbsAst.EnumDecl.class, ast.topDecls().get(7));
        assertEquals(2, enumDecl.cases().size());

        final var callbackDecl = assertInstanceOf(PbsAst.CallbackDecl.class, ast.topDecls().get(8));
        assertEquals(PbsAst.ReturnKind.RESULT, callbackDecl.returnKind());
        assertEquals("Err", callbackDecl.resultErrorType().name());

        final var globalDecl = assertInstanceOf(PbsAst.GlobalDecl.class, ast.topDecls().get(9));
        assertEquals("STATE", globalDecl.name());
        assertInstanceOf(PbsAst.ConstDecl.class, ast.topDecls().get(10));
        final var implementsDecl = assertInstanceOf(PbsAst.ImplementsDecl.class, ast.topDecls().get(11));
        assertEquals(1, implementsDecl.methods().size());
    }

    @Test
    void shouldRejectAttributesInOrdinarySourceAndRecover() {
        final var source = """
            [Host(module = "gfx", name = "draw", version = 1)]
            [Capability(name = "gfx")]
            fn run() -> int { return 1; }
            declare struct Point(x: int,);
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        final PbsAst.File ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertEquals(2, ast.topDecls().size());
        assertInstanceOf(PbsAst.FunctionDecl.class, ast.topDecls().get(0));
        assertInstanceOf(PbsAst.StructDecl.class, ast.topDecls().get(1));
        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(ParseErrors.E_PARSE_ATTRIBUTES_NOT_ALLOWED.name())));
    }

    @Test
    void shouldParseLifecycleMarkersOnlyOnTopLevelFunctions() {
        final var source = """
            [Init]
            fn init() -> void { return; }
            [Frame]
            fn frame() -> void { return; }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        final PbsAst.File ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Lifecycle markers should be accepted on top-level functions");
        final var initDecl = assertInstanceOf(PbsAst.FunctionDecl.class, ast.topDecls().get(0));
        final var frameDecl = assertInstanceOf(PbsAst.FunctionDecl.class, ast.topDecls().get(1));
        assertEquals(PbsAst.LifecycleMarker.INIT, initDecl.lifecycleMarker());
        assertEquals(PbsAst.LifecycleMarker.FRAME, frameDecl.lifecycleMarker());
    }

    @Test
    void shouldRejectErrorKeywordAsCallableAndMemberName() {
        final var source = """
            declare service Log {
              fn error(msg: str) -> void { return; }
              fn error(tag: int, msg: str) -> void { return; }
            }

            fn run() -> void {
              Log.failure("oops");
              Log.failure(7, "oops");
              return;
            }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertFalse(diagnostics.isEmpty(), "Parser should reject reserved keyword 'error' as callable/member name");
    }

    @Test
    void shouldTrackLetNameSpanSeparatelyFromWholeStatementSpan() {
        final var source = """
            fn main() -> void {
              let inferred = 1;
              return;
            }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        final PbsAst.File ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), diagnostics.stream().map(d -> d.getCode() + ":" + d.getMessage()).toList().toString());
        final var functionDecl = assertInstanceOf(PbsAst.FunctionDecl.class, ast.topDecls().getFirst());
        final var letStatement = assertInstanceOf(PbsAst.LetStatement.class, functionDecl.body().statements().getFirst());
        assertEquals("inferred", letStatement.name());
        assertEquals(source.indexOf("inferred"), letStatement.nameSpan().getStart());
        assertEquals(source.indexOf("inferred") + "inferred".length(), letStatement.nameSpan().getEnd());
        assertTrue(letStatement.span().getStart() < letStatement.nameSpan().getStart());
        assertTrue(letStatement.span().getEnd() > letStatement.nameSpan().getEnd());
    }

    @Test
    void shouldAcceptStructFieldTrailingComma() {
        final var source = "declare struct S(a: int,);";
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        final PbsAst.File ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Struct field trailing comma should be accepted");
        final var decl = assertInstanceOf(PbsAst.StructDecl.class, ast.topDecls().getFirst());
        assertEquals(1, decl.fields().size());
        assertEquals("a", decl.fields().getFirst().name());
    }

    @Test
    void shouldParseOptionalAndNamedTupleTypes() {
        final var source = """
            fn shape(a: optional int, b: (x: int, y: optional float)) -> (value: optional int, meta: (k: int, v: float)) {
              return a;
            }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        final PbsAst.File ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Parser should accept optional and named tuple type surfaces");

        final var fn = ast.functions().getFirst();
        final var p0 = fn.parameters().get(0).typeRef();
        assertEquals(PbsAst.TypeRefKind.OPTIONAL, p0.kind());
        assertEquals(PbsAst.TypeRefKind.SIMPLE, p0.inner().kind());

        final var p1 = fn.parameters().get(1).typeRef();
        assertEquals(PbsAst.TypeRefKind.NAMED_TUPLE, p1.kind());
        assertEquals(2, p1.fields().size());

        final var returnType = fn.returnType();
        assertEquals(PbsAst.TypeRefKind.NAMED_TUPLE, returnType.kind());
    }

    @Test
    void shouldRejectNamedTupleTypeWithArityAboveSix() {
        final var source = """
            fn tooWide(v: (a: int, b: int, c: int, d: int, e: int, f: int, g: int)) -> int {
              return 0;
            }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(ParseErrors.E_PARSE_INVALID_TYPE_SURFACE.name())));
    }

    @Test
    void shouldRejectReservedDeclareHostAndBuiltinType() {
        final var source = """
            declare host HostApi { fn run() -> int; }
            declare builtin type Vec(x: int) { fn mag() -> float; }
            fn ok() -> int { return 1; }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        final PbsAst.File ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.hasErrors(), "Parser should reject reserved declare forms in ordinary source");
        assertEquals(3, ast.topDecls().size());
        assertInstanceOf(PbsAst.InvalidDecl.class, ast.topDecls().get(0));
        assertInstanceOf(PbsAst.InvalidDecl.class, ast.topDecls().get(1));
        assertInstanceOf(PbsAst.FunctionDecl.class, ast.topDecls().get(2));
    }

    @Test
    void shouldParseReservedDeclarationsInInterfaceModuleMode() {
        final var source = """
            [BuiltinType(name = "Color", version = 1)]
            declare builtin type Color(
              pub r: int,
              pub g: int,
              pub b: int,
            ) {
              [IntrinsicCall(name = "color.pack", version = 1)]
              fn pack() -> int;
            }
            declare host Gfx {
              [Host(module = "gfx", name = "draw_pixel", version = 1)]
              [Capability(name = "gfx")]
              fn draw(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 fileId = new FileId(0);

        final PbsAst.File ast = PbsParser.parse(
                PbsLexer.lex(source, fileId, diagnostics),
                fileId,
                diagnostics,
                PbsParser.ParseMode.INTERFACE_MODULE);

        assertTrue(diagnostics.isEmpty(), "Parser should accept reserved interface declarations in interface-module mode");
        assertEquals(3, ast.topDecls().size());

        final var builtinType = assertInstanceOf(PbsAst.BuiltinTypeDecl.class, ast.topDecls().get(0));
        assertEquals(1, builtinType.attributes().size());
        assertEquals(3, builtinType.fields().size());
        assertEquals(1, builtinType.signatures().size());
        assertEquals(0, builtinType.fields().getFirst().attributes().size());
        assertEquals(1, builtinType.signatures().getFirst().attributes().size());

        final var hostDecl = assertInstanceOf(PbsAst.HostDecl.class, ast.topDecls().get(1));
        assertEquals(1, hostDecl.signatures().size());
        assertEquals(2, hostDecl.signatures().getFirst().attributes().size());

        final var constDecl = assertInstanceOf(PbsAst.ConstDecl.class, ast.topDecls().get(2));
        assertEquals(1, constDecl.attributes().size());
    }

    @Test
    void shouldRecoverFromMalformedReservedDeclarationsInInterfaceModuleMode() {
        final var source = """
            declare host Gfx {
              [Host(module = "gfx", name = "draw_pixel", version = 1)]
              [Capability(name = "gfx")]
              fn draw(x: int, y: int) -> void;
              unexpected_token
            }
            fn ok() -> int { return 1; }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        final PbsAst.File ast = PbsParser.parse(
                PbsLexer.lex(source, fileId, diagnostics),
                fileId,
                diagnostics,
                PbsParser.ParseMode.INTERFACE_MODULE);

        assertTrue(diagnostics.hasErrors(), "Malformed reserved forms must still emit deterministic parse diagnostics");
        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(ParseErrors.E_PARSE_INVALID_DECL_SHAPE.name())));
        assertEquals(2, ast.topDecls().size());
        assertInstanceOf(PbsAst.HostDecl.class, ast.topDecls().get(0));
        assertInstanceOf(PbsAst.FunctionDecl.class, ast.topDecls().get(1));
    }

    @Test
    void shouldRejectFunctionElseFallbackSurface() {
        final var source = """
            fn run() -> int else 1 {
              return 1;
            }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.hasErrors(), "Function-level else fallback syntax is not valid PBS core syntax");
    }

    @Test
    void shouldRejectSlotAttributeInInterfaceModuleMode() {
        final var source = """
            [BuiltinType(name = "Color", version = 1)]
            declare builtin type Color(
              [Slot(index = 0)] pub raw: int
            ) {
            }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        PbsParser.parse(
                PbsLexer.lex(source, fileId, diagnostics),
                fileId,
                diagnostics,
                PbsParser.ParseMode.INTERFACE_MODULE);

        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(ParseErrors.E_PARSE_INVALID_DECL_SHAPE.name())));
    }

    @Test
    void shouldReportEnumMixedAndDuplicateCases() {
        final var source = """
            declare enum Mixed(A, B = 1);
            declare enum Duplicated(Idle = 0, Idle = 1, Run = 1);
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(ParseErrors.E_PARSE_INVALID_ENUM_FORM.name())));
        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(ParseErrors.E_PARSE_DUPLICATE_ENUM_CASE_LABEL.name())));
        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(ParseErrors.E_PARSE_DUPLICATE_ENUM_CASE_ID.name())));
    }

    @Test
    void shouldRecoverWithInvalidDeclarationNode() {
        final var source = """
            declare ;
            fn ok() -> int { return 1; }
            """;
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);
        final var tokens = PbsLexer.lex(source, fileId, diagnostics);

        final PbsAst.File ast = PbsParser.parse(tokens, fileId, diagnostics);

        assertTrue(diagnostics.hasErrors(), "Parser should report malformed declaration");
        assertEquals(2, ast.topDecls().size());
        assertInstanceOf(PbsAst.InvalidDecl.class, ast.topDecls().get(0));
        assertInstanceOf(PbsAst.FunctionDecl.class, ast.topDecls().get(1));
    }

    @Test
    void shouldPropagateUtf8ByteOffsetsToAstSpans() {
        final var source = "fn cafe\u00E9() -> int { return 1; }";
        final var diagnostics = DiagnosticSink.empty();
        final var fileId = new FileId(0);

        final PbsAst.File ast = PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);

        assertTrue(diagnostics.isEmpty(), "Valid non-ASCII source should parse without diagnostics");

        final int expectedEnd = source.getBytes(StandardCharsets.UTF_8).length;
        final var fn = ast.functions().getFirst();
        assertEquals(expectedEnd, ast.span().getEnd());
        assertEquals(expectedEnd, fn.span().getEnd());
    }
}
