#!/usr/bin/env python3
"""
Summarize / reformat a meeting's enhanced transcript with a LOCAL LLM.

Standalone on purpose: this is NOT wired into process_meeting.py. The pipeline
(transcribe + frames + OCR) stays fully deterministic and offline; this is the
one non-deterministic, network-*capable* step, so it gets its own entry point.
By default it talks to a local vLLM OpenAI-compatible server (no cloud), but the
--base-url swap lets you point it at a company-sanctioned endpoint instead.

The steering instruction is a first-class argument — pass any nuance you want
("focus on names and their roles", "read the closing signals", "reformat as a
decisions+action-items table, English output"). The instruction is threaded
into every stage (map, extract, reduce), not just the final synthesis, so the
per-chunk pass never discards the detail you asked to keep.

Architecture (ports the "let the architecture carry correctness" rule to the
summarization failure mode — hallucinated names/facts + long-input drift):
  1. map     — chunk the transcript, summarize each chunk under the instruction
  2. extract — emit a validated JSON of facts (participants/roles/decisions/...)
  3. reduce  — write the final output from the validated facts + chunk notes

Usage:
    # start the local server first (`~/wdir/llm/serve.sh qwen7b`), then:
    python ctrl/summarize/summarize_meeting.py output/<run>/<stem>_enhanced.txt \\
        "focus on the names mentioned and their roles, output in English" \\
        -o output/<run>/summary_en.md

Run it under the venv that has the `openai` client (e.g. ~/wdir/llm/.venv).
"""
import argparse
import json
import re
import sys
from pathlib import Path


DEFAULT_BASE_URL = "http://localhost:11000/v1"
DEFAULT_MODEL = "Qwen/Qwen2.5-7B-Instruct-AWQ"

# Rough char->token heuristic so we don't need a tokenizer dependency. Mixed
# ES/EN prose lands around ~3.6 chars/token; 4.0 keeps us conservative (we
# under-fill rather than overflow the context window).
CHARS_PER_TOKEN = 4.0

# Shared rules injected into every stage. This is the hallucination guard: the
# model condenses what is present, it never invents — especially names/roles.
GROUNDING_RULES = """\
Rules you must follow:
- Be faithful to the transcript. Never invent names, roles, numbers, or facts.
- If something is unclear or not stated, say so — do not guess.
- Preserve proper nouns exactly as written (people, companies, tools).
- When you state a concrete claim, anchor it to its [mm:ss] timestamp.
- The transcript is machine-generated and may contain ASR errors; prefer the
  most consistent reading across the whole transcript over any single garbled
  line, and flag a name/term you are unsure about rather than normalizing it
  silently.
- Pay special attention to acronyms and product/tool names: a slightly-off
  acronym or a homophone that makes no sense in context is almost certainly an
  ASR mis-hearing — flag it as "(heard: X — likely Y?)" using context to infer
  the intended term, rather than repeating the nonsensical form."""

MAP_SYSTEM = """\
You are condensing ONE chunk of a longer meeting transcript.

The user's instruction for the final output is:
<instruction>
{instruction}
</instruction>

Produce dense, factual notes for THIS chunk that preserve everything relevant to
that instruction. Keep every name, role, decision, date, number, and notable
quote with its [mm:ss]. Do not write a polished summary yet — these notes are
raw material for a later synthesis pass, so keep detail over readability.

{rules}"""

EXTRACT_SYSTEM = """\
Extract structured facts from the meeting notes/transcript as STRICT JSON only —
no prose, no markdown fences. Use exactly this schema; use [] or "" when unknown:

{{
  "participants": [{{"name": "", "role": "", "org": "", "evidence_ts": ""}}],
  "people_mentioned": [{{"name": "", "role": "", "org": "", "evidence_ts": ""}}],
  "orgs": [{{"name": "", "what": ""}}],
  "decisions": [{{"decision": "", "ts": ""}}],
  "action_items": [{{"item": "", "owner": "", "due": "", "ts": ""}}],
  "dates": [{{"what": "", "when": "", "ts": ""}}],
  "key_quotes": [{{"speaker": "", "ts": "", "quote": ""}}],
  "open_questions": [""]
}}

Only include entries actually supported by the text. Distinguish participants
(present on the call) from people merely mentioned. This JSON is the source of
truth for names/roles in the final output, so be precise and do not invent.

{rules}"""

