#!/usr/bin/env python3
"""
entity_detector.py — Auto-detect people and projects from file content.

Uses ``from __future__ import annotations`` so PEP 604 union syntax
(``dict | None``) works on the Python 3.9 baseline.

Two-pass approach:
  Pass 1: scan files, extract entity candidates with signal counts
  Pass 2: score and classify each candidate as person, project, or uncertain

Used by mempalace init before mining begins.
The confirmed entity map feeds the miner as the taxonomy.

Multi-language support:
    All lexical patterns (person verbs, pronouns, dialogue markers, project
    verbs, stopwords, and the candidate-extraction character class) live in
    the ``entity`` section of ``mempalace/i18n/<lang>.json``. Every public
    function accepts a ``languages`` tuple and applies the union of the
    requested locales' patterns. The default is ``("en",)`` — existing
    English-only callers behave exactly as before.

    To add a new language: add an ``entity`` section to that locale's JSON.
    No code changes required.

Usage:
    from mempalace.entity_detector import detect_entities, confirm_entities
    candidates = detect_entities(file_paths)                    # English only
    candidates = detect_entities(paths, languages=("en", "pt-br"))
    confirmed = confirm_entities(candidates)  # interactive review
"""

from __future__ import annotations

import json
import re
import os
import functools
from pathlib import Path
from collections import defaultdict

from mempalace.i18n import get_entity_patterns


# ==================== COCA CONTENT-WORD FILTER (Tier 2 linguistics cleanup) ====================
#
# Common English content words that frequently appear capitalized (sentence
# start, headings, markdown emphasis) but are NOT proper nouns. Filtering
# these at candidate-extraction time prevents false-positive entity detection
# of words like "Code", "Brutal", "Phase", "Chat", "Note", "Line", etc.
#
# The data file lives at ``mempalace/data/coca_content_words.json``. Loaded
# once on first call via ``_get_coca_filter``. Matching is case-insensitive:
# callers must lowercase the candidate before lookup.
#
# Tier 3 (planned) will add a known-systems lexicon that protects compound
# names like "Claude Code" — for now, the multi-word path in
# ``extract_candidates`` is intentionally NOT filtered, so legitimate
# compounds remain detectable.


@functools.lru_cache(maxsize=1)
def _get_coca_filter() -> frozenset[str]:
    """Return the COCA content-word filter set (lowercased).

    Loads ``mempalace/data/coca_content_words.json`` on first call and
    caches the resulting frozenset. Subsequent calls are O(1). Returns
    an empty frozenset if the data file is missing or malformed —
    extraction behavior then degrades gracefully (no filter applied)
    rather than crashing.
    """
    data_path = Path(__file__).parent / "data" / "coca_content_words.json"
    try:
        raw = json.loads(data_path.read_text(encoding="utf-8"))
        words = raw.get("words", [])
        return frozenset(w.lower() for w in words if isinstance(w, str))
    except (OSError, json.JSONDecodeError, AttributeError, TypeError):
        return frozenset()


# ==================== KNOWN-SYSTEMS COMPOUND LEXICON (Tier 3 linguistics cleanup) ====================
#
# Multi-word product / system names that must be detected atomically — NOT
# decomposed into their constituent words. When "Claude Code" appears in
# content, the entity detector counts the compound, not the parts. Without
# this pre-pass, the single-word loop would split "Claude Code" into
# "Claude" + "Code", and the COCA filter (Tier 2) would drop "Code" as a
# content word — leaving "Claude" alone with the wrong attribution.
#
# Data file: ``mempalace/data/known_systems.json``. Loaded once on first
# call via ``_get_known_systems``. Matching is case-insensitive with word
# boundaries.


