#!/usr/bin/env python3
"""Capture Kimi Code's supported plaintext thinking channel from a local session."""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

KIMI_HOME = Path.home() / ".kimi-code"


def _session_dir(session_id: str) -> Path:
    matches: list[Path] = []
    index = KIMI_HOME / "session_index.jsonl"
    for line in index.read_text(encoding="utf-8").splitlines():
        record = json.loads(line)
        if record.get("sessionId") == session_id:
            matches.append(Path(record["sessionDir"]))
    if not matches:
        raise RuntimeError(f"session not found in {index}: {session_id}")
    return matches[-1]


def _extract_wire(session_dir: Path) -> tuple[list[str], list[str], dict[str, int]]:
    wire = session_dir / "agents" / "main" / "wire.jsonl"
    thoughts: list[str] = []
    answers: list[str] = []
    usage = {
        "input_other": 0,
        "input_cache_read": 0,
        "input_cache_creation": 0,
        "output": 0,
    }
    for line in wire.read_text(encoding="utf-8").splitlines():
        event = json.loads(line)
        if event.get("type") == "context.append_loop_event":
            part = event.get("event", {}).get("part", {})
            if isinstance(part.get("think"), str):
                thoughts.append(part["think"])
            if isinstance(part.get("text"), str):
                answers.append(part["text"])
        elif event.get("type") == "usage.record" and event.get("usageScope") == "turn":
            source = event.get("usage", {})
            usage["input_other"] += int(source.get("inputOther", 0) or 0)
            usage["input_cache_read"] += int(source.get("inputCacheRead", 0) or 0)
            usage["input_cache_creation"] += int(source.get("inputCacheCreation", 0) or 0)
            usage["output"] += int(source.get("output", 0) or 0)
    return thoughts, answers, usage


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--prompt-file", type=Path, required=True)
    parser.add_argument("--label", required=True)
    parser.add_argument("--model", default="kimi-code/k3")
    parser.add_argument("--effort", choices=("low", "high", "max"), default="high")
    parser.add_argument("--kimi", default="kimi")
    parser.add_argument("--output-dir", type=Path, default=Path("traces/kimi"))
    parser.add_argument("--cd", type=Path, default=Path.cwd())
    args = parser.parse_args()

    prompt = args.prompt_file.read_text(encoding="utf-8")
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    args.output_dir.mkdir(parents=True, exist_ok=True)
    destination = args.output_dir / f"{args.label}.json"

    env_effort = f"KIMI_MODEL_THINKING_EFFORT={args.effort}"
    print(f"running {args.model} ({args.effort}) · {args.label}", flush=True)
    # Kimi accepts its effort override via the environment. Use `env` so the
    # setting applies only to this child process and never edits user config.
    command = [
        "env",
        env_effort,
        args.kimi,
        "--model",
        args.model,
        "--prompt",
        prompt,
        "--output-format",
        "stream-json",
    ]
    process = subprocess.Popen(
        command,
        cwd=args.cd.resolve(),
        stdout=subprocess.PIPE,
        stderr=None,
        text=True,
        bufsize=1,
    )
    assert process.stdout is not None
    streamed: list[dict[str, object]] = []
    session_id: str | None = None
    for line in process.stdout:
        event = json.loads(line)
        streamed.append(event)
        if event.get("type") == "session.resume_hint":
            session_id = event.get("session_id")
    returncode = process.wait()
    if returncode:
        raise RuntimeError(f"Kimi exited with status {returncode}")
    if not session_id:
        raise RuntimeError("Kimi did not emit a session ID")

    thoughts, answers, usage = _extract_wire(_session_dir(session_id))
    if not thoughts or not answers:
        raise RuntimeError("session did not contain both thinking and answer text")

    public_trace = {
        "schema": "trace-atlas.kimi-plaintext-thinking.v1",
        "captured_at": timestamp,
        "source": {
            "client": "Kimi Code",
            "client_version": next(
                (item["version"] for item in streamed if item.get("type") == "system.version"),
                "unknown",
            ),
            "model": args.model,
            "thinking_effort": args.effort,
            "session_id": "redacted-for-publication",
        },
        "prompt": prompt,
        "thinking": thoughts,
        "answer": "\n".join(answers),
        "usage": usage,
        "publication_note": (
            "Contains the plaintext thinking channel returned to and persisted by "
            "Kimi Code. It is not a claim about unreturned server-side computation."
        ),
    }
    destination.write_text(
        json.dumps(public_trace, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    thought_chars = sum(map(len, thoughts))
    print(
        f"saved {destination} · {len(thoughts)} thinking chunks · "
        f"{thought_chars:,} thinking characters · {usage['output']:,} output tokens"
    )


if __name__ == "__main__":
    try:
        main()
    except (OSError, RuntimeError, json.JSONDecodeError) as error:
        print(f"error: {error}", file=sys.stderr)
        raise SystemExit(1)