#!/usr/bin/env python3
"""
Wrapper around `codex` that captures the most recent session transcript and
writes it as prettified JSON into the current working directory.
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Iterable


SESSION_PATTERN = "rollout-*.jsonl"
DEFAULT_OUTDIR = Path.home() / "codex-logs"


def gather_session_files(root: Path) -> set[Path]:
    if not root.exists():
        return set()
    return {path for path in root.rglob(SESSION_PATTERN) if path.is_file()}


def pick_latest(paths: Iterable[Path]) -> Path | None:
    try:
        return max(paths, key=lambda p: p.stat().st_mtime)
    except ValueError:
        return None


def export_jsonl(jsonl_path: Path, destination: Path) -> None:
    lines = []
    for raw in jsonl_path.read_text(encoding="utf-8").splitlines():
        raw = raw.strip()
        if not raw:
            continue
        lines.append(json.loads(raw))
    destination.write_text(json.dumps(lines, ensure_ascii=False, indent=2), encoding="utf-8")


def resolve_codex_binary() -> list[str]:
    override = os.environ.get("CODEX_CAPTURE_BINARY")
    if override:
        return [override]
    return [os.environ.get("CODEX_CLI_PATH", "codex")]


def main(argv: list[str]) -> int:
    session_root = Path(os.environ.get("CODEX_SESSION_DIR", Path.home() / ".codex" / "sessions"))
    before = gather_session_files(session_root)

    exit_code = subprocess.call([*resolve_codex_binary(), *argv])

    after = gather_session_files(session_root)
    new_files = sorted(after - before, key=lambda p: p.stat().st_mtime)
    target_session = new_files[-1] if new_files else pick_latest(after)

    if target_session is None:
        print("codex_capture: no session transcript found", file=sys.stderr)
        return exit_code

    timestamp = time.strftime("%Y%m%d-%H%M%S")
    out_root = Path(os.environ.get("CODEX_CAPTURE_OUTDIR", DEFAULT_OUTDIR))
    out_root.mkdir(parents=True, exist_ok=True)
    out_file = out_root / f"codex-session-{timestamp}.json"
    export_jsonl(target_session, out_file)
    print(f"codex_capture: saved {target_session} -> {out_file}")

    return exit_code


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))