@functools.lru_cache(maxsize=1)
def _get_known_systems() -> tuple[tuple[str, "re.Pattern[str]"], ...]:
    """Return the known-systems compound tuple — pairs of (canonical name,
    pre-compiled case-insensitive word-bounded regex).

    Loads ``mempalace/data/known_systems.json`` on first call, compiles a
    regex for each valid compound, and caches the resulting tuple of
    pairs. Subsequent calls are O(1) and skip both the disk read AND the
    regex compilation. Returns an empty tuple if the data file is missing
    or malformed — extraction behavior then degrades gracefully
    (compounds detected only by the existing multi-word regex) rather
    than crashing.

    Entries are sorted by length descending so the compound matcher
    prefers longer matches first (e.g. "Visual Studio Code" wins over
    a hypothetical "Visual Studio" if both were in the lexicon).
    """
    data_path = Path(__file__).parent / "data" / "known_systems.json"
    try:
        raw = json.loads(data_path.read_text(encoding="utf-8"))
        compounds = raw.get("compounds", [])
        valid = [c for c in compounds if isinstance(c, str) and c.strip()]
        # Sort by length descending so longest-match-wins during the
        # pre-pass scan (longer compounds get masked first, so a shorter
        # compound contained within a longer one doesn't double-count).
        sorted_compounds = sorted(valid, key=len, reverse=True)

        compiled: list[tuple[str, re.Pattern[str]]] = []
        for c in sorted_compounds:
            # Word-boundary, case-insensitive. Compound may contain
            # hyphens or spaces — re.escape handles special chars; word
            # boundaries on each side prevent partial-word matches
            # (e.g. "GPT-4" must not match "GPT-40").
            pattern = r"(?<!\w)" + re.escape(c) + r"(?!\w)"
            try:
                compiled.append((c, re.compile(pattern, re.IGNORECASE)))
            except re.error:
                continue
        return tuple(compiled)
    except (OSError, json.JSONDecodeError, AttributeError, TypeError):
        return ()


def _apply_known_systems_prepass(text: str) -> tuple[str, dict[str, int]]:
    """Scan ``text`` for known-systems compounds, return a working copy
    with matched spans masked to whitespace plus a dict of detected
    compound counts.

    Returning the counts (instead of mutating a caller-supplied container)
    lets the three call sites (``extract_candidates`` at init-time,
    ``palace.build_closet_lines`` at closet construction, and
    ``miner._extract_entities_for_metadata`` at per-drawer tagging) use
    whichever container shape they already maintain.

    Compounds are matched case-insensitively with word boundaries; the
    canonical (lexicon) casing is what gets counted, regardless of how
    the compound appears in the source text. Regexes are pre-compiled
    once in ``_get_known_systems`` so this function does no compilation.
    """
    compounds = _get_known_systems()
    if not compounds:
        return text, {}
    working = text
    compound_counts: dict[str, int] = {}
    for compound, rx in compounds:
        matches = list(rx.finditer(working))
        if not matches:
            continue
        compound_counts[compound] = compound_counts.get(compound, 0) + len(matches)
        # Mask matched spans with spaces so the subsequent regex passes
        # don't re-decompose. Replacing right-to-left keeps earlier
        # indices stable.
        for m in reversed(matches):
            start, end = m.span()
            working = working[:start] + (" " * (end - start)) + working[end:]
    return working, compound_counts


# ==================== LANGUAGE-AWARE PATTERN LOADING ====================


def _normalize_langs(languages) -> tuple:
    """Coerce a language input into a non-empty hashable tuple."""
    if not languages:
        return ("en",)
    if isinstance(languages, str):
        return (languages,)
    return tuple(languages)


@functools.lru_cache(maxsize=32)
def _get_stopwords(languages: tuple) -> frozenset:
    """Return the union of stopwords across the given languages."""
    patterns = get_entity_patterns(languages)
    return frozenset(patterns["stopwords"])


# ==================== BACKWARD-COMPAT MODULE CONSTANTS ====================
#
# These mirror the old module-level constants so existing imports keep working.
# They reflect the English defaults and are populated at import time from
# ``mempalace/i18n/en.json``. Callers that need multi-language behavior should
# pass the ``languages`` parameter to the public functions below.

