"""Candidate pair generation for the Lint pass.

For each source claim atom, we pull its nearest neighbors via the Atomic
/api/atoms/{id}/similar endpoint (which uses pgvector kNN under the hood).
Only Type/Claim neighbors are kept; Type/Document neighbors are dropped so
we only compare claim-to-claim.

Pairs are canonicalized as (lower_id, higher_id) so (A, B) and (B, A)
produce the same CandidatePair and don't get evaluated twice.
"""

from __future__ import annotations

from dataclasses import dataclass

from shared.atomic_api import AtomicClient
from shared.logging import get_logger

log = get_logger(__name__)

# Atoms with similarity below this are considered too unrelated to be worth
# an NLI call. Between 0.55 and 0.95 is the interesting band — below is
# "probably different topic", above is "almost certainly the same text".
MIN_PAIR_SIMILARITY = 0.55
NEIGHBORS_PER_CLAIM = 15


@dataclass(slots=True, frozen=True)
class CandidatePair:
    """A pair of claim atom ids to evaluate, canonicalized lo<hi."""

    atom_a_id: str
    atom_b_id: str
    similarity: float


def _is_claim_hit(hit_tags: list[dict]) -> bool:
    return any(t.get("name") == "Claim" for t in hit_tags)


async def generate_candidates(
    atomic: AtomicClient,
    *,
    source_atom_ids: list[str],
    min_similarity: float = MIN_PAIR_SIMILARITY,
    neighbors_per: int = NEIGHBORS_PER_CLAIM,
) -> list[CandidatePair]:
    """Walk each source atom and collect canonical candidate pairs.

    Limitation: Atomic's /similar endpoint returns Document and Claim atoms
    mixed together. We filter client-side for Type/Claim. It might be worth
    upstream adding a tag filter later, but for this corpus size the current
    approach is fine.
    """
    seen: set[tuple[str, str]] = set()
    out: list[CandidatePair] = []

    for idx, src_id in enumerate(source_atom_ids):
        try:
            hits = await atomic.find_similar(
                src_id, threshold=min_similarity, limit=neighbors_per
            )
        except Exception as e:  # noqa: BLE001
            log.warning(
                "find_similar_failed", atom_id=src_id, error=f"{type(e).__name__}:{e}"
            )
            continue

        for h in hits:
            if h.atom_id == src_id:
                continue
            if not _is_claim_hit(h.tags):
                continue
            lo, hi = sorted((src_id, h.atom_id))
            key = (lo, hi)
            if key in seen:
                continue
            seen.add(key)
            out.append(
                CandidatePair(
                    atom_a_id=lo, atom_b_id=hi, similarity=round(h.similarity, 4)
                )
            )

        if (idx + 1) % 50 == 0:
            log.info(
                "candidates_progress",
                processed=idx + 1,
                total=len(source_atom_ids),
                pairs=len(out),
            )

    return out