REDUCE_SYSTEM = """\
You are writing the FINAL output of a meeting from (a) the user's instruction,
(b) a validated JSON of facts, and (c) per-chunk notes.

The user's instruction is authoritative — follow it for focus, structure, and
language:
<instruction>
{instruction}
</instruction>

Use the validated facts JSON as the source of truth for all names, roles, and
dates (the notes may contain ASR noise; the JSON has been checked). Write only
what the instruction asks for. Output clean Markdown.

{rules}"""


def estimate_tokens(text: str) -> int:
    return int(len(text) / CHARS_PER_TOKEN)


def default_output(transcript: Path, kind: str) -> Path:
    """Write next to the transcript, in the same run folder, following the
    pipeline's <stem>_<kind> naming (e.g. keneth_aponte_summary.md)."""
    stem = transcript.stem
    if stem.endswith("_enhanced"):
        stem = stem[: -len("_enhanced")]
    return transcript.parent / f"{stem}_{kind}.md"


def load_transcript(path: Path, keep_frames: bool) -> str:
    text = path.read_text()
    if keep_frames:
        return text
    # Drop the "Frame: <path>.jpg" noise lines and their "SCREEN CONTENT:"
    # headers — a text-only model can't use a file path, and they waste tokens.
    # Real OCR text (if any) does not match these patterns and is kept.
    lines = text.splitlines()
    out = []
    skip_next_blank = False
    for line in lines:
        if re.match(r"\s*\[\d+:\d+\]\s+SCREEN CONTENT:\s*$", line):
            skip_next_blank = True
            continue
        if re.match(r"\s*Frame:\s+.*\.(jpg|jpeg|png)\s*$", line):
            continue
        out.append(line)
    return "\n".join(out)


def chunk_transcript(text: str, chunk_tokens: int) -> list:
    """Split on blank-line block boundaries, packing blocks up to chunk_tokens
    so we never cut a speaker turn in half."""
    blocks = re.split(r"\n\s*\n", text)
    chunks, cur, cur_tok = [], [], 0
    for block in blocks:
        bt = estimate_tokens(block)
        if cur and cur_tok + bt > chunk_tokens:
            chunks.append("\n\n".join(cur))
            cur, cur_tok = [], 0
        cur.append(block)
        cur_tok += bt
    if cur:
        chunks.append("\n\n".join(cur))
    return chunks


def make_client(base_url: str, api_key: str):
    try:
        from openai import OpenAI
    except ImportError:
        sys.exit(
            "ERROR: the `openai` client is not installed in this interpreter.\n"
            "Run this under the venv that has it, e.g.:\n"
            "  ~/wdir/llm/.venv/bin/python summarize_meeting.py ..."
        )
    return OpenAI(base_url=base_url, api_key=api_key)


def call(client, model, system, user, temperature, max_tokens):
    resp = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return resp.choices[0].message.content.strip()


def parse_json_lenient(raw: str):
    """vLLM models sometimes wrap JSON in ``` fences or add a stray prefix."""
    raw = re.sub(r"^```(?:json)?\s*|\s*```$", "", raw.strip(), flags=re.MULTILINE)
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        m = re.search(r"\{.*\}", raw, re.DOTALL)
        if m:
            try:
                return json.loads(m.group(0))
            except json.JSONDecodeError:
                pass
    return None