_EN = get_entity_patterns(("en",))

PERSON_VERB_PATTERNS = list(_EN["person_verb_patterns"])
PRONOUN_PATTERNS = list(_EN["pronoun_patterns"])
PRONOUN_RE = re.compile("|".join(PRONOUN_PATTERNS), re.IGNORECASE) if PRONOUN_PATTERNS else None
DIALOGUE_PATTERNS = list(_EN["dialogue_patterns"])
PROJECT_VERB_PATTERNS = list(_EN["project_verb_patterns"])
STOPWORDS = set(_EN["stopwords"])


# ==================== EXTENSION POINTS (not language-scoped) ====================

# For entity detection — prose only, no code files
# Code files have too many capitalized names (classes, functions) that aren't entities
PROSE_EXTENSIONS = {
    ".txt",
    ".md",
    ".rst",
    ".csv",
    ".tex",
    ".bib",
}

READABLE_EXTENSIONS = {
    ".txt",
    ".md",
    ".py",
    ".js",
    ".ts",
    ".json",
    ".yaml",
    ".yml",
    ".csv",
    ".rst",
    ".toml",
    ".sh",
    ".rb",
    ".go",
    ".rs",
}

SKIP_DIRS = {
    ".git",
    "node_modules",
    "__pycache__",
    ".venv",
    "venv",
    "env",
    "dist",
    "build",
    ".next",
    "coverage",
    ".mempalace",
    ".terraform",
    "vendor",
    "target",
}

# Files whose content is boilerplate prose — poisons entity detection.
# Matched by stem (case-insensitive), with or without an extension.
SKIP_FILENAMES = {
    "license",
    "licence",
    "copying",
    "copyright",
    "notice",
    "authors",
    "patents",
    "third_party_notices",
    "third-party-notices",
}


# ==================== CANDIDATE EXTRACTION ====================


def extract_candidates(text: str, languages=("en",)) -> dict:
    """
    Extract all capitalized proper noun candidates from text.
    Returns {name: frequency} for names appearing 3+ times.

    Each language contributes its own character-class pattern (e.g. ASCII
    for English, Latin+diacritics for pt-br, Cyrillic for Russian,
    Devanagari for Hindi). Matches from all languages are unioned.
    """
    langs = _normalize_langs(languages)
    patterns = get_entity_patterns(langs)
    stopwords = _get_stopwords(langs)
    coca_filter = _get_coca_filter()

    counts: defaultdict = defaultdict(int)

    # Tier 3 — known-systems compound pre-pass. Find compound product names
    # ("Claude Code", "GitHub Copilot", ...) FIRST and mask them out of the
    # working text so the subsequent single-word + multi-word loops don't
    # re-decompose them into their constituent tokens.
    working_text, compound_counts = _apply_known_systems_prepass(text)
    for compound, n in compound_counts.items():
        counts[compound] += n

    # Single-word candidates — one pre-wrapped pattern per language
    for wrapped_pat in patterns["candidate_patterns"]:
        try:
            rx = re.compile(wrapped_pat)
        except re.error:
            continue
        for word in rx.findall(working_text):
            wl = word.lower()
            if wl in stopwords:
                continue
            # Tier 2 linguistics cleanup: block common English content words
            # (Code, Brutal, Phase, Line, Note, ...) from entity candidacy.
            # Multi-word path below is intentionally not filtered so
            # compound names like "Claude Code" still get detected.
            if wl in coca_filter:
                continue
            if len(word) < 2:
                continue
            counts[word] += 1

    # Multi-word candidates — one pre-wrapped pattern per language.
    # Runs against the working_text (compounds already masked) so an
    # unknown two-word phrase like "Jane Smith" still gets caught by
    # the regex without competing with known compounds.
    for wrapped_pat in patterns["multi_word_patterns"]:
        try:
            rx = re.compile(wrapped_pat)
        except re.error:
            continue
        for phrase in rx.findall(working_text):
            if any(w.lower() in stopwords for w in phrase.split()):
                continue
            counts[phrase] += 1

    return {name: count for name, count in counts.items() if count >= 3}


