"""
Transcript correction — the second opinion on what was said.

Whisper hears a P25 vocoder through a narrowband channel and guesses at proper
nouns it has no reason to know: street names, business names, unit call signs.
It guesses confidently, so the output reads like speech and is wrong in exactly
the places that matter downstream — "Cool Parts, Illinois" and "Shout out to
Optum" both became incident locations.

Correction used to be a line in intelligence.py's EXTRACTION_PROMPT, which put
it in the wrong place twice over (server-26#36): the same model call that
extracted units, location and severity emitted the correction *afterwards*, so
extraction reasoned over uncorrected text; and it sat behind
`correlation_enabled`, so during a cost-controlled STT-only window nothing was
ever corrected at all. It belongs here, between transcription and everything
that consumes a transcript.

WHY A SEPARATE PASS AND NOT A WHISPER PROMPT: Whisper treats its prompt as
preceding transcript text and will happily continue a pattern it finds there —
an enumerated ten-code prompt made it emit "10-4. 10-5. 10-6. …" over silence
(see transcription.py). Vocabulary can never be a transcription prior. A
corrector that receives an already-produced transcript plus a reference list has
no series to extend; it can only substitute what it was given.

SCOPE RESOLUTION: reference data is merged from the talkgroup and the system,
**talkgroup first**. The specific beats the general — a system spanning several
counties may have one talkgroup covering a single municipality, and that
municipality's streets must not be buried under a county-wide list. A
single-municipality system is the degenerate case: populate the system level and
every talkgroup inherits it.
"""

import asyncio
import json
from typing import Any, Optional

from app.config import settings
from app.internal import area_context
from app.internal import firestore as fstore
from app.internal import place_verifier
from app.internal.logger import logger

# A transcript this short has no proper nouns to get wrong — "10-4.", "6-2,
# stand by." — and 9 of 29 calls in the 2026-08-23 sample sat at or under this.
# Skipping them is most of the cost saving for none of the value.
MIN_WORDS_FOR_CORRECTION = 4

_PROMPT = """You are correcting a police/fire radio transcript produced by an automatic speech recogniser.

The recogniser hears a low-bitrate vocoded radio channel. It reliably mishears proper nouns — street names, business names, town names, unit call signs — and substitutes common words that sound similar. Your job is to put back what was almost certainly said.

{context_block}
Rules:
- Change ONLY what is likely a mishearing. If a phrase is already plausible radio traffic, leave it exactly as it is.
- Prefer a name from the reference lists above when the transcript contains something that sounds like it. That is the entire point of this pass.
- NEVER add information. No new sentences, no invented units, no addresses that are not implied by the audio's own words.
- Keep radio language as radio language. Do NOT expand ten-codes or signals into plain English: "10-4" stays "10-4".
- Keep the speaker's structure and order. This is not a rewrite or a summary.
- If the text is clearly not speech at all — a counting run like "10-11. 10-12. 10-13.", or one phrase repeating many times over static — set not_speech to true.

Return JSON:
  corrected:  the corrected transcript, or null if nothing needed changing
  segments:   REQUIRED when numbered transmissions are given below — the corrected
              text for each one, as an array of exactly the same length and order.
              Never merge, split, reorder or drop a transmission; an unchanged one
              is returned verbatim. Omit this field entirely when no transmissions
              are numbered.
  not_speech: true if this is recogniser noise rather than a transmission
  changed:    list of ["heard" -> "corrected"] pairs you applied, for audit
  locations:  every place name in your corrected output, exactly as it appears
              there — streets, intersections, businesses, schools, towns,
              landmarks. Include ones you are unsure of; that is the point.
              A unit call sign or a person's name is NOT a location.

{transcript}"""


def _render_input(text: str, segments: Optional[list[dict]]) -> str:
    """Numbered transmissions when we have them, so corrections stay aligned."""
    if segments and len(segments) > 1:
        lines = [f"{i + 1}. {s.get('text', '')}" for i, s in enumerate(segments)]
        body = "\n".join(lines)
        return f"Transmissions ({len(segments)}):\n{body}"
    return f"Transcript:\n{text}"


def _dedupe(items: list[str]) -> list[str]:
    """Preserve order, drop case-insensitive duplicates."""
    seen: set[str] = set()
    out: list[str] = []
    for item in items:
        key = (item or "").strip().lower()
        if key and key not in seen:
            seen.add(key)
            out.append(item.strip())
    return out


