"""
Per-system vocabulary learning for STT accuracy improvement.

Three mechanisms:
1. Bootstrap  — one-shot GPT-4o call generates local knowledge at system setup:
                agencies + abbreviations, unit naming, streets, acronyms.
2. Correction — diffs admin transcript edits, extracts corrected tokens → vocabulary.
3. Induction  — background loop samples N tokens of transcripts per system,
                asks GPT-4o-mini to propose new terms → queued as pending for review.

Firestore schema additions on system documents:
  vocabulary:            list[str]   — approved terms; injected into Whisper + GPT prompts
  vocabulary_pending:    list[dict]  — induction suggestions awaiting admin review
                                       each: {term, source, added_at}
  vocabulary_bootstrapped: bool      — bootstrap has been run at least once
"""
import asyncio
import difflib
import json
import random
import re
from datetime import datetime, timezone, timedelta
from typing import Any, Optional
from app.internal.logger import logger
from app.internal import area_context
from app.internal import firestore as fstore
from app.config import settings


# ─────────────────────────────────────────────────────────────────────────────
# Prompt templates
# ─────────────────────────────────────────────────────────────────────────────

_BOOTSTRAP_PROMPT = """\
You are building a radio vocabulary dictionary to improve speech-to-text accuracy for a P25 \
public-safety radio monitoring system in a specific area. The STT model has no local knowledge, \
so common terms like "YVAC" get transcribed as "why vac", "5-baker" as "5 acre", etc.

System name: {system_name}
System type: {system_type}
Area context: {area_hint}

Return ONLY a JSON object:
{{"vocabulary": [list of strings]}}

Include terms you are confident about for this area:
- Agency names and their radio abbreviations (e.g. "YVAC" = Yorktown Volunteer Ambulance Corps)
- Unit ID examples using the local naming convention (e.g. "5-baker", "5-charlie", "1-david";
  many departments use APCO phonetics: adam, baker, charles, david, edward, frank, george,
  henry, ida, john, king, lincoln, mary, nora, ocean, paul, queen, robert, sam, tom, union,
  victor, william, x-ray, young, zebra)
- Major routes, roads, and key intersections
- Local landmarks and geographic references dispatchers use
- Agency-specific codes that differ from standard APCO

Return a flat list of strings — abbreviations, proper names, unit IDs, street names.
Do NOT include common English words. Max 80 terms. Only include what you are confident is \
accurate for this specific area; return fewer terms rather than guessing."""

_INDUCTION_PROMPT = """\
You are analyzing P25 emergency radio transcripts from ONE talkgroup (a single radio channel) \
to find local terms that should be added to improve future speech-to-text accuracy for that \
channel.

System: {system_name}
Channel: {talkgroup_name}
Area: {area_hint}
Terms this channel already knows (do not re-propose these): {existing_vocab}

Sampled transcripts:
{transcript_block}

Find terms that are LIKELY STT errors or local terms missing from the list:
- Unit IDs that appear garbled (e.g. "5 acre" → "5-baker")
- Agency acronyms spelled out phonetically (e.g. "why vac" → "YVAC")
- Street names or locations that look misspelled or oddly transcribed
- Callsigns or local codes not yet known

Return ONLY a JSON object:
{{"new_terms": [{{"term": "YVAC", "meaning": "Yorktown Volunteer Ambulance Corps"}}, ...]}}

`meaning` is what the term refers to — an agency, a road, a unit type. Omit it or use null \
when you genuinely do not know; a term with no meaning is still worth proposing.

Only propose what is specific to THIS channel and this area. Do not propose a term just \
because it appears often. Return {{"new_terms": []}} if nothing new is found."""


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