# ==================== SIGNAL SCORING ====================


@functools.lru_cache(maxsize=256)
def _build_patterns(name: str, languages: tuple = ("en",)) -> dict:
    """Pre-compile all regex patterns for a single entity name, per language set."""
    n = re.escape(name)
    langs = _normalize_langs(languages)
    sources = get_entity_patterns(langs)

    def _compile_each(raw_patterns, flags=re.IGNORECASE):
        compiled = []
        for p in raw_patterns:
            try:
                compiled.append(re.compile(p.format(name=n), flags))
            except (re.error, KeyError, IndexError):
                continue
        return compiled

    direct_sources = sources.get("direct_address_patterns") or []
    direct_compiled = []
    for raw in direct_sources:
        try:
            direct_compiled.append(re.compile(raw.format(name=n), re.IGNORECASE))
        except (re.error, KeyError, IndexError):
            continue

    return {
        "dialogue": _compile_each(sources["dialogue_patterns"], re.MULTILINE | re.IGNORECASE),
        "person_verbs": _compile_each(sources["person_verb_patterns"]),
        "project_verbs": _compile_each(sources["project_verb_patterns"]),
        "direct": direct_compiled,
        "versioned": re.compile(rf"\b{n}[-_]v?\d+(?:\.\d+)*\b", re.IGNORECASE),
        "code_ref": re.compile(rf"\b{n}\.(py|js|ts|yaml|yml|json|sh)\b", re.IGNORECASE),
    }


@functools.lru_cache(maxsize=32)
def _pronoun_re(languages: tuple):
    """Compile a combined pronoun regex for the given languages."""
    langs = _normalize_langs(languages)
    patterns = get_entity_patterns(langs)
    pronouns = patterns.get("pronoun_patterns") or []
    if not pronouns:
        return None
    try:
        return re.compile("|".join(pronouns), re.IGNORECASE)
    except re.error:
        return None


def score_entity(name: str, text: str, lines: list, languages=("en",)) -> dict:
    """
    Score a candidate entity as person vs project.
    Returns scores and the signals that fired.
    """
    langs = _normalize_langs(languages)
    patterns = _build_patterns(name, langs)
    pronoun_re = _pronoun_re(langs)
    person_score = 0
    project_score = 0
    person_signals = []
    project_signals = []

    # --- Person signals ---

    # Dialogue markers (strong signal).
    # The bare `^NAME:\s` colon-prefix pattern matches metadata lines like
    # `Created: 2026-04-21`, so we require >= 2 hits for it to count as dialogue
    # (real speaker markers repeat; single-line metadata doesn't).
    for rx in patterns["dialogue"]:
        matches = len(rx.findall(text))
        if matches == 0:
            continue
        is_bare_colon = rx.pattern.endswith(r":\s") and not rx.pattern.endswith(r"[:\s]")
        if is_bare_colon and matches < 2:
            continue
        person_score += matches * 3
        person_signals.append(f"dialogue marker ({matches}x)")

    # Person verbs
    for rx in patterns["person_verbs"]:
        matches = len(rx.findall(text))
        if matches > 0:
            person_score += matches * 2
            person_signals.append(f"'{name} ...' action ({matches}x)")

    # Pronoun proximity — pronouns within 3 lines of the name
    if pronoun_re is not None:
        name_lower = name.lower()
        name_line_indices = [i for i, line in enumerate(lines) if name_lower in line.lower()]
        pronoun_hits = 0
        for idx in name_line_indices:
            window_text = " ".join(lines[max(0, idx - 2) : idx + 3])
            if pronoun_re.search(window_text):
                pronoun_hits += 1
        if pronoun_hits > 0:
            person_score += pronoun_hits * 2
            person_signals.append(f"pronoun nearby ({pronoun_hits}x)")

    # Direct address
    direct_hits = 0
    for rx in patterns["direct"]:
        direct_hits += len(rx.findall(text))
    if direct_hits > 0:
        person_score += direct_hits * 4
        person_signals.append(f"addressed directly ({direct_hits}x)")

    # --- Project signals ---

    for rx in patterns["project_verbs"]:
        matches = len(rx.findall(text))
        if matches > 0:
            project_score += matches * 2
            project_signals.append(f"project verb ({matches}x)")

    versioned = len(patterns["versioned"].findall(text))
    if versioned > 0:
        project_score += versioned * 3
        project_signals.append(f"versioned/hyphenated ({versioned}x)")

    code_ref = len(patterns["code_ref"].findall(text))
    if code_ref > 0:
        project_score += code_ref * 3
        project_signals.append(f"code file reference ({code_ref}x)")

    return {
        "person_score": person_score,
        "project_score": project_score,
        "person_signals": person_signals[:3],
        "project_signals": project_signals[:3],
    }


