#!/usr/bin/env python3
"""
swl-ses token-analysis
Analiza el uso de tokens de Claude Code desde archivos JSONL en ~/.claude/projects/.
Genera reporte Markdown con desglose por proyecto, sesiones costosas y subagentes.

Uso:
  python scripts/token-analysis.py                    # Todas las sesiones
  SINCE_DAYS=30 python scripts/token-analysis.py      # Últimos 30 días
  SINCE_DATE=2026-04-01 python scripts/token-analysis.py  # Desde fecha específica

Salida:
  .planning/analysis/token_report.md   — Reporte principal
  .planning/analysis/prompts/          — Prompts del usuario por proyecto
"""

import json
import os
import sys
from pathlib import Path
from collections import defaultdict
from datetime import datetime, timedelta, timezone

PROJECTS_DIR = Path.home() / ".claude" / "projects"

# Directorio de salida: .planning/analysis/ del proyecto actual, o fallback a home
CWD = Path.cwd()
PLANNING_DIR = CWD / ".planning"
if PLANNING_DIR.is_dir():
    OUTPUT_DIR = PLANNING_DIR / "analysis" / "tokens"
else:
    OUTPUT_DIR = Path.home() / ".claude" / "analysis" / "tokens"

# Filtro temporal vía variables de entorno
SINCE_DAYS = int(os.environ.get("SINCE_DAYS", "0")) or None
SINCE_DATE = os.environ.get("SINCE_DATE")


