#!/usr/bin/env python3
"""
Manage and extract session transcript backups for the /log skill.

Modes:
  --register BACKUP_NAME --session-id SID --cwd CWD
      Register a new backup file in tracking.json (called by pre-compact hook).

  --list CWD
      List unprocessed transcripts for a project as JSON (small output, for /log skill).

  --extract BACKUP_NAME
      Extract readable conversation text from one backup (for subagent use).

  --mark-processed BACKUP_NAME LOG_FILE
      Mark a backup as processed, recording which log file captured it.

  --mark-all-processed CWD LOG_FILE
      Mark ALL unprocessed backups for a cwd as processed.

Tracking file: ~/.claude/transcript-backups/tracking.json
"""

import json
import os
import sys
from datetime import datetime, timezone
from pathlib import Path

BACKUP_DIR = Path.home() / ".claude" / "transcript-backups"
TRACKING_FILE = BACKUP_DIR / "tracking.json"
MAX_TEXT_LEN = 3000  # max chars per message before truncating


def load_tracking():
    if TRACKING_FILE.exists():
        with open(TRACKING_FILE) as f:
            return json.load(f)
    return {}


def save_tracking(tracking):
    BACKUP_DIR.mkdir(parents=True, exist_ok=True)
    with open(TRACKING_FILE, "w") as f:
        json.dump(tracking, f, indent=2)


def now_utc():
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def extract_text_from_content(content):
    """Extract plain text from a message content field (string or list of blocks)."""
    if isinstance(content, str):
        return content.strip()
    if isinstance(content, list):
        parts = []
        for block in content:
            if not isinstance(block, dict):
                continue
            btype = block.get("type")
            if btype == "text":
                text = block.get("text", "").strip()
                if text:
                    parts.append(text)
            elif btype == "tool_use":
                name = block.get("name", "?")
                inp = block.get("input", {})
                keys = list(inp.keys())[:3] if isinstance(inp, dict) else []
                summary = ", ".join(f"{k}=..." for k in keys) if keys else "..."
                parts.append(f"[tool: {name}({summary})]")
        return " ".join(parts)
    return ""


def extract_conversation(jsonl_path):
    """Read a JSONL transcript; return list of (role, text) tuples."""
    turns = []
    try:
        with open(jsonl_path) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                except json.JSONDecodeError:
                    continue

                if entry.get("type") not in ("user", "assistant"):
                    continue
                if entry.get("isSidechain"):
                    continue

                msg = entry.get("message", {})
                if not isinstance(msg, dict):
                    continue

                role = msg.get("role", entry.get("type", "?"))
                content = msg.get("content", "")
                text = extract_text_from_content(content)
                if not text:
                    continue

                if len(text) > MAX_TEXT_LEN:
                    text = text[:MAX_TEXT_LEN] + f"... [truncated, {len(text)} chars total]"

                turns.append((role, text))
    except Exception as e:
        turns.append(("error", f"Failed to read transcript: {e}"))
    return turns


# ---------------------------------------------------------------------------
# Modes
# ---------------------------------------------------------------------------

def cmd_register(args):
    """Register a new backup in tracking.json."""
    try:
        idx = args.index("--register")
        backup_name = args[idx + 1]
        session_id = args[args.index("--session-id") + 1]
        cwd = args[args.index("--cwd") + 1]
    except (ValueError, IndexError) as e:
        print(f"Usage: --register BACKUP_NAME --session-id SID --cwd CWD", file=sys.stderr)
        sys.exit(1)

    tracking = load_tracking()
    tracking[backup_name] = {
        "session_id": session_id,
        "cwd": cwd,
        "saved_at": now_utc(),
        "processed": False,
        "log_file": None,
    }
    save_tracking(tracking)
    print(f"Registered {backup_name} for {cwd}")


def cmd_list(cwd):
    """List unprocessed transcripts for a cwd as compact JSON."""
    tracking = load_tracking()
    unprocessed = [
        {
            "backup": fname,
            "session_id": entry["session_id"][:8],
            "saved_at": entry.get("saved_at", "?"),
            "path": str(BACKUP_DIR / fname),
            "exists": (BACKUP_DIR / fname).exists(),
        }
        for fname, entry in sorted(tracking.items(), key=lambda x: x[1].get("saved_at", ""))
        if entry.get("cwd") == cwd and not entry.get("processed", False)
    ]
    print(json.dumps(unprocessed, indent=2))