def summarize(client, args, transcript: str) -> str:
    rules = GROUNDING_RULES
    instr = args.instruction
    budget = args.chunk_tokens

    def log(msg):
        if not args.quiet:
            print(f"[summarize] {msg}", file=sys.stderr)

    total_tok = estimate_tokens(transcript)
    single_pass = args.no_map_reduce or total_tok <= budget

    # --- map ---------------------------------------------------------------
    if single_pass:
        log(f"single-pass (~{total_tok} tok <= chunk budget {budget})")
        notes = transcript
    else:
        chunks = chunk_transcript(transcript, budget)
        log(f"map: {len(chunks)} chunks (~{total_tok} tok total)")
        chunk_notes = []
        for i, ch in enumerate(chunks, 1):
            log(f"  map chunk {i}/{len(chunks)}")
            note = call(
                client, args.model,
                MAP_SYSTEM.format(instruction=instr, rules=rules),
                ch, args.temperature, args.max_tokens,
            )
            chunk_notes.append(f"### Chunk {i} notes\n{note}")
        notes = "\n\n".join(chunk_notes)

    # --- extract -----------------------------------------------------------
    facts_json = "{}"
    if not args.no_extract:
        log("extract: pulling structured facts")
        raw = call(
            client, args.model,
            EXTRACT_SYSTEM.format(rules=rules),
            notes if single_pass else notes + "\n\n" + transcript[: budget * 4],
            0.0, args.max_tokens,
        )
        facts = parse_json_lenient(raw)
        if facts is None:
            log("  WARNING: extraction did not return valid JSON; continuing without it")
        else:
            facts_json = json.dumps(facts, ensure_ascii=False, indent=2)
        if args.extract_only:
            return facts_json

    # --- reduce ------------------------------------------------------------
    log("reduce: writing final output")
    user = (
        f"VALIDATED FACTS (source of truth for names/roles/dates):\n{facts_json}\n\n"
        f"NOTES / TRANSCRIPT:\n{notes}"
    )
    return call(
        client, args.model,
        REDUCE_SYSTEM.format(instruction=instr, rules=rules),
        user, args.temperature, args.max_tokens,
    )


def main():
    p = argparse.ArgumentParser(
        description="Summarize/reformat a meeting enhanced-transcript with a local LLM.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("transcript", type=Path, help="path to *_enhanced.txt (or any text)")
    p.add_argument(
        "instruction", nargs="?", default="Summarize this meeting clearly.",
        help='steering instruction, e.g. "focus on names and roles, English output"',
    )
    p.add_argument("-o", "--output", type=Path, help="write here (default: <run>/<stem>_summary.md next to the transcript)")
    p.add_argument("--stdout", action="store_true", help="print to stdout instead of writing a file")
    p.add_argument("--base-url", default=DEFAULT_BASE_URL, help=f"OpenAI-compatible endpoint (default: {DEFAULT_BASE_URL})")
    p.add_argument("--model", default=DEFAULT_MODEL, help=f"model id (default: {DEFAULT_MODEL})")
    p.add_argument("--api-key", default="local", help="ignored by vLLM; set for a real provider")
    p.add_argument("--chunk-tokens", type=int, default=6000, help="map-reduce chunk budget (default: 6000)")
    p.add_argument("--max-tokens", type=int, default=4096, help="generation cap per call (default: 4096)")
    p.add_argument("--temperature", type=float, default=0.2, help="sampling temperature (default: 0.2)")
    p.add_argument("--no-map-reduce", action="store_true", help="force single-pass (short transcripts)")
    p.add_argument("--no-extract", action="store_true", help="skip the structured-facts grounding pass")
    p.add_argument("--extract-only", action="store_true", help="print only the extracted JSON facts and exit")
    p.add_argument("--keep-frames", action="store_true", help="keep 'Frame: <path>' lines (default: strip them)")
    p.add_argument("-q", "--quiet", action="store_true", help="suppress progress on stderr")
    args = p.parse_args()

    if not args.transcript.is_file():
        sys.exit(f"ERROR: transcript not found: {args.transcript}")

    transcript = load_transcript(args.transcript, args.keep_frames)
    if not transcript.strip():
        sys.exit("ERROR: transcript is empty after loading.")

    client = make_client(args.base_url, args.api_key)
    result = summarize(client, args, transcript)

    if args.stdout:
        print(result)
    else:
        out = args.output or default_output(args.transcript, "summary")
        out.write_text(result + "\n")
        if not args.quiet:
            print(f"[summarize] wrote {out}", file=sys.stderr)


if __name__ == "__main__":
    main()