def extract_text_content(content):
    """Extrae texto del contenido de un mensaje (string o lista)."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict):
                if item.get("type") == "text":
                    parts.append(item.get("text", ""))
            elif isinstance(item, str):
                parts.append(item)
        return "\n".join(parts).strip()
    return ""


def is_human_prompt(msg_obj):
    """Verifica si es un prompt del usuario (no resultado de herramienta)."""
    content = msg_obj.get("message", {}).get("content", "")
    if isinstance(content, list):
        types = [i.get("type") for i in content if isinstance(i, dict)]
        if types and all(t == "tool_result" for t in types):
            return False
    return True


def parse_session(jsonl_path, is_subagent=False):
    """Parsea un archivo JSONL de sesión."""
    usage_total = defaultdict(int)
    prompts = []
    agent_id = None
    session_id = None
    timestamp_start = None
    subagent_sessions = []

    try:
        with open(jsonl_path, encoding="utf-8", errors="replace") as f:
            lines = f.readlines()
    except Exception:
        return None

    for line in lines:
        try:
            obj = json.loads(line)
        except json.JSONDecodeError:
            continue

        msg_type = obj.get("type")
        ts = obj.get("timestamp")
        if ts and not timestamp_start:
            timestamp_start = ts

        if not agent_id:
            agent_id = obj.get("agentId")
        if not session_id:
            session_id = obj.get("sessionId")

        if msg_type == "assistant":
            usage = obj.get("message", {}).get("usage", {})
            usage_total["input_tokens"] += usage.get("input_tokens", 0)
            usage_total["cache_creation_input_tokens"] += usage.get("cache_creation_input_tokens", 0)
            usage_total["cache_read_input_tokens"] += usage.get("cache_read_input_tokens", 0)
            usage_total["output_tokens"] += usage.get("output_tokens", 0)

        elif msg_type == "user":
            user_type = obj.get("userType", "")
            is_sidechain = obj.get("isSidechain", False)
            content = obj.get("message", {}).get("content", "")
            text = extract_text_content(content)

            if text and not is_sidechain and is_human_prompt(obj) and user_type != "tool":
                prompts.append({
                    "text": text,
                    "timestamp": obj.get("timestamp"),
                    "entrypoint": obj.get("entrypoint", ""),
                })

    # Buscar sesiones de subagentes
    session_dir = jsonl_path.parent / jsonl_path.stem
    if session_dir.is_dir():
        subagents_dir = session_dir / "subagents"
        if subagents_dir.is_dir():
            for sub_file in subagents_dir.glob("*.jsonl"):
                sub_data = parse_session(sub_file, is_subagent=True)
                if sub_data:
                    sub_data["subagent_file"] = str(sub_file.name)
                    subagent_sessions.append(sub_data)

    total_tokens = (
        usage_total["input_tokens"]
        + usage_total["cache_creation_input_tokens"]
        + usage_total["cache_read_input_tokens"]
        + usage_total["output_tokens"]
    )

    return {
        "file": str(jsonl_path),
        "session_id": session_id or jsonl_path.stem,
        "agent_id": agent_id,
        "is_subagent": is_subagent,
        "timestamp_start": timestamp_start,
        "usage": dict(usage_total),
        "total_tokens": total_tokens,
        "prompts": prompts,
        "subagent_sessions": subagent_sessions,
    }


def get_project_name(project_dir_name):
    """Convierte nombre de directorio a nombre legible de proyecto."""
    name = project_dir_name
    # Eliminar prefijo de ruta del usuario (Windows o Linux)
    home_name = Path.home().name
    prefixes_to_strip = [
        f"-Users-{home_name}-",
        f"Users-{home_name}-",
        f"-home-{home_name}-",
        f"home-{home_name}-",
        # Windows con drive letter
        "C--Users-",
        "D--",
        "C--",
    ]
    for p in prefixes_to_strip:
        if name.startswith(p):
            name = name[len(p):]
            break
    # Reemplazar separadores restantes
    name = name.replace("-", "/", 3).replace("-", "/")
    return name or project_dir_name


def get_cutoff():
    """Retorna datetime UTC de corte, o None para todas las sesiones."""
    if SINCE_DATE:
        return datetime.fromisoformat(SINCE_DATE).replace(tzinfo=timezone.utc)
    if SINCE_DAYS:
        return datetime.now(timezone.utc) - timedelta(days=SINCE_DAYS)
    return None


def session_in_range(session, cutoff):
    if not cutoff or not session["timestamp_start"]:
        return True
    ts_str = session["timestamp_start"]
    try:
        ts = datetime.fromisoformat(ts_str.replace("Z", "+00:00"))
        return ts >= cutoff
    except ValueError:
        return True


def analyze_all():
    """Analiza todos los proyectos y sesiones."""
    projects = defaultdict(list)
    cutoff = get_cutoff()

    if not PROJECTS_DIR.is_dir():
        print(f"No se encontró {PROJECTS_DIR}")
        sys.exit(1)

    for project_dir in sorted(PROJECTS_DIR.iterdir()):
        if not project_dir.is_dir():
            continue
        project_name = get_project_name(project_dir.name)

        for jsonl_file in sorted(project_dir.glob("*.jsonl")):
            session = parse_session(jsonl_file)
            if session and session["total_tokens"] > 0 and session_in_range(session, cutoff):
                projects[project_name].append(session)

    return projects


def format_tokens(n):
    """Formatea conteo de tokens con comas."""
    return f"{n:,}"


def estimate_cost_usd(usage):
    """Estima costo en USD basado en precios de Claude Sonnet 4."""
    # Precios por millón de tokens (Claude Sonnet 4, aprox.)
    input_price = 3.0    # $3/M input tokens
    output_price = 15.0  # $15/M output tokens
    cache_read_price = 0.30  # $0.30/M cache read
    cache_write_price = 3.75  # $3.75/M cache write

    cost = (
        usage.get("input_tokens", 0) * input_price / 1_000_000
        + usage.get("output_tokens", 0) * output_price / 1_000_000
        + usage.get("cache_read_input_tokens", 0) * cache_read_price / 1_000_000
        + usage.get("cache_creation_input_tokens", 0) * cache_write_price / 1_000_000
    )
    return cost


def summarize_projects(projects):
    """Construye resumen por proyecto."""
    summaries = []
    for project_name, sessions in projects.items():
        total = defaultdict(int)
        all_subagent_tokens = 0
        subagent_count = 0

        for session in sessions:
            for k, v in session["usage"].items():
                total[k] += v
            for sub in session["subagent_sessions"]:
                all_subagent_tokens += sub["total_tokens"]
                subagent_count += 1

        grand_total = sum(total.values())
        summaries.append({
            "project": project_name,
            "sessions": len(sessions),
            "usage": dict(total),
            "total_tokens": grand_total,
            "cost_usd": estimate_cost_usd(dict(total)),
            "subagent_tokens": all_subagent_tokens,
            "subagent_count": subagent_count,
        })

    summaries.sort(key=lambda x: x["total_tokens"], reverse=True)
    return summaries


def find_costly_sessions(projects, top_n=20):
    """Encuentra las sesiones más costosas."""
    all_sessions = []
    for project_name, sessions in projects.items():
        for session in sessions:
            all_sessions.append((project_name, session))

    all_sessions.sort(key=lambda x: x[1]["total_tokens"], reverse=True)
    return all_sessions[:top_n]


def find_costly_subagents(projects, top_n=20):
    """Encuentra los subagentes más costosos."""
    all_subs = []
    for project_name, sessions in projects.items():
        for session in sessions:
            for sub in session["subagent_sessions"]:
                all_subs.append((project_name, session["session_id"], sub))

    all_subs.sort(key=lambda x: x[2]["total_tokens"], reverse=True)
    return all_subs[:top_n]


def write_report(projects, summaries):
    """Genera el reporte principal en Markdown."""
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    report_path = OUTPUT_DIR / "token_report.md"

    lines = []
    cutoff = get_cutoff()
    date_range = f"Desde {cutoff.strftime('%Y-%m-%d')}" if cutoff else "Todas las sesiones"
    lines.append("# Análisis de uso de tokens — Claude Code")
    lines.append(f"\nGenerado: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} | Rango: {date_range}\n")

    # Totales globales
    grand_input = sum(s["usage"].get("input_tokens", 0) for s in summaries)
    grand_cache_create = sum(s["usage"].get("cache_creation_input_tokens", 0) for s in summaries)
    grand_cache_read = sum(s["usage"].get("cache_read_input_tokens", 0) for s in summaries)
    grand_output = sum(s["usage"].get("output_tokens", 0) for s in summaries)
    grand_total = sum(s["total_tokens"] for s in summaries)
    grand_cost = sum(s["cost_usd"] for s in summaries)
    total_sessions = sum(s["sessions"] for s in summaries)
    total_subagent_tokens = sum(s["subagent_tokens"] for s in summaries)
    total_subagent_count = sum(s["subagent_count"] for s in summaries)

    lines.append("## Totales globales\n")
    lines.append(f"- **Proyectos**: {len(summaries)}")
    lines.append(f"- **Sesiones**: {total_sessions:,}")
    lines.append(f"- **Tokens totales**: {format_tokens(grand_total)}")
    lines.append(f"  - Input: {format_tokens(grand_input)}")
    lines.append(f"  - Cache creation: {format_tokens(grand_cache_create)}")
    lines.append(f"  - Cache read: {format_tokens(grand_cache_read)}")
    lines.append(f"  - Output: {format_tokens(grand_output)}")
    lines.append(f"- **Costo estimado**: ${grand_cost:,.2f} USD")
    lines.append(f"- **Subagentes**: {total_subagent_count:,} sesiones ({format_tokens(total_subagent_tokens)} tokens)")
    lines.append("")

    # Por proyecto
    lines.append("## Por proyecto\n")
    lines.append("| Proyecto | Sesiones | Tokens | Costo USD | Input | Cache Read | Output | Subagentes |")
    lines.append("|----------|----------|--------|-----------|-------|------------|--------|------------|")

    for s in summaries:
        u = s["usage"]
        lines.append(
            f"| {s['project']} | {s['sessions']} "
            f"| {format_tokens(s['total_tokens'])} "
            f"| ${s['cost_usd']:,.2f} "
            f"| {format_tokens(u.get('input_tokens', 0))} "
            f"| {format_tokens(u.get('cache_read_input_tokens', 0))} "
            f"| {format_tokens(u.get('output_tokens', 0))} "
            f"| {s['subagent_count']} ({format_tokens(s['subagent_tokens'])}) |"
        )

    lines.append("")

    # Sesiones más costosas
    lines.append("## Sesiones más costosas\n")
    costly = find_costly_sessions(projects, top_n=25)

    for i, (proj, session) in enumerate(costly, 1):
        u = session["usage"]
        cost = estimate_cost_usd(u)
        lines.append(f"### {i}. {proj} — {format_tokens(session['total_tokens'])} tokens (${cost:,.2f})")
        lines.append(f"- **Sesión**: `{session['session_id']}`")
        if session["timestamp_start"]:
            lines.append(f"- **Inicio**: {session['timestamp_start'][:19].replace('T', ' ')}")
        lines.append(f"- **Tokens**: input={format_tokens(u.get('input_tokens', 0))}, cache_create={format_tokens(u.get('cache_creation_input_tokens', 0))}, cache_read={format_tokens(u.get('cache_read_input_tokens', 0))}, output={format_tokens(u.get('output_tokens', 0))}")
        lines.append(f"- **Subagentes**: {len(session['subagent_sessions'])}")

        if session["prompts"]:
            first = session["prompts"][0]["text"][:400].replace("\n", " ")
            lines.append(f"- **Primer prompt**: {first}")
        lines.append("")

    # Subagentes más costosos
    lines.append("## Subagentes más costosos\n")
    costly_subs = find_costly_subagents(projects, top_n=20)

    lines.append("| # | Proyecto | Sesión padre | Archivo subagente | Tokens | Input | Output |")
    lines.append("|---|----------|-------------|-------------------|--------|-------|--------|")

    for i, (proj, session_id, sub) in enumerate(costly_subs, 1):
        u = sub["usage"]
        input_total = u.get("input_tokens", 0) + u.get("cache_creation_input_tokens", 0) + u.get("cache_read_input_tokens", 0)
        lines.append(
            f"| {i} | {proj} | `{session_id[:8]}...` "
            f"| `{sub.get('subagent_file', '?')}` "
            f"| {format_tokens(sub['total_tokens'])} "
            f"| {format_tokens(input_total)} "
            f"| {format_tokens(u.get('output_tokens', 0))} |"
        )

    lines.append("")

    with open(report_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))

    print(f"Reporte: {report_path}")
    return report_path


def write_prompts_by_project(projects):
    """Exporta prompts del usuario por proyecto."""
    prompts_dir = OUTPUT_DIR / "prompts"
    prompts_dir.mkdir(parents=True, exist_ok=True)

    for project_name, sessions in projects.items():
        all_prompts = []
        for session in sessions:
            for prompt in session["prompts"]:
                all_prompts.append({
                    "session_id": session["session_id"],
                    "timestamp": prompt.get("timestamp", ""),
                    "entrypoint": prompt.get("entrypoint", ""),
                    "text": prompt["text"],
                })

        if not all_prompts:
            continue

        all_prompts.sort(key=lambda x: x["timestamp"] or "")

        safe_name = project_name.replace("/", "_").replace("\\", "_").replace(" ", "_")[:80]
        out_path = prompts_dir / f"{safe_name}.md"

        lines = []
        lines.append(f"# Prompts: {project_name}")
        lines.append(f"\n{len(all_prompts)} prompts en {len(sessions)} sesiones\n")

        for i, p in enumerate(all_prompts, 1):
            ts = p["timestamp"][:19].replace("T", " ") if p["timestamp"] else "?"
            lines.append(f"## {i}. [{ts}] Sesión `{p['session_id'][:8]}`")
            if p["entrypoint"]:
                lines.append(f"*entrypoint: {p['entrypoint']}*")
            lines.append("")
            lines.append(p["text"][:2000])
            lines.append("")

        with open(out_path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))

    print(f"Prompts: {prompts_dir}")


def print_summary(summaries, projects):
    """Imprime resumen en consola."""
    grand_total = sum(s["total_tokens"] for s in summaries)
    grand_cost = sum(s["cost_usd"] for s in summaries)
    total_sessions = sum(s["sessions"] for s in summaries)

    print(f"\nTotal: {format_tokens(grand_total)} tokens, ${grand_cost:,.2f} USD, {total_sessions} sesiones, {len(summaries)} proyectos\n")
    print(f"{'Proyecto':<55} {'Sesiones':>8} {'Tokens':>14} {'Costo':>10} {'SubAg':>6}")
    print("-" * 97)

    for s in summaries[:30]:
        print(
            f"{s['project']:<55} {s['sessions']:>8,} {format_tokens(s['total_tokens']):>14} ${s['cost_usd']:>8,.2f} {s['subagent_count']:>6,}"
        )

    print("\nTop 10 sesiones más costosas:")
    for proj, session in find_costly_sessions(projects, top_n=10):
        ts = session["timestamp_start"][:10] if session["timestamp_start"] else "?"
        cost = estimate_cost_usd(session["usage"])
        first_prompt = ""
        if session["prompts"]:
            first_prompt = session["prompts"][0]["text"][:70].replace("\n", " ")
        print(f"  [{ts}] {proj}: {format_tokens(session['total_tokens'])} (${cost:,.2f}) — {first_prompt}")


def main():
    print("Analizando sesiones de Claude Code...")
    projects = analyze_all()

    print(f"Encontrados {len(projects)} proyectos")
    summaries = summarize_projects(projects)

    print_summary(summaries, projects)

    report_path = write_report(projects, summaries)
    write_prompts_by_project(projects)

    print(f"\nReporte completo: {report_path}")


if __name__ == "__main__":
    main()