#!/usr/bin/env python3
"""
Compile a long meeting/training enhanced-transcript into a detailed technical
reference, using a LOCAL multimodal LLM that reads frames ON DEMAND.

This is NOT summarization — it RETAINS workflow/architecture detail and
reorganizes it out of conversation order. It uses the REFINE pattern: walk the
transcript top-to-bottom in windows, carrying a running compiled document as the
only large context. The running doc IS the memory; the raw transcript is never
held whole (which is why a 4-hour recording fits a small model).

Frames are consulted the way a human note-taker does: while reading each window,
the model decides which referenced frames it actually needs to see, and only
those images are attached (on demand) — webcam/transition frames cost nothing.

Standalone on purpose: not wired into process_meeting.py. Talks to a local
OpenAI-compatible server (vLLM or llama.cpp — see ~/wdir/llm/serve.sh); --base-url
swaps it.

Usage (start `~/wdir/llm/serve.sh qwen-vl` first, then):
    ~/wdir/llm/.venv/bin/python compile_meeting.py \\
        output/<run>/<stem>_enhanced.txt \\
        "compile every deployment/data-flow workflow and the system architecture \\
         as a technical reference; note the [mm:ss] each was shown on screen" \\
        -o output/<run>/reference.md

Frame modes:
    --frames ondemand  (default) two-step: model lists which frames it needs,
                       then only those are attached. Cheapest on vision tokens.
    --frames window    attach every frame in the current window; model uses the
                       relevant ones. Simpler, more tokens.
    --frames none      ignore frames entirely (text-only; for non-VL models / A-B).
"""
import argparse
import base64
import json
import re
import sys
from pathlib import Path

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

REFINE_SYS = """\
You are building up a DETAILED REFERENCE DOCUMENT from a meeting/training
transcript, one window at a time. The user's intent — follow it; otherwise use
your own judgment for structure, depth, ordering, and emphasis:
<instruction>
{instruction}
</instruction>

You get the document so far and the next window of transcript (sometimes with
screen frames). Fold the new material into the document and return the COMPLETE
updated document, dropping nothing important from before. Keep concrete detail —
names, numbers, steps, configs, specifics — rather than collapsing to general
ideas; this is a reference, not a recap. Stay faithful to the transcript and
don't invent. It's machine-transcribed, so use your own judgment on garbled
spots (an odd acronym is probably a mis-hearing). Beyond that, write and organize
it however reads best to you."""

TRIAGE_SYS = """\
You're compiling a detailed reference and reading this window of transcript. The
frames listed below are screenshots referenced in it (id + [mm:ss]). List the ids
of any you'd find worth actually looking at — lean toward looking whenever a frame
might carry detail the words alone don't. Reply with JSON only:
  {{"need": ["<frame-id>", ...]}}
Empty list if none seem useful."""

FRAME_RE = re.compile(r"Frame:\s+(\S+\.(?:jpg|jpeg|png))", re.IGNORECASE)
TS_RE = re.compile(r"\[(\d+):(\d+)\]")


def estimate_tokens(text):
    return int(len(text) / CHARS_PER_TOKEN)


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


def content_tokens(content):
    """Estimate input tokens of a chat content (str or multimodal list)."""
    if isinstance(content, str):
        return estimate_tokens(content)
    total = 0
    for part in content:
        if part.get("type") == "text":
            total += estimate_tokens(part.get("text", ""))
        else:  # image_url etc. — rough per-image vision-token allowance
            total += 800
    return total


def parse_windows(path, window_tokens):
    """Split the enhanced transcript into windows of ~window_tokens, packing
    blank-line-separated blocks whole. Each window keeps the frame refs that
    fall inside it: {text, frames:[{id, ts, path}]}."""
    raw = path.read_text()
    blocks = re.split(r"\n\s*\n", raw)
    windows, cur_text, cur_frames, cur_tok = [], [], [], 0

    def flush():
        if cur_text:
            windows.append({"text": "\n\n".join(cur_text),
                            "frames": list(cur_frames)})

    last_ts = "00:00"
    for block in blocks:
        ts_m = TS_RE.search(block)
        if ts_m:
            last_ts = f"{ts_m.group(1)}:{ts_m.group(2)}"
        fm = FRAME_RE.search(block)
        bt = estimate_tokens(block)
        if cur_text and cur_tok + bt > window_tokens:
            flush()
            cur_text, cur_frames, cur_tok = [], [], 0
        if fm:
            p = fm.group(1)
            cur_frames.append({"id": Path(p).stem, "ts": last_ts, "path": p})
            # keep a compact ref line in the text instead of the bare path
            cur_text.append(f"[{last_ts}] (frame {Path(p).stem})")
        else:
            cur_text.append(block)
        cur_tok += bt
    flush()
    return windows


