"""TF-IDF keyword distiller tests."""

from __future__ import annotations

import json
from typing import Iterator

from aborist.distill import TfidfKeywordDistiller
from aborist.distill.runner import distill_existing
from aborist.document import Document
from aborist.ingest import ingest_source
from aborist.merkle import proof_from_dict, verify_proof
from aborist.source import Source
from aborist.store import connect


class FakeSource(Source):
    source_type = "test"

    def __init__(self, docs: list[Document]):
        self.docs = docs

    def iter_documents(self) -> Iterator[Document]:
        yield from self.docs


def test_tfidf_picks_distinctive_terms():
    """A term concentrated in one chunk should outrank one spread across all."""
    distiller = TfidfKeywordDistiller(top_k=4)
    chunks = [
        "merkle merkle merkle providence cryptography",
        "the quick brown fox common common common common",
        "wikipedia article history history history",
    ]
    src = Document(uri="x", content="", source_type="t", title="t")
    result = distiller.distill(src, chunks)
    keywords = result.core.content.split(", ")
    # 'merkle' (3x in chunk 0, 0 elsewhere) should rank high.
    assert "merkle" in keywords
    # 'history' should also rank high (3x in chunk 2).
    assert "history" in keywords


def test_tfidf_deterministic():
    """Same input -> same output."""
    d = TfidfKeywordDistiller(top_k=8)
    src = Document(uri="x", content="", source_type="t", title="t")
    chunks = [
        "merkle providence proves provenance cryptographically",
        "providence cache stores answer records",
    ]
    a = d.distill(src, chunks)
    b = d.distill(src, chunks)
    assert a.core.content == b.core.content
    assert a.contributing_chunk_indices == b.contributing_chunk_indices


def test_tfidf_distillation_full_round_trip(tmp_path):
    db = tmp_path / "tfidf.db"
    conn = connect(db)
    try:
        long_text = (
            "merkle providence reverse rag verifies provenance cryptographically. " * 30
            + "\n\n"
            + "wikipedia article from 2003 about anarchism. " * 30
            + "\n\n"
            + "the eight forms of capital include living, social, intellectual. " * 30
        )
        ingest_source(conn, FakeSource([Document(uri="t://x", content=long_text, source_type="test", title="X")]))
        result = distill_existing(conn, TfidfKeywordDistiller(top_k=8))
        assert result["distilled"] == 1

        core = conn.execute(
            "SELECT * FROM documents WHERE kind='core'"
        ).fetchone()
        assert core["source_type"] == "core:tfidf-keywords-v1"
        assert "[KEYWORDS]" in (core["title"] or "")

        # Core content is comma-separated tokens.
        chunk_row = conn.execute(
            "SELECT content FROM chunks WHERE document_root=? ORDER BY idx",
            (core["document_root"],),
        ).fetchone()
        keywords = chunk_row["content"].split(", ")
        assert len(keywords) <= 8
        # Distinctive terms should appear.
        assert any(k in {"merkle", "providence", "anarchism", "wikipedia", "capital"} for k in keywords)

        # Per-chunk Merkle proofs in derivation must reconstruct source root.
        deriv = conn.execute("SELECT proof_blob, src_root FROM derivations").fetchone()
        blob = json.loads(deriv["proof_blob"])
        for entry in blob["contributing"]:
            p = proof_from_dict(entry["proof"])
            assert verify_proof(p)
            assert p.root.hex() == deriv["src_root"]
    finally:
        conn.close()
