"""Unit tests for databridge_core.linker — cross-file entity linking.

No external dependencies or API keys required. All tests use synthetic data.
"""
from __future__ import annotations

import json
import os
import tempfile
from pathlib import Path
from typing import Any, Dict, List

import pytest

# ── Types ────────────────────────────────────────────────────────────────────

from databridge_core.linker._types import (
    ConflictType,
    EntityCluster,
    EntityConflict,
    EntityLink,
    EntityMap,
    EntityMention,
    LinkBatchSummary,
    MatchStrategy,
    _new_id,
)


class TestTypes:
    """Test Pydantic model construction and serialization."""

    def test_new_id_prefix(self):
        id_ = _new_id("link")
        assert id_.startswith("link_")
        assert len(id_) == 5 + 12  # prefix + underscore + 12 hex chars

    def test_new_id_unique(self):
        ids = {_new_id("test") for _ in range(100)}
        assert len(ids) == 100

    def test_entity_mention_defaults(self):
        m = EntityMention()
        assert m.file_path == ""
        assert m.raw_name == ""
        assert m.normalized_name == ""
        assert m.mention_id.startswith("mention_")

    def test_entity_mention_full(self):
        m = EntityMention(
            file_path="data/test.xlsx",
            archetype="Financial Report",
            raw_name="Revenue_West",
            normalized_name="revenue_west",
            cell_ref="Sheet1!B12",
            formula="=SUM(B2:B11)",
            business_meaning="Total Western Region Revenue",
            source_type="formula_intent",
        )
        d = m.model_dump()
        assert d["raw_name"] == "Revenue_West"
        assert d["source_type"] == "formula_intent"

    def test_entity_link(self):
        link = EntityLink(
            mention_a_id="mention_aaa",
            mention_b_id="mention_bbb",
            composite_score=0.85,
            component_scores={"name_similarity": 0.9},
            strategies_used=["name_similarity"],
        )
        assert link.composite_score == 0.85

    def test_entity_conflict(self):
        c = EntityConflict(
            cluster_id="cluster_test",
            conflict_type=ConflictType.FORMULA_MISMATCH,
            description="Different formulas",
            mention_ids=["m1", "m2"],
            severity="high",
        )
        assert c.conflict_type == ConflictType.FORMULA_MISMATCH

    def test_entity_cluster(self):
        cluster = EntityCluster(
            canonical_name="Revenue",
            domain="revenue",
            file_count=3,
        )
        assert cluster.canonical_name == "Revenue"
        assert cluster.cluster_id.startswith("cluster_")

    def test_entity_map(self):
        em = EntityMap(
            total_mentions=10,
            total_clusters=3,
        )
        assert em.total_mentions == 10
        assert em.map_id.startswith("map_")
        assert em.timestamp  # auto-generated

    def test_link_batch_summary(self):
        s = LinkBatchSummary(
            total_mentions=100,
            total_clusters=20,
            domain_counts={"revenue": 5, "expense": 3},
        )
        assert s.avg_cluster_size == 0.0  # default

    def test_entity_mention_partition_fields(self):
        m = EntityMention(
            file_path="data/test.xlsx",
            raw_name="Revenue",
            normalized_name="revenue",
            partition="Sheet1",
            partition_type="data",
            workbook="test.xlsx",
        )
        d = m.model_dump()
        assert d["partition"] == "Sheet1"
        assert d["partition_type"] == "data"
        assert d["workbook"] == "test.xlsx"

    def test_entity_mention_partition_defaults(self):
        m = EntityMention()
        assert m.partition == ""
        assert m.partition_type == ""
        assert m.workbook == ""

    def test_match_strategy_enum(self):
        assert MatchStrategy.NAME_SIMILARITY.value == "name_similarity"
        assert MatchStrategy.FORMULA_SIMILARITY.value == "formula_similarity"
        assert MatchStrategy.PARTITION_COMPATIBILITY.value == "partition_compatibility"

    def test_conflict_type_enum(self):
        assert ConflictType.FORMULA_MISMATCH.value == "formula_mismatch"
        assert ConflictType.SIGN_REVERSAL.value == "sign_reversal"