def resolve_path(ref_path, transcript_path):
    p = Path(ref_path)
    if p.is_absolute() and p.exists():
        return p
    # paths in the transcript are usually relative to the run dir
    cand = transcript_path.parent / ref_path
    return cand if cand.exists() else p


def encode_image(path, max_side):
    data = path.read_bytes()
    mime = "image/png" if path.suffix.lower() == ".png" else "image/jpeg"
    if max_side:
        try:
            from PIL import Image
            import io
            img = Image.open(io.BytesIO(data))
            if max(img.size) > max_side:
                img.thumbnail((max_side, max_side))
                buf = io.BytesIO()
                img.convert("RGB").save(buf, format="JPEG", quality=85)
                data, mime = buf.getvalue(), "image/jpeg"
        except ImportError:
            pass  # PIL absent: send original (more tokens, still works)
    b64 = base64.b64encode(data).decode()
    return f"data:{mime};base64,{b64}"


def make_client(base_url, api_key, timeout):
    try:
        from openai import OpenAI
    except ImportError:
        sys.exit("ERROR: `openai` not installed here. Run under ~/wdir/llm/.venv")
    # timeout=None => no limit (a slow mixed GPU/CPU model can take >>10min per call;
    # the client default of ~600s is what raised RequestTimedOut). max_retries=0 so a
    # rare hiccup doesn't silently re-send a 40-minute generation.
    return OpenAI(base_url=base_url, api_key=api_key, timeout=timeout, max_retries=0)


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


def call_fit(client, model, system, content, temperature, ctx, init_out):
    """Like call(), but bulletproof against token-estimate error: if the server
    rejects the request for exceeding context, parse the REAL input-token count
    from its error and retry with an exactly-fitting output budget. Re-raises the
    context error only when the input alone fills the window (the doc-too-big
    case the caller handles by stopping)."""
    out = init_out
    last = None
    for _ in range(4):
        try:
            return call(client, model, system, content, temperature, out)
        except Exception as e:
            last = e
            m = re.search(r"at least (\d+) input tokens", str(e))
            if "maximum context length" in str(e) and m:
                real_in = int(m.group(1))
                out = ctx - real_in - 64
                if out < 256:
                    break  # input itself ~fills the window — genuine overflow
                continue
            raise
    raise last


def parse_need(raw, valid_ids):
    raw = re.sub(r"^```(?:json)?|```$", "", raw.strip(), flags=re.MULTILINE)
    m = re.search(r"\{.*\}", raw, re.DOTALL)
    if not m:
        return []
    try:
        ids = json.loads(m.group(0)).get("need", [])
    except json.JSONDecodeError:
        return []
    return [i for i in ids if i in valid_ids]