def cmd_extract(backup_name):
    """Extract readable conversation from one backup file."""
    backup_path = BACKUP_DIR / backup_name
    if not backup_path.exists():
        print(f"Backup not found: {backup_path}", file=sys.stderr)
        sys.exit(1)

    tracking = load_tracking()
    entry = tracking.get(backup_name, {})

    print(f"{'='*72}")
    print(f"Transcript: {backup_name}  (pre-compaction snapshot)")
    print(f"Project:    {entry.get('cwd', 'unknown')}")
    print(f"Session:    {entry.get('session_id', '?')[:8]}  |  Saved: {entry.get('saved_at', '?')}")
    print(f"{'='*72}")

    turns = extract_conversation(backup_path)
    if not turns:
        print("[no readable conversation found]")
        return

    for role, text in turns:
        label = "USER" if role == "user" else "ASST"
        print(f"\n[{label}] {text}")

    print(f"\n[{len(turns)} turns extracted]")


def cmd_mark_processed(backup_name, log_file):
    """Mark one backup as processed."""
    tracking = load_tracking()
    if backup_name not in tracking:
        print(f"Warning: {backup_name} not found in tracking, adding entry", file=sys.stderr)
        tracking[backup_name] = {}
    tracking[backup_name]["processed"] = True
    tracking[backup_name]["processed_at"] = now_utc()
    tracking[backup_name]["log_file"] = log_file
    save_tracking(tracking)
    print(f"Marked {backup_name} as processed → {log_file}")


def cmd_mark_all_processed(cwd, log_file):
    """Mark all unprocessed backups for a cwd as processed."""
    tracking = load_tracking()
    count = 0
    for fname, entry in tracking.items():
        if entry.get("cwd") == cwd and not entry.get("processed", False):
            tracking[fname]["processed"] = True
            tracking[fname]["processed_at"] = now_utc()
            tracking[fname]["log_file"] = log_file
            count += 1
    save_tracking(tracking)
    print(f"Marked {count} transcript(s) as processed → {log_file}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    args = sys.argv[1:]

    if not args:
        print(__doc__)
        sys.exit(1)

    if "--register" in args:
        cmd_register(args)
    elif "--list" in args:
        idx = args.index("--list")
        cwd = args[idx + 1] if idx + 1 < len(args) else ""
        if not cwd:
            print("Usage: --list CWD", file=sys.stderr)
            sys.exit(1)
        cmd_list(cwd)
    elif "--extract" in args:
        idx = args.index("--extract")
        backup_name = args[idx + 1] if idx + 1 < len(args) else ""
        if not backup_name:
            print("Usage: --extract BACKUP_NAME", file=sys.stderr)
            sys.exit(1)
        cmd_extract(backup_name)
    elif "--mark-all-processed" in args:
        idx = args.index("--mark-all-processed")
        cwd = args[idx + 1] if idx + 1 < len(args) else ""
        log_idx = args.index("--log-file") if "--log-file" in args else -1
        log_file = args[log_idx + 1] if log_idx >= 0 and log_idx + 1 < len(args) else ""
        if not cwd or not log_file:
            print("Usage: --mark-all-processed CWD --log-file LOG_FILE", file=sys.stderr)
            sys.exit(1)
        cmd_mark_all_processed(cwd, log_file)
    elif "--mark-processed" in args:
        idx = args.index("--mark-processed")
        backup_name = args[idx + 1] if idx + 1 < len(args) else ""
        log_idx = args.index("--log-file") if "--log-file" in args else -1
        log_file = args[log_idx + 1] if log_idx >= 0 and log_idx + 1 < len(args) else ""
        if not backup_name or not log_file:
            print("Usage: --mark-processed BACKUP_NAME --log-file LOG_FILE", file=sys.stderr)
            sys.exit(1)
        cmd_mark_processed(backup_name, log_file)
    else:
        print(f"Unknown arguments: {args}", file=sys.stderr)
        print(__doc__)
        sys.exit(1)


if __name__ == "__main__":
    main()