# ==================== CLASSIFY ====================


def classify_entity(name: str, frequency: int, scores: dict) -> dict:
    """
    Given scores, classify as person / project / uncertain.
    Returns entity dict with confidence.
    """
    ps = scores["person_score"]
    prs = scores["project_score"]
    total = ps + prs

    if total == 0:
        # No strong signals — frequency-only candidate, uncertain
        confidence = min(0.4, frequency / 50)
        return {
            "name": name,
            "type": "uncertain",
            "confidence": round(confidence, 2),
            "frequency": frequency,
            "signals": [f"appears {frequency}x, no strong type signals"],
        }

    person_ratio = ps / total if total > 0 else 0

    # Require TWO different signal categories to confidently classify as a person.
    # One signal type with many hits (e.g. "Click, click, click...") is not enough —
    # it just means that word appears often in a particular syntactic position.
    signal_categories = set()
    for s in scores["person_signals"]:
        if "dialogue" in s:
            signal_categories.add("dialogue")
        elif "action" in s:
            signal_categories.add("action")
        elif "pronoun" in s:
            signal_categories.add("pronoun")
        elif "addressed" in s:
            signal_categories.add("addressed")

    has_two_signal_types = len(signal_categories) >= 2
    # Single-category pronoun signal still classifies as person when the
    # evidence is overwhelming — a diary's main character is referenced
    # with pronouns, not dialogue markers. Require both: many pronoun hits
    # AND a high pronoun-to-frequency ratio so common sentence-start words
    # (Never, Before, etc.) with incidental pronoun proximity don't qualify.
    pronoun_hits = 0
    for s in scores["person_signals"]:
        m = re.search(r"pronoun nearby \((\d+)x\)", s)
        if m:
            pronoun_hits = int(m.group(1))
            break
    strong_pronoun_signal = pronoun_hits >= 5 and frequency > 0 and pronoun_hits / frequency >= 0.2

    if person_ratio >= 0.7 and (has_two_signal_types and ps >= 5 or strong_pronoun_signal):
        entity_type = "person"
        confidence = min(0.99, 0.5 + person_ratio * 0.5)
        signals = scores["person_signals"] or [f"appears {frequency}x"]
    elif person_ratio >= 0.7:
        # Weak single-category person signal — downgrade to uncertain
        entity_type = "uncertain"
        confidence = 0.4
        signals = scores["person_signals"] + [f"appears {frequency}x — weak person signal"]
    elif person_ratio <= 0.3:
        entity_type = "project"
        confidence = min(0.99, 0.5 + (1 - person_ratio) * 0.5)
        signals = scores["project_signals"] or [f"appears {frequency}x"]
    else:
        entity_type = "uncertain"
        confidence = 0.5
        signals = (scores["person_signals"] + scores["project_signals"])[:3]
        signals.append("mixed signals — needs review")

    return {
        "name": name,
        "type": entity_type,
        "confidence": round(confidence, 2),
        "frequency": frequency,
        "signals": signals,
    }