def _talkgroup_entry(system_doc: dict, talkgroup_id: Optional[int]) -> dict:
    """The config.talkgroups[] entry for this talkgroup, or {}."""
    if talkgroup_id is None:
        return {}
    try:
        wanted = int(talkgroup_id)
    except (TypeError, ValueError):
        return {}
    for tg in (system_doc.get("config") or {}).get("talkgroups", []) or []:
        try:
            if int(tg.get("id", -1)) == wanted:
                return tg
        except (TypeError, ValueError):
            continue
    return {}


def _area_lines(area: dict) -> list[str]:
    """
    Render a merged area_context as prompt lines. Empty when nothing is set.

    One block, not one per scope: by the time this runs the two scopes have
    already been merged with talkgroup ahead of system, and showing the model
    two competing lists invites it to pick from the wrong one.
    """
    if not area:
        return []
    lines: list[str] = []
    place = ", ".join(
        str(area[f]) for f in area_context.PLACE_FIELDS if area.get(f)
    )
    if place:
        lines.append(f"Area covered by this channel: {place}")
    knowledge = area.get("local_knowledge") or []
    if knowledge:
        lines.append("Local names heard on this channel:")
        lines.extend(
            f"  {e['term']} — {e['meaning']}" if e.get("meaning") else f"  {e['term']}"
            for e in knowledge
        )
    return lines


async def resolve_context(system_id: Optional[str], talkgroup_id: Optional[int]) -> dict:
    """
    Merge the reference data a corrector needs, talkgroup ahead of system.

    Returns {"vocabulary", "ten_codes", "area_lines", "area", "system_area",
    "tg_area"}. Empty everywhere is legitimate — a system nobody has configured
    yet. The two raw scopes come back alongside the merge because the place
    verifier needs them to pick an anchor (server-26#37).
    """
    empty: dict[str, Any] = {
        "vocabulary": [], "ten_codes": {}, "area_lines": [],
        "area": {}, "system_area": {}, "tg_area": {},
    }
    if not system_id:
        return empty

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

    tg = _talkgroup_entry(system_doc, talkgroup_id)

    # Talkgroup terms first so they survive any downstream truncation.
    vocabulary = _dedupe(
        list(tg.get("vocabulary") or []) + list(system_doc.get("vocabulary") or [])
    )

    # Ten-codes: system-wide reference, with talkgroup entries overriding a
    # code that means something different on this channel.
    ten_codes = dict(system_doc.get("ten_codes") or {})
    ten_codes.update(tg.get("ten_codes") or {})

    system_area = system_doc.get("area_context") or {}
    tg_area = tg.get("area_context") or {}
    area = area_context.effective(system_area, tg_area)

    return {
        "vocabulary": vocabulary,
        "ten_codes": ten_codes,
        "area_lines": _area_lines(area),
        "area": area,
        "system_area": system_area,
        "tg_area": tg_area,
    }


def build_context_block(context: dict, talkgroup_name: Optional[str]) -> str:
    """Render resolved context into the prompt's reference section."""
    lines: list[str] = []
    if talkgroup_name:
        lines.append(f"Channel: {talkgroup_name}")
    lines.extend(context.get("area_lines") or [])
    vocabulary = context.get("vocabulary") or []
    if vocabulary:
        lines.append("Known local names and terms: " + ", ".join(vocabulary))
    ten_codes = context.get("ten_codes") or {}
    if ten_codes:
        rendered = ", ".join(f"{code}={meaning}" for code, meaning in sorted(ten_codes.items()))
        lines.append(f"Ten-codes used on this system: {rendered}")
    return ("\n".join(lines) + "\n") if lines else ""


def _sync_gemini(model_name: str, prompt: str) -> dict:
    import google.generativeai as genai  # lazy import — only when needed

    genai.configure(api_key=settings.gemini_api_key)
    model = genai.GenerativeModel(
        model_name,
        generation_config={"response_mime_type": "application/json"},
    )
    return json.loads(model.generate_content(prompt).text)


