#!/usr/bin/env python3
"""
parse-ai-session — read a JSONL session log from Claude Code, Kilo, Codex,
or any AI tool that stores sessions in JSONL format. Extracts messages and
metadata, produces a structured summary, and passes it to update-session-docs
which appends it to the relevant companion files.

Usage: parse-ai-session <session.jsonl> <tool-name>
  tool-name: any string — appears in the Session Log entry header
             common values: claude, kilo, codex, aider

For tools that do not produce JSONL (ChatGPT exports, Kimi, Ollama frontends,
plain text logs), use parse-ghostty-session as a generic text log parser, or
write a dedicated parse-{toolname} script following the same pattern.

Install: copy to ~/.local/bin/parse-ai-session and chmod +x

Config (via ~/.devme/config.json):
  vault_dir   — where session archives are saved (default: ~/Documents/sessions)
  tool_paths  — override default log directories per tool (see session-close)
"""

import json
import os
import re
import subprocess
import sys
from datetime import datetime
from pathlib import Path


# ── Config ────────────────────────────────────────────────────────────────────

def _load_config() -> dict:
    defaults = {"vault_dir": str(Path.home() / "Documents" / "sessions")}
    cfg_path = Path.home() / ".devme" / "config.json"
    legacy_cfg_path = Path.home() / ".ash" / "config.json"
    if not cfg_path.exists() and legacy_cfg_path.exists():
        cfg_path = legacy_cfg_path
    if cfg_path.exists():
        try:
            loaded = json.loads(cfg_path.read_text())
            if "vault_dir" in loaded:
                defaults["vault_dir"] = loaded["vault_dir"]
        except Exception:
            pass
    return defaults

CFG       = _load_config()
VAULT_DIR = Path(CFG["vault_dir"]).expanduser()
CLAUDE_BIN = Path.home() / ".local/bin/claude"


# ── Transcript parsing ────────────────────────────────────────────────────────

def extract_messages(path: Path):
    """Parse JSONL transcript → readable messages + session metadata."""
    messages = []
    meta = {
        "model": None, "slug": None, "cwd": None,
        "git_branch": None, "session_id": None,
        "timestamp": None, "version": None,
    }

    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue

            for field, key in [
                ("session_id", "sessionId"), ("slug", "slug"), ("cwd", "cwd"),
                ("git_branch", "gitBranch"), ("version", "version"),
            ]:
                if not meta[field] and key in obj:
                    meta[field] = obj[key]

            if not meta["timestamp"] and "timestamp" in obj:
                meta["timestamp"] = obj["timestamp"]

            if obj.get("type") in ("user", "assistant") and "message" in obj:
                msg  = obj["message"]
                role = msg.get("role", "")
                if not meta["model"] and msg.get("model"):
                    meta["model"] = msg["model"]
                content = msg.get("content", "")
                if isinstance(content, list):
                    text = "\n".join(
                        c.get("text", "") for c in content
                        if isinstance(c, dict) and c.get("type") == "text"
                    )
                else:
                    text = str(content)
                text = text.strip()
                if text:
                    messages.append(f"{role.upper()}: {text}")

    return messages, meta


def build_prompt(tool: str, messages: list, meta: dict) -> str:
    if len(messages) > 60:
        sample = messages[:30] + ["\n[... middle of session omitted for length ...]\n"] + messages[-30:]
    else:
        sample = messages

    transcript = "\n\n---\n\n".join(sample)
    return f"""You are summarizing a {tool} AI dev session for the developer's personal archive. Be specific and concrete — this is reference material they'll search later.

Session metadata:
- Tool: {tool}
- Model: {meta.get("model") or "unknown"}
- Working directory: {meta.get("cwd") or "unknown"}
- Branch: {meta.get("git_branch") or "unknown"}
- Session: {meta.get("slug") or "unknown"}

Transcript:
{transcript}

Write a session summary using exactly these sections (omit a section only if there is genuinely nothing to say for it):

## Key Points
- [3–7 concise bullets: what was accomplished, what was discussed, what changed]

## Decisions & Reasoning
[Significant decisions made, criteria used, tradeoffs considered, rejected approaches. Use bullets.]

## Instructionals & Code
[Reproduce verbatim any how-tos, commands, config snippets, or significant code written. Use fenced code blocks with language tags.]

## Takeaway
[1–2 sentences, plain English. What was this session fundamentally about and what was the outcome?]"""