# ==================== MAIN DETECT ====================


def detect_entities(
    file_paths: list,
    max_files: int = 10,
    languages=("en",),
    corpus_origin: dict | None = None,
) -> dict:
    """
    Scan files and detect entity candidates.

    Args:
        file_paths: List of Path objects to scan
        max_files: Max files to read (for speed)
        languages: Tuple of language codes whose entity patterns should be
            applied (union). Defaults to ``("en",)``.
        corpus_origin: Optional corpus-origin context (the dict produced
            by ``mempalace.corpus_origin`` and persisted to
            ``<palace>/.mempalace/origin.json`` by ``mempalace init``).
            When supplied and the corpus is identified as AI-dialogue with
            known agent persona names, candidates whose name matches an
            agent persona are moved out of ``people``/``uncertain`` and
            into a new ``agent_personas`` bucket. Shape:
            ``{"schema_version": 1, "result": {"agent_persona_names": [...], ...}}``.

    Returns:
        {
            "people":   [...entity dicts...],
            "projects": [...entity dicts...],
            "topics":   [...entity dicts...],
            "uncertain":[...entity dicts...],
            # Only present when corpus_origin reclassifies at least one
            # candidate as an agent persona:
            "agent_personas": [...entity dicts...],
        }
    """
    langs = _normalize_langs(languages)

    # Collect text from files
    all_text = []
    all_lines = []
    files_read = 0

    MAX_BYTES_PER_FILE = 5_000  # first 5KB per file — enough to catch recurring entities

    for filepath in file_paths:
        if files_read >= max_files:
            break
        try:
            with open(filepath, encoding="utf-8", errors="replace") as f:
                content = f.read(MAX_BYTES_PER_FILE)
            all_text.append(content)
            all_lines.extend(content.splitlines())
            files_read += 1
        except OSError:
            continue

    combined_text = "\n".join(all_text)

    # Extract candidates
    candidates = extract_candidates(combined_text, languages=langs)

    if not candidates:
        return _apply_corpus_origin(
            {"people": [], "projects": [], "topics": [], "uncertain": []},
            corpus_origin,
        )

    # Score and classify each candidate
    people = []
    projects = []
    uncertain = []

    for name, frequency in sorted(candidates.items(), key=lambda x: x[1], reverse=True):
        scores = score_entity(name, combined_text, all_lines, languages=langs)
        entity = classify_entity(name, frequency, scores)

        if entity["type"] == "person":
            people.append(entity)
        elif entity["type"] == "project":
            projects.append(entity)
        else:
            uncertain.append(entity)

    # Sort by confidence descending
    people.sort(key=lambda x: x["confidence"], reverse=True)
    projects.sort(key=lambda x: x["confidence"], reverse=True)
    uncertain.sort(key=lambda x: x["frequency"], reverse=True)

    detected = {
        "people": people[:15],
        "projects": projects[:10],
        "topics": [],
        "uncertain": uncertain[:8],
    }

    return _apply_corpus_origin(detected, corpus_origin)