# ── Extractor ────────────────────────────────────────────────────────────────

from databridge_core.linker._extractor import (
    _extract_names_from_text,
    _normalize_name,
    _parse_partition_from_cell_ref,
    extract_mentions_from_logic_dna,
    load_and_extract_mentions,
)


class TestExtractor:
    """Test entity mention extraction from Logic DNA data."""

    def test_normalize_name_basic(self):
        assert _normalize_name("Revenue_West") == "revenue_west"

    def test_normalize_name_spaces(self):
        assert _normalize_name("Total Q1 Revenue") == "total_q1_revenue"

    def test_normalize_name_special_chars(self):
        assert _normalize_name("Acct#101-Rev") == "acct_101_rev"

    def test_normalize_name_empty(self):
        assert _normalize_name("") == ""

    def test_extract_names_from_text_quoted(self):
        text = 'The "Revenue West" account feeds into "Total Revenue"'
        names = _extract_names_from_text(text)
        assert "Revenue West" in names
        assert "Total Revenue" in names

    def test_extract_names_from_text_underscored(self):
        text = "The Revenue_West column maps to Acct_101"
        names = _extract_names_from_text(text)
        assert "Revenue_West" in names
        assert "Acct_101" in names

    def test_extract_names_from_text_camelcase(self):
        text = "RevenueWest is linked to TotalMargin"
        names = _extract_names_from_text(text)
        assert "RevenueWest" in names
        assert "TotalMargin" in names

    def test_extract_names_from_text_accounting_patterns(self):
        text = "Total Revenue for the period"
        names = _extract_names_from_text(text)
        assert any("Total Revenue" in n for n in names)

    def test_extract_mentions_from_logic_dna_basic(self):
        dna = {
            "file_path": "data/test.xlsx",
            "archetype": "Financial Report",
            "formula_intents": [
                {
                    "cell_ref": "Sheet1!B12",
                    "raw_formula": "=SUM(B2:B11)",
                    "business_meaning": 'Total "Revenue_West" for Q1',
                },
            ],
            "cross_references": [],
        }
        mentions = extract_mentions_from_logic_dna(dna)
        assert len(mentions) >= 1
        assert all(m.file_path == "data/test.xlsx" for m in mentions)
        assert all(m.archetype == "Financial Report" for m in mentions)

    def test_extract_mentions_max_cap(self):
        dna = {
            "file_path": "data/test.xlsx",
            "archetype": "Data Extract",
            "formula_intents": [
                {
                    "cell_ref": f"Sheet1!B{i}",
                    "raw_formula": f"=A{i}+B{i}",
                    "business_meaning": f'"Entity_{i}" total for region',
                }
                for i in range(100)
            ],
            "cross_references": [],
        }
        mentions = extract_mentions_from_logic_dna(dna, max_mentions=10)
        assert len(mentions) <= 10

    def test_extract_mentions_from_cross_references(self):
        dna = {
            "file_path": "data/test.xlsx",
            "archetype": "Financial Report",
            "formula_intents": [],
            "cross_references": [
                {
                    "source_cell": "Sheet1!A1",
                    "target_cell": "Sheet2!B1",
                    "relationship": '"Revenue_Total" feeds into "Gross_Margin"',
                },
            ],
        }
        mentions = extract_mentions_from_logic_dna(dna)
        assert len(mentions) >= 1
        assert any(m.source_type == "cross_reference" for m in mentions)

    def test_extract_mentions_empty_dna(self):
        dna = {
            "file_path": "data/empty.xlsx",
            "archetype": "Unknown",
            "formula_intents": [],
            "cross_references": [],
        }
        mentions = extract_mentions_from_logic_dna(dna)
        assert mentions == []

    def test_load_and_extract_mentions_from_dir(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            # Write a Logic DNA file
            dna = {
                "file_path": "data/test.xlsx",
                "archetype": "Financial Report",
                "formula_intents": [
                    {
                        "cell_ref": "Sheet1!B12",
                        "raw_formula": "=SUM(B2:B11)",
                        "business_meaning": '"Revenue_West" total',
                    },
                ],
                "cross_references": [],
            }
            dna_file = Path(tmpdir) / "test_logic_dna.json"
            with open(dna_file, "w") as f:
                json.dump(dna, f)

            mentions = load_and_extract_mentions(logic_dna_dir=tmpdir)
            assert len(mentions) >= 1

    def test_load_and_extract_missing_dir(self):
        mentions = load_and_extract_mentions(logic_dna_dir="/nonexistent/dir")
        assert mentions == []

    def test_parse_partition_from_cell_ref_with_sheet(self):
        assert _parse_partition_from_cell_ref("Sheet1!B12") == "Sheet1"

    def test_parse_partition_from_cell_ref_summary(self):
        assert _parse_partition_from_cell_ref("Summary!A1") == "Summary"

    def test_parse_partition_from_cell_ref_no_sheet(self):
        assert _parse_partition_from_cell_ref("A1") == ""

    def test_parse_partition_from_cell_ref_empty(self):
        assert _parse_partition_from_cell_ref("") == ""

    def test_parse_partition_from_cell_ref_quoted_sheet(self):
        assert _parse_partition_from_cell_ref("'My Sheet'!A1") == "My Sheet"

    def test_extract_mentions_populates_partition(self):
        dna = {
            "file_path": "data/test.xlsx",
            "archetype": "Financial Report",
            "formula_intents": [
                {
                    "cell_ref": "Revenue!B12",
                    "raw_formula": "=SUM(B2:B11)",
                    "business_meaning": '"Revenue_West" total',
                },
            ],
            "cross_references": [],
        }
        mentions = extract_mentions_from_logic_dna(dna)
        assert len(mentions) >= 1
        assert mentions[0].partition == "Revenue"

    def test_extract_mentions_populates_workbook(self):
        dna = {
            "file_path": "data/test.xlsx",
            "archetype": "Financial Report",
            "workbook": "test.xlsx",
            "formula_intents": [
                {
                    "cell_ref": "Sheet1!B1",
                    "raw_formula": "=A1",
                    "business_meaning": '"Revenue" total',
                },
            ],
            "cross_references": [],
        }
        mentions = extract_mentions_from_logic_dna(dna)
        assert len(mentions) >= 1
        assert mentions[0].workbook == "test.xlsx"

    def test_extract_mentions_partition_from_dna_sheet_name(self):
        dna = {
            "file_path": "data/test.xlsx",
            "archetype": "Financial Report",
            "sheet_name": "DataSheet",
            "formula_intents": [
                {
                    "cell_ref": "A1",
                    "raw_formula": "=B1",
                    "business_meaning": '"Revenue" total',
                },
            ],
            "cross_references": [],
        }
        mentions = extract_mentions_from_logic_dna(dna)
        assert len(mentions) >= 1
        assert mentions[0].partition == "DataSheet"


# ── Scorer ───────────────────────────────────────────────────────────────────

from databridge_core.linker._scorer import (
    LINK_THRESHOLD,
    _archetype_compatibility,
    _formula_similarity,
    _meaning_similarity,
    _name_similarity,
    _partition_compatibility,
    score_pair,
)


class TestScorer:
    """Test composite scoring functions."""

    def _make_mention(self, **kwargs) -> EntityMention:
        defaults = {
            "file_path": "data/test.xlsx",
            "archetype": "Financial Report",
            "raw_name": "Revenue",
            "normalized_name": "revenue",
            "formula": "=SUM(A1:A10)",
            "business_meaning": "Total Revenue",
        }
        defaults.update(kwargs)
        return EntityMention(**defaults)

    def test_name_similarity_identical(self):
        a = self._make_mention(normalized_name="revenue_west")
        b = self._make_mention(normalized_name="revenue_west")
        assert _name_similarity(a, b) == 1.0

    def test_name_similarity_similar(self):
        a = self._make_mention(raw_name="Revenue_West", normalized_name="revenue_west")
        b = self._make_mention(raw_name="Revenue_Western", normalized_name="revenue_western")
        score = _name_similarity(a, b)
        assert 0.7 < score < 1.0

    def test_name_similarity_different(self):
        a = self._make_mention(raw_name="Revenue", normalized_name="revenue")
        b = self._make_mention(raw_name="Headcount", normalized_name="headcount")
        score = _name_similarity(a, b)
        assert score < 0.4

    def test_name_similarity_empty(self):
        a = self._make_mention(normalized_name="")
        b = self._make_mention(normalized_name="revenue")
        assert _name_similarity(a, b) == 0.0

    def test_formula_similarity_identical(self):
        a = self._make_mention(formula="=SUM(B2:B11)")
        b = self._make_mention(formula="=SUM(B2:B11)")
        assert _formula_similarity(a, b) == 1.0

    def test_formula_similarity_empty(self):
        a = self._make_mention(formula="")
        b = self._make_mention(formula="=SUM(B2:B11)")
        assert _formula_similarity(a, b) == 0.0

    def test_meaning_similarity_identical(self):
        a = self._make_mention(business_meaning="Total Revenue")
        b = self._make_mention(business_meaning="Total Revenue")
        assert _meaning_similarity(a, b) == 1.0

    def test_meaning_similarity_case_insensitive(self):
        a = self._make_mention(business_meaning="Total Revenue")
        b = self._make_mention(business_meaning="total revenue")
        assert _meaning_similarity(a, b) == 1.0

    def test_archetype_compatibility_same_group(self):
        a = self._make_mention(archetype="Financial Report")
        b = self._make_mention(archetype="Budget")
        assert _archetype_compatibility(a, b) == 1.0

    def test_archetype_compatibility_related(self):
        a = self._make_mention(archetype="Financial Report")
        b = self._make_mention(archetype="Dashboard")
        score = _archetype_compatibility(a, b)
        assert score == 0.7

    def test_archetype_compatibility_unrelated(self):
        a = self._make_mention(archetype="Financial Report")
        b = self._make_mention(archetype="Data Extract")
        score = _archetype_compatibility(a, b)
        assert score < 0.5

    def test_score_pair_identical(self):
        a = self._make_mention()
        b = self._make_mention(file_path="data/other.xlsx")
        composite, components, strategies = score_pair(a, b)
        assert composite >= 0.9
        assert len(strategies) > 0

    def test_score_pair_above_threshold(self):
        a = self._make_mention(
            normalized_name="revenue_west",
            business_meaning="Total Western Region Revenue",
        )
        b = self._make_mention(
            file_path="data/other.xlsx",
            normalized_name="revenue_western",
            business_meaning="Western Region Revenue Total",
        )
        composite, _, _ = score_pair(a, b)
        assert composite >= LINK_THRESHOLD

    def test_score_pair_below_threshold(self):
        a = self._make_mention(
            raw_name="Revenue_West",
            normalized_name="revenue_west",
            formula="=SUM(A1:A10)",
            business_meaning="Western Revenue",
        )
        b = self._make_mention(
            raw_name="Headcount_East",
            normalized_name="headcount_east",
            formula="=COUNT(B1:B10)",
            business_meaning="Eastern Headcount",
        )
        composite, _, _ = score_pair(a, b)
        assert composite < LINK_THRESHOLD

    def test_link_threshold_value(self):
        assert LINK_THRESHOLD == 0.65

    def test_partition_compatibility_same_sheet(self):
        a = self._make_mention(file_path="file.xlsx", partition="Sheet1")
        b = self._make_mention(file_path="file.xlsx", partition="Sheet1")
        assert _partition_compatibility(a, b) == 0.0

    def test_partition_compatibility_cross_sheet_same_file(self):
        a = self._make_mention(file_path="file.xlsx", partition="Sheet1")
        b = self._make_mention(file_path="file.xlsx", partition="Summary")
        assert _partition_compatibility(a, b) == 0.6

    def test_partition_compatibility_cross_file_same_type(self):
        a = self._make_mention(file_path="a.xlsx", partition="Data1", partition_type="data")
        b = self._make_mention(file_path="b.xlsx", partition="Data2", partition_type="data")
        assert _partition_compatibility(a, b) == 0.8

    def test_partition_compatibility_cross_file_different_type(self):
        a = self._make_mention(file_path="a.xlsx", partition="Data1", partition_type="data")
        b = self._make_mention(file_path="b.xlsx", partition="Summary", partition_type="summary")
        assert _partition_compatibility(a, b) == 0.4

    def test_partition_compatibility_neutral(self):
        a = self._make_mention(file_path="a.xlsx", partition="", partition_type="")
        b = self._make_mention(file_path="b.xlsx", partition="", partition_type="")
        assert _partition_compatibility(a, b) == 0.5

    def test_partition_compatibility_one_type_missing(self):
        a = self._make_mention(file_path="a.xlsx", partition="Sheet1", partition_type="data")
        b = self._make_mention(file_path="b.xlsx", partition="Sheet2", partition_type="")
        assert _partition_compatibility(a, b) == 0.5

    def test_score_pair_with_partition(self):
        """Verify partition contributes to composite score."""
        a = self._make_mention(
            file_path="a.xlsx", partition="Data1", partition_type="data",
        )
        b = self._make_mention(
            file_path="b.xlsx", partition="Data2", partition_type="data",
        )
        composite, components, strategies = score_pair(a, b)
        assert MatchStrategy.PARTITION_COMPATIBILITY.value in components
        assert components[MatchStrategy.PARTITION_COMPATIBILITY.value] == 0.8


# ── Linker ───────────────────────────────────────────────────────────────────

from databridge_core.linker._linker import EntityLinker, _UnionFind, _detect_conflicts, _infer_domain


class TestUnionFind:
    """Test Union-Find data structure."""

    def test_basic_union(self):
        uf = _UnionFind()
        uf.make_set("a")
        uf.make_set("b")
        uf.union("a", "b")
        assert uf.find("a") == uf.find("b")

    def test_transitive_union(self):
        uf = _UnionFind()
        for x in "abcde":
            uf.make_set(x)
        uf.union("a", "b")
        uf.union("b", "c")
        uf.union("d", "e")
        assert uf.find("a") == uf.find("c")
        assert uf.find("d") == uf.find("e")
        assert uf.find("a") != uf.find("d")

    def test_clusters(self):
        uf = _UnionFind()
        for x in "abcde":
            uf.make_set(x)
        uf.union("a", "b")
        uf.union("b", "c")
        clusters = uf.clusters()
        # Should have 3 clusters: {a,b,c}, {d}, {e}
        sizes = sorted(len(v) for v in clusters.values())
        assert sizes == [1, 1, 3]


class TestConflictDetection:
    """Test conflict detection within clusters."""

    def _make_mention(self, **kwargs) -> EntityMention:
        defaults = {
            "file_path": "data/test.xlsx",
            "archetype": "Financial Report",
            "raw_name": "Revenue",
            "normalized_name": "revenue",
            "formula": "",
            "business_meaning": "Revenue",
        }
        defaults.update(kwargs)
        return EntityMention(**defaults)

    def test_no_conflicts_same_formula(self):
        mentions = [
            self._make_mention(formula="=SUM(A1:A10)"),
            self._make_mention(formula="=SUM(A1:A10)"),
        ]
        conflicts = _detect_conflicts("c1", mentions)
        assert len(conflicts) == 0

    def test_formula_mismatch(self):
        mentions = [
            self._make_mention(formula="=SUM(A1:A10)"),
            self._make_mention(formula="=AVERAGE(B1:B10)"),
        ]
        conflicts = _detect_conflicts("c1", mentions)
        formula_conflicts = [c for c in conflicts if c.conflict_type == ConflictType.FORMULA_MISMATCH]
        assert len(formula_conflicts) >= 1

    def test_sign_reversal(self):
        mentions = [
            self._make_mention(formula="=SUM(A1:A10)", business_meaning="Revenue total"),
            self._make_mention(formula="-B5", business_meaning="SUBTRACT from total"),
        ]
        conflicts = _detect_conflicts("c1", mentions)
        sign_conflicts = [c for c in conflicts if c.conflict_type == ConflictType.SIGN_REVERSAL]
        assert len(sign_conflicts) >= 1

    def test_no_conflicts_empty_formulas(self):
        mentions = [
            self._make_mention(formula=""),
            self._make_mention(formula=""),
        ]
        conflicts = _detect_conflicts("c1", mentions)
        assert len(conflicts) == 0


class TestDomainInference:
    """Test domain inference from mentions."""

    def _make_mention(self, **kwargs) -> EntityMention:
        defaults = {"file_path": "", "archetype": "", "raw_name": "", "normalized_name": ""}
        defaults.update(kwargs)
        return EntityMention(**defaults)

    def test_revenue_domain(self):
        mentions = [self._make_mention(normalized_name="revenue_west", business_meaning="revenue total")]
        assert _infer_domain(mentions) == "revenue"

    def test_expense_domain(self):
        mentions = [self._make_mention(normalized_name="cost_center", business_meaning="total expense")]
        assert _infer_domain(mentions) == "expense"

    def test_general_domain(self):
        mentions = [self._make_mention(normalized_name="xyz_abc", business_meaning="something")]
        assert _infer_domain(mentions) == "general"


class TestEntityLinker:
    """Test the main EntityLinker class."""

    def _make_mention(self, file_path: str, name: str, **kwargs) -> EntityMention:
        defaults = {
            "file_path": file_path,
            "archetype": "Financial Report",
            "raw_name": name,
            "normalized_name": name.lower().replace(" ", "_"),
            "formula": "=SUM(A1:A10)",
            "business_meaning": f"Total {name}",
        }
        defaults.update(kwargs)
        return EntityMention(**defaults)

    def test_link_empty(self):
        linker = EntityLinker()
        result = linker.link([])
        assert result.total_clusters == 0

    def test_link_single_file_no_clusters(self):
        """Mentions from the same file should NOT be linked."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue"),
            self._make_mention("file_a.xlsx", "Expense"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        # No cross-file links -> no multi-mention clusters
        assert result.total_clusters == 0

    def test_link_cross_file_identical(self):
        """Identical mentions across files should cluster."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue"),
            self._make_mention("file_b.xlsx", "Revenue"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        assert result.total_clusters >= 1
        # The cluster should contain mentions from both files
        cluster = result.clusters[0]
        files = {m.file_path for m in cluster.mentions}
        assert len(files) == 2

    def test_link_cross_file_similar(self):
        """Similar mentions across files should cluster."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue_West"),
            self._make_mention("file_b.xlsx", "Revenue_Western"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        assert result.total_clusters >= 1

    def test_link_cross_file_different(self):
        """Very different mentions should NOT cluster."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue", formula="=SUM(A:A)", business_meaning="Total sales income"),
            self._make_mention("file_b.xlsx", "Headcount", formula="=COUNT(B:B)", business_meaning="Number of employees"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        assert result.total_clusters == 0

    def test_link_three_files_transitive(self):
        """If A<->B and B<->C link, all three should be in one cluster."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue"),
            self._make_mention("file_b.xlsx", "Revenue"),
            self._make_mention("file_c.xlsx", "Revenue"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        assert result.total_clusters == 1
        assert len(result.clusters[0].mentions) == 3

    def test_link_canonical_name(self):
        """Canonical name should be set on the cluster."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue_West"),
            self._make_mention("file_b.xlsx", "Revenue_West"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        assert result.clusters[0].canonical_name != ""

    def test_link_domain_inference(self):
        """Domain should be inferred from mention names."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue_West", business_meaning="Western region revenue"),
            self._make_mention("file_b.xlsx", "Revenue_West", business_meaning="West revenue total"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        assert result.clusters[0].domain == "revenue"

    def test_link_with_conflict(self):
        """Different formulas in same cluster should produce a conflict."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue", formula="=SUM(A1:A10)"),
            self._make_mention("file_b.xlsx", "Revenue", formula="=AVERAGE(B1:B10)"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        assert result.total_clusters >= 1
        assert result.total_conflicts >= 1

    def test_link_custom_threshold(self):
        """Higher threshold should be stricter."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue_West"),
            self._make_mention("file_b.xlsx", "Revenue_East"),
        ]
        strict = EntityLinker(threshold=0.95)
        result = strict.link(mentions)
        # Very high threshold -- names are somewhat different
        # This may or may not cluster depending on other scores
        # Just verify it runs without error
        assert isinstance(result, EntityMap)

    def test_entity_map_serialization(self):
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue"),
            self._make_mention("file_b.xlsx", "Revenue"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        d = result.model_dump(mode="json")
        assert "clusters" in d
        assert "total_mentions" in d
        assert "map_id" in d

    def test_link_cross_sheet_same_file(self):
        """Entities on different sheets of the same file should cluster."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue", partition="Sheet1"),
            self._make_mention("file_a.xlsx", "Revenue", partition="Summary"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        assert result.total_clusters >= 1
        cluster = result.clusters[0]
        assert len(cluster.mentions) == 2

    def test_link_same_sheet_suppressed(self):
        """Entities on the same sheet of the same file should NOT cluster."""
        mentions = [
            self._make_mention("file_a.xlsx", "Revenue", partition="Sheet1"),
            self._make_mention("file_a.xlsx", "Expense", partition="Sheet1"),
        ]
        linker = EntityLinker()
        result = linker.link(mentions)
        # Same partition -> grouped together -> no cross-partition comparison
        assert result.total_clusters == 0


# ── Public API ───────────────────────────────────────────────────────────────

from databridge_core.linker import (
    find_entity,
    get_entity_cluster,
    get_entity_map,
    get_link_summary,
    link_entities,
)


class TestPublicAPI:
    """Test the public API functions."""

    def _write_dna_files(self, tmpdir: str) -> None:
        """Write two Logic DNA files with a shared entity for testing."""
        dna_a = {
            "file_path": "data/file_a.xlsx",
            "archetype": "Financial Report",
            "formula_intents": [
                {
                    "cell_ref": "Sheet1!B12",
                    "raw_formula": "=SUM(B2:B11)",
                    "business_meaning": '"Revenue_West" total for Q1',
                },
                {
                    "cell_ref": "Sheet1!B15",
                    "raw_formula": "=SUM(C2:C11)",
                    "business_meaning": '"Expense_Admin" total for Q1',
                },
            ],
            "cross_references": [],
        }
        dna_b = {
            "file_path": "data/file_b.xlsx",
            "archetype": "Financial Report",
            "formula_intents": [
                {
                    "cell_ref": "Summary!A1",
                    "raw_formula": "=SUM(A2:A11)",
                    "business_meaning": '"Revenue_West" consolidated',
                },
            ],
            "cross_references": [],
        }
        for name, dna in [("file_a_logic_dna.json", dna_a), ("file_b_logic_dna.json", dna_b)]:
            with open(Path(tmpdir) / name, "w") as f:
                json.dump(dna, f)

    def test_link_entities_basic(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            dna_dir = Path(tmpdir) / "debate"
            dna_dir.mkdir()
            out_dir = Path(tmpdir) / "linker"

            self._write_dna_files(str(dna_dir))

            result = link_entities(
                logic_dna_dir=str(dna_dir),
                output_dir=str(out_dir),
            )
            assert "summary" in result
            assert "output_files" in result
            assert (out_dir / "entity_map.json").exists()
            assert (out_dir / "entity_clusters.jsonl").exists()

    def test_link_entities_empty_dir(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            result = link_entities(logic_dna_dir=tmpdir)
            assert "error" in result

    def test_link_entities_missing_dir(self):
        result = link_entities(logic_dna_dir="/nonexistent/dir")
        assert "error" in result

    def test_get_entity_map(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            dna_dir = Path(tmpdir) / "debate"
            dna_dir.mkdir()
            out_dir = Path(tmpdir) / "linker"

            self._write_dna_files(str(dna_dir))
            link_entities(logic_dna_dir=str(dna_dir), output_dir=str(out_dir))

            result = get_entity_map(output_dir=str(out_dir))
            assert "map_id" in result
            assert "total_clusters" in result

    def test_get_entity_map_missing(self):
        result = get_entity_map(output_dir="/nonexistent/dir")
        assert "error" in result

    def test_get_entity_cluster(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            dna_dir = Path(tmpdir) / "debate"
            dna_dir.mkdir()
            out_dir = Path(tmpdir) / "linker"

            self._write_dna_files(str(dna_dir))
            link_result = link_entities(logic_dna_dir=str(dna_dir), output_dir=str(out_dir))

            # Get the first cluster ID from the sample
            if link_result.get("sample_clusters"):
                cid = link_result["sample_clusters"][0]["cluster_id"]
                cluster = get_entity_cluster(cluster_id=cid, output_dir=str(out_dir))
                assert "cluster_id" in cluster
                assert cluster["cluster_id"] == cid

    def test_get_entity_cluster_missing(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            # Write a minimal entity map
            out_dir = Path(tmpdir) / "linker"
            out_dir.mkdir()
            with open(out_dir / "entity_map.json", "w") as f:
                json.dump({"clusters": []}, f)

            result = get_entity_cluster(cluster_id="nonexistent", output_dir=str(out_dir))
            assert "error" in result

    def test_get_link_summary(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            dna_dir = Path(tmpdir) / "debate"
            dna_dir.mkdir()
            out_dir = Path(tmpdir) / "linker"

            self._write_dna_files(str(dna_dir))
            link_entities(logic_dna_dir=str(dna_dir), output_dir=str(out_dir))

            summary = get_link_summary(output_dir=str(out_dir))
            assert "total_clusters" in summary
            assert "domain_counts" in summary

    def test_get_link_summary_missing(self):
        result = get_link_summary(output_dir="/nonexistent/dir")
        assert "error" in result

    def test_find_entity_basic(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            dna_dir = Path(tmpdir) / "debate"
            dna_dir.mkdir()
            out_dir = Path(tmpdir) / "linker"

            self._write_dna_files(str(dna_dir))
            link_entities(logic_dna_dir=str(dna_dir), output_dir=str(out_dir))

            result = find_entity(query="Revenue", output_dir=str(out_dir))
            assert "results" in result
            assert result["query"] == "Revenue"

    def test_find_entity_missing(self):
        result = find_entity(query="test", output_dir="/nonexistent/dir")
        assert "error" in result

    def test_link_with_partition_metadata_enrichment(self):
        """Partition metadata should enrich mention partition_type fields."""
        with tempfile.TemporaryDirectory() as tmpdir:
            dna_dir = Path(tmpdir) / "debate"
            dna_dir.mkdir()
            out_dir = Path(tmpdir) / "linker"

            self._write_dna_files(str(dna_dir))

            partition_metadata = {
                "Sheet1": {"type": "data"},
                "Summary": {"type": "summary"},
            }

            result = link_entities(
                logic_dna_dir=str(dna_dir),
                output_dir=str(out_dir),
                partition_metadata=partition_metadata,
            )
            assert "summary" in result
            # Verify the entity map was created with partition info
            map_file = Path(out_dir) / "entity_map.json"
            assert map_file.exists()
            with open(map_file) as f:
                data = json.load(f)
            # Check that at least some mentions have partition_type set
            all_mentions = []
            for cluster in data.get("clusters", []):
                all_mentions.extend(cluster.get("mentions", []))
            # Even if no clusters, the run should succeed without error
            assert "error" not in result
