package p.studio.compiler.pbs.linking;

import org.junit.jupiter.api.Test;
import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.lexer.PbsLexer;
import p.studio.compiler.pbs.parser.PbsBarrelParser;
import p.studio.compiler.pbs.parser.PbsParser;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.FileId;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

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

class PbsModuleVisibilityTest {

    @Test
    void shouldReportMissingBarrelPerModule() {
        final var diagnostics = DiagnosticSink.empty();
        final var nextFileId = new AtomicInteger(0);
        final var module = module("core", "math", List.of(
                """
                fn sum(a: int, b: int) -> int { return a + b; }
                """
        ), (String) null, nextFileId, diagnostics);

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(module)), diagnostics);

        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(PbsLinkErrors.E_LINK_MISSING_BARREL.name())));
    }

    @Test
    void shouldReportBarrelWithoutSourceFiles() {
        final var diagnostics = DiagnosticSink.empty();
        final var nextFileId = new AtomicInteger(0);
        final var module = module("core", "math", List.of(), """
                pub fn run() -> int;
                """, nextFileId, diagnostics);

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(module)), diagnostics);

        assertTrue(diagnostics.stream().anyMatch(d -> d.getCode().equals(PbsLinkErrors.E_LINK_BARREL_WITHOUT_SOURCE.name())));
    }

    @Test
    void shouldReportDuplicateBarrelEntriesByKindAndSignature() {
        final var diagnostics = DiagnosticSink.empty();
        final var nextFileId = new AtomicInteger(0);
        final var module = module("core", "math", List.of(
                """
                fn foo(x: int) -> int { return x; }
                declare struct State(v: int);
                """
        ), """
                pub fn foo(a: int) -> int;
                pub fn foo(b: int) -> int;
                pub struct State;
                pub struct State;
                """, nextFileId, diagnostics);

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(module)), diagnostics);

        final long duplicateCount = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsLinkErrors.E_LINK_DUPLICATE_BARREL_ENTRY.name()))
                .count();
        assertEquals(2, duplicateCount);
    }

    @Test
    void shouldReportDuplicateBarrelFiles() {
        final var diagnostics = DiagnosticSink.empty();
        final var nextFileId = new AtomicInteger(0);
        final var module = module("core", "math", List.of(
                """
                fn run() -> int { return 1; }
                """
        ), List.of(
                """
                pub fn run() -> int;
                """,
                """
                pub fn run() -> int;
                """
        ), nextFileId, diagnostics);

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(module)), diagnostics);

        final long duplicateBarrelCount = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsLinkErrors.E_LINK_DUPLICATE_BARREL_FILE.name()))
                .count();
        assertEquals(1, duplicateBarrelCount);
    }

    @Test
    void shouldReportUnresolvedBarrelEntries() {
        final var diagnostics = DiagnosticSink.empty();
        final var nextFileId = new AtomicInteger(0);
        final var module = module("core", "math", List.of(
                """
                fn foo() -> int { return 1; }
                """
        ), """
                pub fn missing() -> int;
                pub enum MissingEnum;
                """, nextFileId, diagnostics);

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(module)), diagnostics);

        final long unresolvedCount = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsLinkErrors.E_LINK_UNRESOLVED_BARREL_ENTRY.name()))
                .count();
        assertEquals(2, unresolvedCount);
    }

    @Test
    void shouldRejectImportOfNonPublicSymbols() {
        final var diagnostics = DiagnosticSink.empty();
        final var nextFileId = new AtomicInteger(0);

        final var moduleMath = module("core", "math", List.of(
                """
                fn secret() -> int { return 1; }
                fn open() -> int { return 2; }
                """
        ), """
                mod fn secret() -> int;
                pub fn open() -> int;
                """, nextFileId, diagnostics);

        final var moduleApp = module("core", "app", List.of(
                """
                import { secret, open } from @core:math;
                fn use() -> int { return open(); }
                """
        ), """
                pub fn use() -> int;
                """, nextFileId, diagnostics);

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(moduleMath, moduleApp)), diagnostics);

        final var importVisibilityDiagnostics = diagnostics.stream()
                .filter(d -> d.getCode().equals(PbsLinkErrors.E_LINK_IMPORT_SYMBOL_NOT_PUBLIC.name()))
                .toList();
        assertEquals(1, importVisibilityDiagnostics.size());
        assertTrue(importVisibilityDiagnostics.getFirst().getMessage().contains("secret"));
    }

    @Test
    void shouldRejectImportFromMissingModule() {
        final var diagnostics = DiagnosticSink.empty();
        final var nextFileId = new AtomicInteger(0);

        final var moduleApp = module("core", "app", List.of(
                """
                import { open } from @core:missing;
                fn use() -> int { return 1; }
                """
        ), """
                pub fn use() -> int;
                """, nextFileId, diagnostics);

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(moduleApp)), diagnostics);

        assertTrue(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(PbsLinkErrors.E_LINK_IMPORT_MODULE_NOT_FOUND.name())));
    }

    @Test
    void shouldRejectImportOfUnknownSymbol() {
        final var diagnostics = DiagnosticSink.empty();
        final var nextFileId = new AtomicInteger(0);

        final var moduleMath = module("core", "math", List.of(
                """
                fn open() -> int { return 2; }
                """
        ), """
                pub fn open() -> int;
                """, nextFileId, diagnostics);

        final var moduleApp = module("core", "app", List.of(
                """
                import { missing } from @core:math;
                fn use() -> int { return 1; }
                """
        ), """
                pub fn use() -> int;
                """, nextFileId, diagnostics);

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(moduleMath, moduleApp)), diagnostics);

        assertTrue(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(PbsLinkErrors.E_LINK_IMPORT_SYMBOL_UNRESOLVED.name())));
    }

    @Test
    void shouldResolveBarrelFunctionWithResultReturnSurface() {
        final var diagnostics = DiagnosticSink.empty();
        final var nextFileId = new AtomicInteger(0);
        final var module = module("core", "math", List.of(
                """
                fn run(input: int) -> result<Err> (left: int, right: int) { return input; }
                """
        ), """
                pub fn run(v: int) -> result<Err> (a: int, b: int);
                """, nextFileId, diagnostics);

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(module)), diagnostics);

        assertTrue(diagnostics.stream().noneMatch(d ->
                d.getCode().equals(PbsLinkErrors.E_LINK_UNRESOLVED_BARREL_ENTRY.name())));
    }

    @Test
    void shouldResolveHostBarrelEntryWhenHostDeclarationExistsInInterfaceModule() {
        final var diagnostics = DiagnosticSink.empty();
        final var sourceFileId = new FileId(100);
        final var barrelFileId = new FileId(101);
        final var sourceAst = parseSource(
                """
                declare host Gfx {
                  [Host(module = "gfx", name = "draw_pixel", version = 1)]
                  [Capability(name = "gfx")]
                  fn draw_pixel(x: int, y: int) -> void;
                }
                """,
                sourceFileId,
                diagnostics,
                PbsParser.ParseMode.INTERFACE_MODULE);
        final var barrelAst = parseBarrel(
                """
                pub host Gfx;
                """,
                barrelFileId,
                diagnostics);
        final var module = new PbsModuleVisibilityValidator.ModuleUnit(
                new PbsModuleVisibilityValidator.ModuleCoordinates("sdk", ReadOnlyList.wrap(List.of("gfx"))),
                ReadOnlyList.wrap(List.of(new PbsModuleVisibilityValidator.SourceFile(sourceFileId, sourceAst))),
                ReadOnlyList.wrap(List.of(new PbsModuleVisibilityValidator.BarrelFile(barrelFileId, barrelAst))));

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(module)), diagnostics);

        assertTrue(diagnostics.stream().noneMatch(d ->
                d.getCode().equals(PbsLinkErrors.E_LINK_UNRESOLVED_BARREL_ENTRY.name())));
    }

    @Test
    void shouldResolveBuiltinTypeBarrelStructEntryWhenDeclarationExistsInInterfaceModule() {
        final var diagnostics = DiagnosticSink.empty();
        final var sourceFileId = new FileId(105);
        final var barrelFileId = new FileId(106);
        final var sourceAst = parseSource(
                """
                [BuiltinType(name = "Color", version = 1)]
                declare builtin type Color(
                  pub r: int,
                  pub g: int,
                  pub b: int
                ) {
                  [IntrinsicCall(name = "core.color.pack", version = 1)]
                  fn pack() -> int;
                }
                """,
                sourceFileId,
                diagnostics,
                PbsParser.ParseMode.INTERFACE_MODULE);
        final var barrelAst = parseBarrel(
                """
                pub struct Color;
                """,
                barrelFileId,
                diagnostics);
        final var module = new PbsModuleVisibilityValidator.ModuleUnit(
                new PbsModuleVisibilityValidator.ModuleCoordinates("core", ReadOnlyList.wrap(List.of("color"))),
                ReadOnlyList.wrap(List.of(new PbsModuleVisibilityValidator.SourceFile(sourceFileId, sourceAst))),
                ReadOnlyList.wrap(List.of(new PbsModuleVisibilityValidator.BarrelFile(barrelFileId, barrelAst))));

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(module)), diagnostics);

        assertTrue(diagnostics.stream().noneMatch(d ->
                d.getCode().equals(PbsLinkErrors.E_LINK_UNRESOLVED_BARREL_ENTRY.name())));
    }

    @Test
    void shouldRejectHostBarrelEntryWhenHostDeclarationIsMissing() {
        final var diagnostics = DiagnosticSink.empty();
        final var sourceFileId = new FileId(110);
        final var barrelFileId = new FileId(111);
        final var sourceAst = parseSource(
                """
                declare enum Color(Red);
                """,
                sourceFileId,
                diagnostics);
        final var barrelAst = parseBarrel(
                """
                pub host Gfx;
                """,
                barrelFileId,
                diagnostics);
        final var module = new PbsModuleVisibilityValidator.ModuleUnit(
                new PbsModuleVisibilityValidator.ModuleCoordinates("sdk", ReadOnlyList.wrap(List.of("gfx"))),
                ReadOnlyList.wrap(List.of(new PbsModuleVisibilityValidator.SourceFile(sourceFileId, sourceAst))),
                ReadOnlyList.wrap(List.of(new PbsModuleVisibilityValidator.BarrelFile(barrelFileId, barrelAst))));

        new PbsModuleVisibilityValidator().validate(ReadOnlyList.wrap(List.of(module)), diagnostics);

        assertTrue(diagnostics.stream().anyMatch(d ->
                d.getCode().equals(PbsLinkErrors.E_LINK_UNRESOLVED_BARREL_ENTRY.name())));
    }

    private PbsModuleVisibilityValidator.ModuleUnit module(
            final String project,
            final String modulePath,
            final List<String> sourceContents,
            final String barrelContent,
            final AtomicInteger nextFileId,
            final DiagnosticSink diagnostics) {
        return module(project, modulePath, sourceContents, barrelContent == null ? List.of() : List.of(barrelContent), nextFileId, diagnostics);
    }

    private PbsModuleVisibilityValidator.ModuleUnit module(
            final String project,
            final String modulePath,
            final List<String> sourceContents,
            final List<String> barrelContents,
            final AtomicInteger nextFileId,
            final DiagnosticSink diagnostics) {
        final var sources = new ArrayList<PbsModuleVisibilityValidator.SourceFile>();
        for (final var sourceContent : sourceContents) {
            final var fileId = new FileId(nextFileId.getAndIncrement());
            final var ast = parseSource(sourceContent, fileId, diagnostics);
            sources.add(new PbsModuleVisibilityValidator.SourceFile(fileId, ast));
        }

        final var barrels = new ArrayList<PbsModuleVisibilityValidator.BarrelFile>();
        for (final var barrelContent : barrelContents) {
            final var fileId = new FileId(nextFileId.getAndIncrement());
            final var barrelAst = parseBarrel(barrelContent, fileId, diagnostics);
            barrels.add(new PbsModuleVisibilityValidator.BarrelFile(fileId, barrelAst));
        }

        final var coordinates = new PbsModuleVisibilityValidator.ModuleCoordinates(project, modulePathSegments(modulePath));
        return new PbsModuleVisibilityValidator.ModuleUnit(
                coordinates,
                ReadOnlyList.wrap(sources),
                ReadOnlyList.wrap(barrels));
    }

    private PbsAst.File parseSource(
            final String source,
            final FileId fileId,
            final DiagnosticSink diagnostics) {
        return parseSource(source, fileId, diagnostics, PbsParser.ParseMode.ORDINARY);
    }

    private PbsAst.File parseSource(
            final String source,
            final FileId fileId,
            final DiagnosticSink diagnostics,
            final PbsParser.ParseMode parseMode) {
        return PbsParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics, parseMode);
    }

    private PbsAst.BarrelFile parseBarrel(
            final String source,
            final FileId fileId,
            final DiagnosticSink diagnostics) {
        return PbsBarrelParser.parse(PbsLexer.lex(source, fileId, diagnostics), fileId, diagnostics);
    }

    private ReadOnlyList<String> modulePathSegments(final String modulePath) {
        if (modulePath == null || modulePath.isBlank()) {
            return ReadOnlyList.empty();
        }
        return ReadOnlyList.wrap(List.of(modulePath.split("/")));
    }
}