def _apply_corpus_origin(detected: dict, corpus_origin: dict | None) -> dict:
    """Reclassify per-candidate buckets using corpus-origin context.

    When the corpus is identified as AI-dialogue with known agent persona
    names, a candidate whose name case-insensitively matches one of those
    personas is moved from ``people``/``uncertain`` into an
    ``agent_personas`` bucket. The candidate's per-entity ``type`` is also
    rewritten to ``"agent_persona"``.

    No-op when ``corpus_origin`` is ``None`` or contains no usable persona
    names. Pure: returns a new dict, does not mutate the input.
    """
    if not corpus_origin:
        return detected

    origin_result = corpus_origin.get("result") or {}
    raw_personas = origin_result.get("agent_persona_names") or []
    persona_lower = {n.lower() for n in raw_personas if isinstance(n, str)}
    if not persona_lower:
        return detected

    agent_personas: list = []
    new_people: list = []
    new_uncertain: list = []

    for entity in detected.get("people", []):
        if entity["name"].lower() in persona_lower:
            agent_personas.append(_tag_as_persona(entity))
        else:
            new_people.append(entity)

    for entity in detected.get("uncertain", []):
        if entity["name"].lower() in persona_lower:
            agent_personas.append(_tag_as_persona(entity))
        else:
            new_uncertain.append(entity)

    if not agent_personas:
        return detected

    agent_personas.sort(key=lambda x: x.get("confidence", 0), reverse=True)

    return {
        **detected,
        "people": new_people,
        "uncertain": new_uncertain,
        "agent_personas": agent_personas,
    }


def _tag_as_persona(entity: dict) -> dict:
    """Return a new entity dict tagged as agent_persona with provenance signal."""
    existing_signals = entity.get("signals", [])
    return {
        **entity,
        "type": "agent_persona",
        "confidence": max(0.95, entity.get("confidence", 0.0)),
        "signals": ["matched corpus_origin agent_persona_names"] + existing_signals[:2],
    }


# ==================== INTERACTIVE CONFIRM ====================


def _print_entity_list(entities: list, label: str):
    print(f"\n  {label}:")
    if not entities:
        print("    (none detected)")
        return
    for i, e in enumerate(entities):
        confidence_bar = "●" * int(e["confidence"] * 5) + "○" * (5 - int(e["confidence"] * 5))
        signals_str = ", ".join(e["signals"][:2]) if e["signals"] else ""
        print(f"    {i + 1:2}. {e['name']:20} [{confidence_bar}] {signals_str}")


