"""
extractor.analytics.corpus_analytics — Corpus-level entity statistics.

All functions are pure / deterministic; no external dependencies.
"""

from __future__ import annotations

from .entity_index import EntityIndex


def entity_frequency_table(
    index: EntityIndex,
    entity_type: str | None = None,
    min_mentions: int = 1,
) -> list[dict]:
    """
    Return a frequency table for entities, sorted by mention_count descending.

    Each row: {"canonical_id", "canonical_text", "entity_type", "mention_count", "doc_count"}
    """
    rows: list[dict] = []
    for entity in index.top_entities(entity_type=entity_type, n=10_000_000):
        if entity.mention_count < min_mentions:
            continue
        rows.append(
            {
                "canonical_id": entity.canonical_id,
                "canonical_text": entity.canonical_text,
                "entity_type": entity.entity_type,
                "mention_count": entity.mention_count,
                "doc_count": entity.doc_count,
            }
        )
    # already sorted by mention_count desc from top_entities
    return rows


def entity_doc_cooccurrence(
    index: EntityIndex,
    canonical_id_a: str,
    canonical_id_b: str,
) -> list[str]:
    """
    Return list of source_filenames where both entities appear.
    """
    entity_a = index.get(canonical_id_a)
    entity_b = index.get(canonical_id_b)
    if entity_a is None or entity_b is None:
        return []

    docs_a = {m.source_filename for m in entity_a.mentions}
    docs_b = {m.source_filename for m in entity_b.mentions}
    return sorted(docs_a & docs_b)


def top_entities_by_section(
    index: EntityIndex,
    section: str,
    entity_type: str | None = None,
    n: int = 10,
) -> list[dict]:
    """
    Return top entities (by mention count within matching sections), optionally
    filtered by entity_type.

    A mention matches if any element in its section_path contains `section`
    as a case-insensitive substring.
    """
    section_lower = section.lower()
    counts: dict[str, int] = {}

    all_entities = index.top_entities(entity_type=entity_type, n=10_000_000)
    for entity in all_entities:
        cnt = 0
        for mention in entity.mentions:
            if any(section_lower in seg.lower() for seg in mention.section_path):
                cnt += 1
        if cnt > 0:
            counts[entity.canonical_id] = cnt

    # Sort by section mention count desc, then take top n
    sorted_ids = sorted(counts, key=lambda cid: counts[cid], reverse=True)[:n]

    result: list[dict] = []
    for cid in sorted_ids:
        entity = index.get(cid)
        if entity is None:
            continue
        result.append(
            {
                "canonical_id": entity.canonical_id,
                "canonical_text": entity.canonical_text,
                "entity_type": entity.entity_type,
                "mention_count": counts[cid],
                "doc_count": entity.doc_count,
            }
        )
    return result
