#!/usr/bin/env python3
"""
Minimal meeting summarizer: walk an enhanced transcript in order, summarizing as
you go, LOOKING AT EVERY referenced frame for context. No triage, no grounding
rules, no instruction block — just the SYS prompt + transcript + every frame.
Edit SYS below to change the steer.

Talks to a local OpenAI-compatible endpoint (ollama / vLLM / llama.cpp).

Usage:
    ~/wdir/llm/.venv/bin/python ctrl/summarize/summarize_simple.py <stem>_enhanced.txt \\
        --base-url http://localhost:11434/v1 --model gemma3-27b-16k
"""
import argparse, base64, io, re, sys
from pathlib import Path

DEFAULT_BASE_URL = "http://localhost:11434/v1"
DEFAULT_MODEL = "gemma3-27b-16k"
CHARS_PER_TOKEN = 4.0
FRAME_RE = re.compile(r"Frame:\s+(\S+\.(?:jpg|jpeg|png))", re.IGNORECASE)
TS_RE = re.compile(r"\[(\d+):(\d+)\]")

SYS = """\
summarize a meeting/training from its transcript, 
read screen frames interlieved in the dialog"""


def est(t): return int(len(t) / CHARS_PER_TOKEN)


def windows(path, wtok, fcap):
    """Break on EITHER the text budget OR fcap frames, so every referenced frame
    lands in some window and gets attached — dense stretches just become more
    (smaller) windows. Nothing is sampled away."""
    blocks = re.split(r"\n\s*\n", path.read_text())
    out, cur, frames, tok, ts = [], [], [], 0, "00:00"
    def flush():
        if cur: out.append({"text": "\n\n".join(cur), "frames": list(frames)})
    for b in blocks:
        m = TS_RE.search(b)
        if m: ts = f"{m.group(1)}:{m.group(2)}"
        fm = FRAME_RE.search(b)
        bt = est(b)
        if cur and (tok + bt > wtok or (fm and len(frames) >= fcap)):
            flush(); cur, frames, tok = [], [], 0
        if fm:
            frames.append({"ts": ts, "path": fm.group(1)})
            cur.append(f"[{ts}] (frame)")
        else:
            cur.append(b)
        tok += bt
    flush()
    return out


def resolve(p, transcript):
    pp = Path(p)
    if pp.is_absolute() and pp.exists(): return pp
    c = transcript.parent / p
    return c if c.exists() else pp


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


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("transcript", type=Path)
    p.add_argument("-o", "--output", type=Path)
    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("--base-url", default=DEFAULT_BASE_URL)
    p.add_argument("--model", default=DEFAULT_MODEL)
    p.add_argument("--api-key", default="local")
    p.add_argument("--window-tokens", type=int, default=1500, help="transcript tokens per step")
    p.add_argument("--max-frames", type=int, default=6, help="frames per step — also forces a new window, so EVERY frame is still seen")
    p.add_argument("--max-side", type=int, default=768, help="downscale frames to this max side")
    p.add_argument("--ctx", type=int, default=16384)
    p.add_argument("--max-tokens", type=int, default=4096)
    p.add_argument("--temperature", type=float, default=0.3)
    p.add_argument("--timeout", type=float, default=0, help="per-request seconds; 0 = no limit (slow local models)")
    p.add_argument("--checkpoint", type=Path)
    args = p.parse_args()

    if not args.transcript.is_file():
        sys.exit(f"ERROR: not found: {args.transcript}")
    try:
        from openai import OpenAI
    except ImportError:
        sys.exit("ERROR: `openai` not installed here. Run under ~/wdir/llm/.venv")
    client = OpenAI(base_url=args.base_url, api_key=args.api_key, timeout=(args.timeout or None), max_retries=0)

    wins = windows(args.transcript, args.window_tokens, args.max_frames)
    nfr = sum(len(w["frames"]) for w in wins)
    print(f"[simple] {len(wins)} windows, {nfr} frame refs (all inspected)", file=sys.stderr)

    # --output-dir is the PARENT; auto-create the run folder under it (named after
    # the transcript's own run folder), matching process_meeting.py's layout.
    base_dir = (args.output_dir / args.transcript.parent.name) if args.output_dir else args.transcript.parent
    out_path = args.output
    if not out_path:
        stem = args.transcript.stem
        if stem.endswith("_enhanced"): stem = stem[:-9]
        out_path = base_dir / f"{stem}_summary_simple.md"
    elif not out_path.is_absolute():
        out_path = base_dir / out_path
    out_path.parent.mkdir(parents=True, exist_ok=True)

    # MAP + APPEND: summarize each window independently and APPEND it under its
    # timestamp — no carried/re-emitted running summary, so nothing collapses and
    # every segment is kept. The output file IS the accumulator (also the checkpoint).
    doc = ""
    for i, w in enumerate(wins, 1):
        m = TS_RE.search(w["text"])
        start_ts = f"{m.group(1)}:{m.group(2)}" if m else "?"
        content = [{"type": "text", "text": f"PART OF THE MEETING:\n{w['text']}"}]
        attached = 0
        for f in w["frames"]:
            ip = resolve(f["path"], args.transcript)
            if ip.exists():
                content.append({"type": "image_url", "image_url": {"url": encode(ip, args.max_side)}})
                attached += 1
        print(f"[simple]   window {i}/{len(wins)} [{start_ts}]: {attached} frame(s)", file=sys.stderr)

        in_tok = est(str(content)) + est(SYS) + attached * 300  # rough, incl. vision
        out = max(256, min(args.max_tokens, args.ctx - in_tok - max(512, in_tok // 20)))
        try:
            r = client.chat.completions.create(
                model=args.model,
                messages=[{"role": "system", "content": SYS}, {"role": "user", "content": content}],
                temperature=args.temperature, max_tokens=out)
            part = r.choices[0].message.content.strip()
        except Exception as e:
            if "context length" in str(e) or "maximum context" in str(e):
                print(f"[simple]   window {i}: too big for ctx — skipping (lower --max-frames "
                      f"or --window-tokens to avoid). Continuing.", file=sys.stderr)
                continue
            raise

        doc += f"## [{start_ts}]\n\n{part}\n\n"
        out_path.write_text(doc)            # the file grows as we go (= live result)
        if args.checkpoint:
            args.checkpoint.write_text(doc)  # mirror, so `tail -f` keeps working

    print(f"[simple] wrote {out_path}", file=sys.stderr)


if __name__ == "__main__":
    main()