def confirm_entities(detected: dict, yes: bool = False) -> dict:
    """
    Interactive confirmation step.
    User reviews detected entities, removes wrong ones, adds missing ones.
    Returns confirmed {people: [names], projects: [names], topics: [names]}.

    Topics are not surfaced for interactive review — they come from the
    LLM-refined ``TOPIC`` bucket and are passed through verbatim. They
    feed cross-wing tunnel computation at mine time (see
    ``palace_graph.compute_topic_tunnels``); a wrong topic at worst adds
    a low-traffic tunnel and never alters drawer storage.

    Pass yes=True to auto-accept all detected entities without prompting.
    """
    print(f"\n{'=' * 58}")
    print("  MemPalace — Entity Detection")
    print(f"{'=' * 58}")
    print("\n  Scanned your files. Here's what we found:\n")

    _print_entity_list(detected["people"], "PEOPLE")
    _print_entity_list(detected["projects"], "PROJECTS")

    if detected.get("topics"):
        _print_entity_list(detected["topics"], "TOPICS (cross-wing tunnel signal)")

    if detected["uncertain"]:
        _print_entity_list(detected["uncertain"], "UNCERTAIN (need your call)")

    confirmed_people = [e["name"] for e in detected["people"]]
    confirmed_projects = [e["name"] for e in detected["projects"]]
    confirmed_topics = [e["name"] for e in detected.get("topics", [])]

    if yes:
        # Auto-accept: include all detected (skip uncertain — ambiguous without user input)
        print(
            f"\n  Auto-accepting {len(confirmed_people)} people, "
            f"{len(confirmed_projects)} projects, "
            f"{len(confirmed_topics)} topics."
        )
        return {
            "people": confirmed_people,
            "projects": confirmed_projects,
            "topics": confirmed_topics,
        }

    print(f"\n{'─' * 58}")
    print("  Options:")
    print("    [enter]  Accept all")
    print("    [edit]   Remove wrong entries or reclassify uncertain")
    print("    [add]    Add missing people or projects")
    print()

    choice = input("  Your choice [enter/edit/add]: ").strip().lower()

    confirmed_people = [e["name"] for e in detected["people"]]
    confirmed_projects = [e["name"] for e in detected["projects"]]

    if choice == "edit":
        # Handle uncertain first
        if detected["uncertain"]:
            print("\n  Uncertain entities — classify each:")
            for e in detected["uncertain"]:
                ans = input(f"    {e['name']} — (p)erson, (r)project, or (s)kip? ").strip().lower()
                if ans == "p":
                    confirmed_people.append(e["name"])
                elif ans == "r":
                    confirmed_projects.append(e["name"])

        # Remove wrong people
        print(f"\n  Current people: {', '.join(confirmed_people) or '(none)'}")
        remove = input(
            "  Numbers to REMOVE from people (comma-separated, or enter to skip): "
        ).strip()
        if remove:
            to_remove = {int(x.strip()) - 1 for x in remove.split(",") if x.strip().isdigit()}
            confirmed_people = [p for i, p in enumerate(confirmed_people) if i not in to_remove]

        # Remove wrong projects
        print(f"\n  Current projects: {', '.join(confirmed_projects) or '(none)'}")
        remove = input(
            "  Numbers to REMOVE from projects (comma-separated, or enter to skip): "
        ).strip()
        if remove:
            to_remove = {int(x.strip()) - 1 for x in remove.split(",") if x.strip().isdigit()}
            confirmed_projects = [p for i, p in enumerate(confirmed_projects) if i not in to_remove]

    if choice == "add" or input("\n  Add any missing? [y/N]: ").strip().lower() == "y":
        while True:
            name = input("  Name (or enter to stop): ").strip()
            if not name:
                break
            kind = input(f"  Is '{name}' a (p)erson or p(r)oject? ").strip().lower()
            if kind == "p":
                confirmed_people.append(name)
            elif kind == "r":
                confirmed_projects.append(name)

    print(f"\n{'=' * 58}")
    print("  Confirmed:")
    print(f"  People:   {', '.join(confirmed_people) or '(none)'}")
    print(f"  Projects: {', '.join(confirmed_projects) or '(none)'}")
    if confirmed_topics:
        print(f"  Topics:   {', '.join(confirmed_topics)}")
    print(f"{'=' * 58}\n")

    return {
        "people": confirmed_people,
        "projects": confirmed_projects,
        "topics": confirmed_topics,
    }


# ==================== SCAN HELPER ====================


def scan_for_detection(project_dir: str, max_files: int = 10) -> list:
    """
    Collect prose file paths for entity detection.
    Prose only (.txt, .md, .rst, .csv) — code files produce too many false positives.
    Falls back to all readable files if no prose found.
    """
    project_path = Path(project_dir).expanduser().resolve()
    prose_files = []
    all_files = []

    for root, dirs, filenames in os.walk(project_path):
        dirs[:] = [d for d in dirs if d not in SKIP_DIRS]
        for filename in filenames:
            filepath = Path(root) / filename
            if filepath.stem.lower() in SKIP_FILENAMES:
                continue
            ext = filepath.suffix.lower()
            if ext in PROSE_EXTENSIONS:
                prose_files.append(filepath)
            elif ext in READABLE_EXTENSIONS:
                all_files.append(filepath)

    # Prefer prose files — fall back to all readable if too few prose files
    files = prose_files if len(prose_files) >= 3 else prose_files + all_files
    return files[:max_files]


# ==================== CLI ====================

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python entity_detector.py <directory> [lang1,lang2,...]")
        sys.exit(1)

    project_dir = sys.argv[1]
    langs = tuple(sys.argv[2].split(",")) if len(sys.argv) >= 3 else ("en",)
    print(f"Scanning: {project_dir} (languages: {', '.join(langs)})")
    files = scan_for_detection(project_dir)
    print(f"Reading {len(files)} files...")
    detected = detect_entities(files, languages=langs)
    confirmed = confirm_entities(detected)
    print("Confirmed entities:", confirmed)