async def bootstrap_system_vocabulary(system_id: str) -> list[str]:
    """
    One-shot GPT-4o bootstrap: generate local-knowledge vocabulary for a system.
    Merges generated terms into system.vocabulary and sets vocabulary_bootstrapped=True.
    Returns the list of newly generated terms.
    """
    system_doc = await fstore.doc_get("systems", system_id)
    if not system_doc:
        logger.warning(f"Vocabulary bootstrap: system {system_id} not found")
        return []

    system_name = system_doc.get("name", "Unknown")
    system_type = system_doc.get("type", "P25")

    # Prefer the place an operator actually set. Guessing the area from talkgroup
    # names is thin for a single-municipality system and close to useless for a
    # multi-county one (server-26#36), so it is only the fallback now.
    area = system_doc.get("area_context") or {}
    place = ", ".join(str(area[f]) for f in area_context.PLACE_FIELDS if area.get(f))
    if place:
        area_hint = place
    else:
        talkgroups = system_doc.get("config", {}).get("talkgroups", [])
        tg_names = [tg.get("name", "") for tg in talkgroups if tg.get("name")][:8]
        area_hint = f"Talkgroups include: {', '.join(tg_names)}" if tg_names else "Unknown area"

    terms = await asyncio.to_thread(_sync_bootstrap, system_name, system_type, area_hint)
    if not terms:
        return []

    existing = system_doc.get("vocabulary") or []
    existing_lower = {t.lower() for t in existing}
    to_add = [t for t in terms if t.lower() not in existing_lower]
    merged = list(dict.fromkeys(existing + to_add))

    await fstore.doc_set("systems", system_id, {
        "vocabulary": merged,
        "vocabulary_bootstrapped": True,
    })
    logger.info(
        f"Vocabulary bootstrap: {len(to_add)} term(s) generated for system {system_id} "
        f"({system_name})"
    )
    return to_add


async def learn_from_correction(system_id: str, original: str, corrected: str) -> None:
    """
    Diff original and corrected transcripts; append new tokens to the approved vocabulary.
    Called automatically when an admin saves a transcript correction.
    """
    if not system_id or not original or not corrected:
        return

    new_terms = _diff_new_terms(original, corrected)
    if not new_terms:
        return

    system_doc = await fstore.doc_get("systems", system_id)
    if not system_doc:
        return

    existing = system_doc.get("vocabulary") or []
    existing_lower = {t.lower() for t in existing}
    to_add = [t for t in new_terms if t.lower() not in existing_lower]
    if not to_add:
        return

    merged = list(dict.fromkeys(existing + to_add))
    await fstore.doc_set("systems", system_id, {"vocabulary": merged})
    logger.info(
        f"Vocabulary: learned {len(to_add)} term(s) from correction on system {system_id}: "
        f"{to_add}"
    )


async def approve_pending_term(system_id: str, term: str) -> None:
    """Move a pending term into the approved vocabulary."""
    system_doc = await fstore.doc_get("systems", system_id)
    if not system_doc:
        return
    pending = [p for p in (system_doc.get("vocabulary_pending") or []) if p["term"] != term]
    vocab = system_doc.get("vocabulary") or []
    if term.lower() not in {t.lower() for t in vocab}:
        vocab = list(dict.fromkeys(vocab + [term]))
    await fstore.doc_set("systems", system_id, {
        "vocabulary": vocab,
        "vocabulary_pending": pending,
    })


async def dismiss_pending_term(system_id: str, term: str) -> None:
    """Remove a pending term without adding it to vocabulary."""
    system_doc = await fstore.doc_get("systems", system_id)
    if not system_doc:
        return
    pending = [p for p in (system_doc.get("vocabulary_pending") or []) if p["term"] != term]
    await fstore.doc_set("systems", system_id, {"vocabulary_pending": pending})


async def add_term(system_id: str, term: str) -> None:
    """Manually add a term to the approved vocabulary."""
    system_doc = await fstore.doc_get("systems", system_id)
    if not system_doc:
        return
    vocab = system_doc.get("vocabulary") or []
    if term.lower() not in {t.lower() for t in vocab}:
        vocab = list(dict.fromkeys(vocab + [term.strip()]))
        await fstore.doc_set("systems", system_id, {"vocabulary": vocab})