def main():
    p = argparse.ArgumentParser(
        description="Compile a long meeting transcript into a technical reference (refine + on-demand frames).",
        formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__)
    p.add_argument("transcript", type=Path)
    p.add_argument("instruction", nargs="?",
                   default="Compile a detailed technical reference of the workflows and architecture covered.")
    p.add_argument("-o", "--output", type=Path, help="write here (default: <run>/<stem>_reference.md next to the transcript)")
    p.add_argument("--output-dir", type=Path,
                   help="base/parent directory; the run folder (taken from the "
                        "transcript's folder name) is auto-created under it and the "
                        "output written inside (default: the transcript's folder)")
    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)
    p.add_argument("--model", default=DEFAULT_MODEL)
    p.add_argument("--api-key", default="local")
    p.add_argument("--frames", choices=["ondemand", "window", "none"], default="ondemand")
    p.add_argument("--window-tokens", type=int, default=3500, help="transcript tokens per refine step (default 3500)")
    p.add_argument("--max-tokens", type=int, default=8192, help="upper bound on generation; auto-capped to fit --ctx (default 8192)")
    p.add_argument("--ctx", type=int, default=16384, help="model context window; output is auto-capped so input+output fit (default 16384)")
    p.add_argument("--max-image-side", type=int, default=1280, help="downscale frames to this max side (0=off)")
    p.add_argument("--temperature", type=float, default=0.2)
    p.add_argument("--timeout", type=float, default=0, help="per-request timeout in seconds; 0 = no limit (default, for slow local models)")
    p.add_argument("--checkpoint", type=Path, help="write the running doc here after each window (resumable progress)")
    p.add_argument("-q", "--quiet", action="store_true")
    args = p.parse_args()

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

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

    windows = parse_windows(args.transcript, args.window_tokens)
    nframes = sum(len(w["frames"]) for w in windows)
    log(f"{len(windows)} windows, {nframes} frame refs, mode={args.frames}")

    client = make_client(args.base_url, args.api_key, args.timeout or None)
    doc = "# (compilation in progress)\n"

    for wi, w in enumerate(windows, 1):
        wanted = []
        if args.frames != "none" and w["frames"]:
            if args.frames == "window":
                wanted = w["frames"]
            else:  # ondemand: ask the model which frames it needs
                listing = "\n".join(f"- {f['id']} [{f['ts']}]" for f in w["frames"])
                raw = call(client, args.model,
                           TRIAGE_SYS.format(instruction=args.instruction),
                           f"Window transcript:\n{w['text']}\n\nReferenced frames:\n{listing}",
                           0.0, 256)
                valid = {f["id"] for f in w["frames"]}
                keep = set(parse_need(raw, valid))
                wanted = [f for f in w["frames"] if f["id"] in keep]

        # build the refine turn (multimodal if any frames wanted)
        text = (f"CURRENT REFERENCE:\n{doc}\n\n"
                f"NEXT TRANSCRIPT WINDOW:\n{w['text']}")
        if wanted:
            text += "\n\nAttached frames: " + ", ".join(f"{f['id']} [{f['ts']}]" for f in wanted)
            content = [{"type": "text", "text": text}]
            for f in wanted:
                ip = resolve_path(f["path"], args.transcript)
                if ip.exists():
                    content.append({"type": "image_url",
                                    "image_url": {"url": encode_image(ip, args.max_image_side)}})
            log(f"  window {wi}/{len(windows)}: {len(wanted)} frame(s) attached")
        else:
            content = text
            log(f"  window {wi}/{len(windows)}: text-only")

        sys_txt = REFINE_SYS.format(instruction=args.instruction)
        in_tok = content_tokens(content) + estimate_tokens(sys_txt)
        # proportional margin absorbs token-estimate error; call_fit self-corrects
        # from the server's real count if it's still off.
        out_budget = max(256, min(args.max_tokens, args.ctx - in_tok - max(512, in_tok // 20)))
        try:
            doc = call_fit(client, args.model, sys_txt, content,
                           args.temperature, args.ctx, out_budget)
        except Exception as e:
            if "maximum context length" in str(e):
                log(f"  STOP at window {wi}/{len(windows)}: the running doc filled the "
                    f"{args.ctx}-token window. Partial reference is in the checkpoint. A {args.ctx // 1024}k "
                    f"refine cannot hold a whole long-meeting doc — use the chunk-at-breaks + "
                    f"compact-carry + merge design, or a 32k-context profile.")
                break
            raise

        if args.checkpoint:
            args.checkpoint.write_text(doc + "\n")
        if estimate_tokens(doc) > args.window_tokens * 4 and not args.quiet:
            log(f"  note: running doc ~{estimate_tokens(doc)} tok and growing — "
                f"if it nears the context window, switch to a 32k-context profile (qwen14b-gguf)")

    if args.stdout:
        print(doc)
    else:
        # --output-dir is the PARENT; auto-create the run folder under it (named
        # after the transcript's own run folder), matching process_meeting.py.
        base_dir = (args.output_dir / args.transcript.parent.name) if args.output_dir else args.transcript.parent
        if args.output:
            out = args.output
            if not out.is_absolute():
                out = base_dir / out
        else:
            out = default_output(args.transcript, "reference", base_dir)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(doc + "\n")
        log(f"wrote {out}")


if __name__ == "__main__":
    main()