def already_archived(session_id: str) -> bool:
    if not session_id or not VAULT_DIR.exists():
        return False
    short = session_id[:8]
    return any(short in p.name for p in VAULT_DIR.rglob("*.md"))


def parse_timestamp(ts):
    if not ts:
        return datetime.now()
    try:
        return datetime.fromisoformat(ts.replace("Z", "+00:00")).astimezone()
    except ValueError:
        return datetime.now()


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    if len(sys.argv) < 3:
        print("Usage: parse-ai-session <transcript.jsonl> <tool>", file=sys.stderr)
        sys.exit(1)

    transcript_path = Path(sys.argv[1])
    tool            = sys.argv[2]

    if not transcript_path.exists():
        print(f"Transcript not found: {transcript_path}", file=sys.stderr)
        sys.exit(1)

    if not CLAUDE_BIN.exists():
        print(f"claude binary not found at {CLAUDE_BIN}", file=sys.stderr)
        sys.exit(1)

    messages, meta = extract_messages(transcript_path)
    if not messages:
        print("No messages found — skipping", file=sys.stderr)
        sys.exit(0)

    session_id = meta.get("session_id") or ""
    if already_archived(session_id):
        print(f"Session {session_id[:8]} already archived — skipping")
        sys.exit(0)

    prompt = build_prompt(tool, messages, meta)
    env    = os.environ.copy()
    env.pop("CLAUDECODE", None)

    try:
        result = subprocess.run(
            [str(CLAUDE_BIN), "-p", "--no-session-persistence"],
            input=prompt, capture_output=True, text=True, timeout=120, env=env,
        )
    except subprocess.TimeoutExpired:
        print("claude -p timed out after 120s", file=sys.stderr)
        sys.exit(1)

    if result.returncode != 0:
        print(f"claude -p failed (exit {result.returncode}):\n{result.stderr[:500]}", file=sys.stderr)
        sys.exit(1)

    summary = result.stdout.strip()
    if not summary:
        print("claude -p returned empty output — skipping", file=sys.stderr)
        sys.exit(1)

    # Build output path
    ts         = parse_timestamp(meta.get("timestamp"))
    slug       = re.sub(r"[^\w\-]", "_", meta.get("slug") or "unnamed")
    short_id   = session_id[:8] if session_id else "unknown"

    VAULT_DIR.mkdir(parents=True, exist_ok=True)
    month_dir = VAULT_DIR / ts.strftime("%Y-%m")
    month_dir.mkdir(exist_ok=True)

    filename    = f"{ts.strftime('%Y-%m-%d-%H%M%S')}-{tool}-{slug}-{short_id}.md"
    output_path = month_dir / filename

    output = f"""# {slug} ({tool})

| | |
|---|---|
| **Date** | {ts.strftime("%Y-%m-%d %H:%M")} |
| **Tool** | {tool} |
| **Model** | {meta.get("model") or "unknown"} |
| **Project** | {meta.get("cwd") or "unknown"} |
| **Branch** | {meta.get("git_branch") or "unknown"} |
| **Session ID** | `{session_id}` |
| **Transcript** | `{transcript_path}` |

---

{summary}
"""
    output_path.write_text(output)
    print(f"Saved: {output_path}")

    # Update companion files asynchronously
    update_docs = Path.home() / ".local/bin/update-session-docs"
    if update_docs.exists():
        subprocess.Popen(
            [str(update_docs), "--summary", summary,
             "--cwd",  meta.get("cwd") or str(Path.home()),
             "--tool", tool, "--date", ts.isoformat()],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        )


if __name__ == "__main__":
    main()