async def remove_term(system_id: str, term: str) -> None:
    """Remove a term from the approved vocabulary."""
    system_doc = await fstore.doc_get("systems", system_id)
    if not system_doc:
        return
    vocab = [t for t in (system_doc.get("vocabulary") or []) if t.lower() != term.lower()]
    await fstore.doc_set("systems", system_id, {"vocabulary": vocab})


async def get_vocabulary(system_id: str) -> dict:
    """Return vocabulary and pending terms for a system (TTL-cached, 5 min)."""
    doc = await fstore.doc_get_cached("systems", system_id)
    if not doc:
        return {"vocabulary": [], "vocabulary_pending": [], "vocabulary_bootstrapped": False}
    return {
        "vocabulary":             doc.get("vocabulary") or [],
        "vocabulary_pending":     doc.get("vocabulary_pending") or [],
        "vocabulary_bootstrapped": doc.get("vocabulary_bootstrapped", False),
    }


# ─────────────────────────────────────────────────────────────────────────────
# Prompt-injection helpers (called by transcription.py and intelligence.py)
# ─────────────────────────────────────────────────────────────────────────────

def build_whisper_vocab_prompt(vocabulary: list[str]) -> str:
    """
    Format vocabulary for Whisper prompt injection.
    Whisper's prompt field acts as a context prior with a ~224-token limit.
    The base _WHISPER_PROMPT uses ~70 tokens; we budget ~150 tokens (≈550 chars) here.
    """
    if not vocabulary:
        return ""
    char_budget = 550
    terms: list[str] = []
    used = 0
    for term in vocabulary:
        cost = len(term) + 2  # ", "
        if used + cost > char_budget:
            break
        terms.append(term)
        used += cost
    return ", ".join(terms) + ". " if terms else ""


def build_gpt_vocab_block(vocabulary: list[str]) -> str:
    """Format vocabulary for injection into GPT extraction prompts."""
    if not vocabulary:
        return ""
    return f"Known local terms: {', '.join(vocabulary)}\n"


# ─────────────────────────────────────────────────────────────────────────────
# Background induction loop
# ─────────────────────────────────────────────────────────────────────────────

async def vocabulary_induction_loop() -> None:
    from app.internal.feature_flags import get_flags
    interval = settings.vocabulary_induction_interval_hours * 3600
    logger.info(
        f"Vocabulary induction loop started — "
        f"interval: {settings.vocabulary_induction_interval_hours}h, "
        f"sample budget: {settings.vocabulary_induction_sample_tokens} tokens"
    )
    await asyncio.sleep(30)  # short startup grace period before first pass
    while True:
        try:
            flags = await get_flags()
            if flags["vocabulary_learning_enabled"]:
                await _run_induction_pass()
            else:
                logger.info("Vocabulary learning disabled — skipping induction pass")
        except Exception as e:
            logger.error(f"Vocabulary induction pass failed: {e}")
        await asyncio.sleep(interval)


async def _run_induction_pass() -> None:
    systems = await fstore.collection_list("systems")
    if not systems:
        return
    logger.info(f"Vocabulary induction: processing {len(systems)} system(s)")
    for system in systems:
        system_id = system.get("system_id")
        if system_id:
            try:
                await _induct_system(system_id, system)
            except Exception as e:
                logger.warning(f"Induction failed for system {system_id}: {e}")