async def correct(
    call_id: str,
    text: str,
    segments: Optional[list[dict]] = None,
    system_id: Optional[str] = None,
    talkgroup_id: Optional[int] = None,
    talkgroup_name: Optional[str] = None,
) -> tuple[Optional[str], Optional[list[dict]], bool]:
    """
    Second-opinion pass over a transcript.

    Returns (corrected_text, corrected_segments, not_speech). ``None`` for
    either correction means "no change" — the corrector found nothing to fix,
    could not run, or returned segments that did not line up. Callers keep the
    original in that case; correction is an improvement, never a dependency.

    Segments matter as much as the joined text: intelligence.py builds its
    extraction prompt from NUMBERED SEGMENTS whenever there is more than one,
    so a correction that only fixed the joined transcript would never reach the
    model on exactly the multi-transmission calls that carry the most content.
    """
    if not settings.gemini_api_key or not settings.transcript_correction_enabled:
        return None, None, False

    if len((text or "").split()) < MIN_WORDS_FOR_CORRECTION:
        return None, None, False

    context = await resolve_context(system_id, talkgroup_id)
    prompt = _PROMPT.format(
        context_block=build_context_block(context, talkgroup_name),
        transcript=_render_input(text, segments),
    )

    try:
        raw = await asyncio.to_thread(
            _sync_gemini, settings.transcript_correction_model, prompt
        )
    except Exception as e:
        # Never fail the transcript over a failed correction — the raw text is
        # still worth having. ai_health reporting is the caller's business.
        logger.warning(f"Transcript correction failed for call {call_id}: {e}")
        return None, None, False

    not_speech = bool(raw.get("not_speech"))

    corrected = raw.get("corrected")
    if not isinstance(corrected, str) or not corrected.strip():
        corrected = None
    elif corrected.strip() == (text or "").strip():
        corrected = None

    # Segment alignment is non-negotiable: scene extraction maps scenes back to
    # transmissions by INDEX (segment_indices), so a returned array of the wrong
    # length would silently attribute the wrong audio to a scene. Wrong length,
    # wrong type, or any non-string entry and the segments are discarded whole —
    # the joined correction still stands.
    corrected_segments: Optional[list[dict]] = None
    if segments and len(segments) > 1:
        returned = raw.get("segments")
        if (
            isinstance(returned, list)
            and len(returned) == len(segments)
            and all(isinstance(x, str) for x in returned)
        ):
            corrected_segments = [
                {**seg, "text": new.strip() or seg.get("text", "")}
                for seg, new in zip(segments, returned)
            ]
            if all(s["text"] == o.get("text") for s, o in zip(corrected_segments, segments)):
                corrected_segments = None
        elif returned is not None:
            logger.warning(
                f"Transcript correction for call {call_id} returned "
                f"{len(returned) if isinstance(returned, list) else type(returned).__name__} "
                f"segment(s) against {len(segments)} — discarding segment corrections"
            )

    # Maps has the last word on place names (server-26#37). The corrector can
    # only match against the list it was handed, so a plausible-sounding invention
    # — "Cool Parts, Illinois" — reads exactly like a real street to it. The
    # verifier geocodes each location noun against the talkgroup's anchor and,
    # on a miss, looks for a sound-alike that does resolve there. It runs on the
    # corrected copy so it judges the text everything downstream will actually
    # read, and it skips entirely when there is no discriminating anchor.
    if not not_speech:
        locations = [x for x in (raw.get("locations") or []) if isinstance(x, str)]
        try:
            verified_text, verified_segments = await place_verifier.verify(
                call_id,
                corrected or text,
                corrected_segments or segments,
                locations,
                context.get("system_area"),
                context.get("tg_area"),
                system_id=system_id,
                talkgroup_id=talkgroup_id,
            )
        except Exception as e:
            logger.warning(f"Place verification failed for call {call_id}: {e}")
            verified_text, verified_segments = None, None
        if verified_text:
            corrected = verified_text
        if verified_segments:
            corrected_segments = verified_segments

    if corrected or corrected_segments or not_speech:
        changed = raw.get("changed") or []
        logger.info(
            f"Transcript correction ({settings.transcript_correction_model}): call {call_id} "
            f"not_speech={not_speech} segments={'yes' if corrected_segments else 'no'} "
            f"changes={changed if isinstance(changed, list) else '?'}"
        )
    return corrected, corrected_segments, not_speech