async def _induct_system(system_id: str, system_doc: dict) -> None:
    """
    Sample recent transcripts per TALKGROUP and propose local knowledge there.

    Proposals used to land at system level, which is the wrong blast radius
    (server-26#37). A wrong term on a talkgroup misleads one channel; the same
    term at system level misleads every channel on that system — including one
    400km away on a statewide system, which is exactly the context poisoning the
    scope rule exists to prevent. If a term really does apply system-wide,
    carrying it on several talkgroups costs almost nothing, while auto-promoting
    a wrong one is expensive to notice. So: talkgroup-level pending terms only,
    and nothing here ever promotes upward or approves itself.
    """
    system_name = system_doc.get("name", "Unknown")
    system_area = system_doc.get("area_context") or {}

    # Fetch calls from the last 7 days only — avoids scanning the entire history.
    # Active calls have ended_at=None and are excluded by the range filter automatically.
    # Needs a composite index on (system_id ASC, ended_at ASC).
    cutoff = datetime.now(timezone.utc) - timedelta(days=7)
    all_calls = await fstore.collection_where("calls", [
        ("system_id", "==", system_id),
        ("ended_at",  ">=", cutoff),
    ])
    if not all_calls:
        return

    by_tg: dict[Any, list[dict]] = {}
    for call in all_calls:
        tgid = call.get("talkgroup_id")
        if tgid is None:
            continue
        by_tg.setdefault(tgid, []).append(call)

    # The sample budget is per system, split across the talkgroups that have
    # traffic — a channel with 400 calls should not starve one with 12.
    char_budget = max(
        (settings.vocabulary_induction_sample_tokens * 4) // max(len(by_tg), 1), 800
    )

    for talkgroup_id, calls in by_tg.items():
        try:
            await _induct_talkgroup(
                system_id, system_doc, system_name, system_area,
                talkgroup_id, calls, char_budget,
            )
        except Exception as e:
            logger.warning(
                f"Induction failed for talkgroup {talkgroup_id} on system {system_id}: {e}"
            )


async def _induct_talkgroup(
    system_id: str,
    system_doc: dict,
    system_name: str,
    system_area: dict,
    talkgroup_id: Any,
    calls: list[dict],
    char_budget: int,
) -> None:
    tg_entry = area_context.talkgroup_entry(system_doc, talkgroup_id)
    tg_area = tg_entry.get("area_context") or {}
    area = area_context.effective(system_area, tg_area)

    talkgroup_name = (
        tg_entry.get("name")
        or calls[0].get("talkgroup_name")
        or f"TGID {talkgroup_id}"
    )
    known = area_context._known_terms(tg_entry, system_doc)

    random.shuffle(calls)
    transcript_block = ""
    sampled_call_docs: list[dict] = []
    for call in calls:
        text = call.get("transcript_corrected") or call.get("transcript") or ""
        if not text:
            continue
        if len(transcript_block) + len(text) > char_budget:
            break
        transcript_block += f"{text}\n"
        sampled_call_docs.append(call)

    if len(sampled_call_docs) < 3:
        return  # not enough data on this channel to learn from yet

    place = ", ".join(str(area[f]) for f in area_context.PLACE_FIELDS if area.get(f))
    proposed = await asyncio.to_thread(
        _sync_induct,
        system_name, talkgroup_name, place or "not set",
        sorted(known)[:80], transcript_block,
    )
    if not proposed:
        return

    entries = [
        {
            "term": p["term"],
            "meaning": p.get("meaning"),
            "source": "induction",
            "source_call_ids": _find_source_calls(p["term"], sampled_call_docs),
        }
        for p in proposed
        if p.get("term") and p["term"].lower() not in known
    ]
    if entries:
        await area_context.add_pending(system_id, talkgroup_id, entries)


# ─────────────────────────────────────────────────────────────────────────────
# Internal sync helpers
# ─────────────────────────────────────────────────────────────────────────────

def _find_source_calls(term: str, sampled_calls: list[dict], max_results: int = 3) -> list[str]:
    """
    Find which sampled calls most likely produced this induction suggestion.
    Splits the proposed term into tokens and searches call transcripts for overlap.
    Falls back to the first two sampled calls when no token match is found
    (e.g. fully garbled terms like "why vac" → "YVAC" have no word overlap).
    """
    tokens = [t.lower() for t in re.split(r"[^a-zA-Z0-9]+", term) if len(t) >= 2]
    matched: list[str] = []
    if tokens:
        for call in sampled_calls:
            call_id = call.get("call_id")
            if not call_id:
                continue
            text = (call.get("transcript_corrected") or call.get("transcript") or "").lower()
            if any(tok in text for tok in tokens):
                matched.append(call_id)
                if len(matched) >= max_results:
                    break
    if not matched:
        matched = [c["call_id"] for c in sampled_calls[:2] if c.get("call_id")]
    return matched


_STOP_WORDS = {
    "the", "and", "for", "are", "was", "were", "this", "that", "with",
    "have", "has", "had", "but", "not", "from", "they", "will", "what",
    "can", "all", "been", "one", "two", "three", "four", "five", "six",
    "you", "out", "who", "get", "her", "him", "his", "its", "our", "my",
    "via", "per", "any", "now", "got", "she", "let", "did", "may", "yes",
    "sir", "say", "see", "too", "off", "how", "put", "set", "try", "back",
    "just", "like", "into", "than", "them", "then", "some", "also", "onto",
    "went", "over", "copy", "okay", "unit", "post", "road", "lane", "going",
    "being", "doing", "there", "their", "about", "would", "could", "should",
    "route", "north", "south", "east", "west", "avenue", "street", "drive",
}


def _diff_new_terms(original: str, corrected: str) -> list[str]:
    """
    Token-level diff: find tokens in `corrected` that replaced or were inserted
    relative to `original`. These are the admin's intended spellings — good
    candidates for vocabulary.
    """
    orig_tokens = original.split()
    corr_tokens = corrected.split()

    matcher = difflib.SequenceMatcher(None,
        [t.lower() for t in orig_tokens],
        [t.lower() for t in corr_tokens],
    )
    new_terms: list[str] = []
    for tag, _i1, _i2, j1, j2 in matcher.get_opcodes():
        if tag in ("insert", "replace"):
            for tok in corr_tokens[j1:j2]:
                clean = tok.strip(".,!?;:()'\"").strip("-")
                if len(clean) >= 3 and clean.lower() not in _STOP_WORDS:
                    new_terms.append(clean)

    return list(dict.fromkeys(new_terms))


def _sync_bootstrap(system_name: str, system_type: str, area_hint: str) -> list[str]:
    from app.config import settings as cfg
    from openai import OpenAI

    if not cfg.openai_api_key:
        return []

    prompt = _BOOTSTRAP_PROMPT.format(
        system_name=system_name,
        system_type=system_type,
        area_hint=area_hint,
    )
    try:
        client = OpenAI(api_key=cfg.openai_api_key)
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            response_format={"type": "json_object"},
        )
        data = json.loads(response.choices[0].message.content)
        terms = data.get("vocabulary") or []
        return [str(t).strip() for t in terms if str(t).strip()]
    except Exception as e:
        logger.warning(f"Vocabulary bootstrap GPT call failed: {e}")
        return []


def _sync_induct(
    system_name: str,
    talkgroup_name: str,
    area_hint: str,
    existing_vocab: list[str],
    transcript_block: str,
) -> list[dict]:
    """Returns [{term, meaning}] — a bare string is still accepted from the model."""
    from app.config import settings as cfg
    from openai import OpenAI

    if not cfg.openai_api_key:
        return []

    vocab_str = ", ".join(existing_vocab[:80]) if existing_vocab else "(none yet)"
    prompt = _INDUCTION_PROMPT.format(
        system_name=system_name,
        talkgroup_name=talkgroup_name,
        area_hint=area_hint,
        existing_vocab=vocab_str,
        transcript_block=transcript_block[:8000],
    )
    try:
        client = OpenAI(api_key=cfg.openai_api_key)
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            response_format={"type": "json_object"},
        )
        data = json.loads(response.choices[0].message.content)
        return area_context.normalize_local_knowledge(data.get("new_terms") or [])
    except Exception as e:
        logger.warning(f"Vocabulary induction GPT call failed: {e}")
        return []