#!/usr/bin/env python3

import argparse
import json
import math
from collections import Counter
from pathlib import Path
from statistics import mean, median
from typing import Any


BUCKETS = ("input", "output", "cache_read", "cache_write")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Analyze Pi agent session logs for token mix and cache behavior."
    )
    parser.add_argument(
        "root",
        nargs="?",
        default="~/.pi/agent/sessions",
        help="Directory containing Pi session .jsonl files.",
    )
    parser.add_argument(
        "--json-out",
        help="Optional path to write the full report as JSON.",
    )
    return parser.parse_args()


def usage_value(usage: dict[str, Any], *keys: str) -> int:
    for key in keys:
        value = usage.get(key)
        if value is None:
            continue
        if isinstance(value, bool):
            return int(value)
        if isinstance(value, (int, float)):
            return int(value)
    return 0


def safe_div(numerator: float, denominator: float) -> float | None:
    if denominator == 0:
        return None
    return numerator / denominator


def percentile(values: list[float], p: float) -> float | None:
    if not values:
        return None
    if len(values) == 1:
        return values[0]
    ordered = sorted(values)
    rank = (len(ordered) - 1) * p
    low = math.floor(rank)
    high = math.ceil(rank)
    if low == high:
        return ordered[low]
    weight = rank - low
    return ordered[low] * (1 - weight) + ordered[high] * weight


def summarize(values: list[float]) -> dict[str, float | None]:
    return {
        "count": len(values),
        "mean": mean(values) if values else None,
        "median": median(values) if values else None,
        "p10": percentile(values, 0.10),
        "p90": percentile(values, 0.90),
        "min": min(values) if values else None,
        "max": max(values) if values else None,
    }


def session_primary(counter: Counter[str]) -> str | None:
    if not counter:
        return None
    return max(counter.items(), key=lambda item: (item[1], item[0]))[0]


def format_number(value: float | None, digits: int = 3) -> str:
    if value is None:
        return "n/a"
    return f"{value:,.{digits}f}"


def format_pct(value: float | None, digits: int = 1) -> str:
    if value is None:
        return "n/a"
    return f"{value * 100:.{digits}f}%"


def analyze_session(path: Path) -> dict[str, Any]:
    totals = {bucket: 0 for bucket in BUCKETS}
    usage_messages = 0
    assistant_messages = 0
    models = Counter()
    providers = Counter()
    apis = Counter()
    line_count = 0
    parse_errors = 0

    with path.open("r", encoding="utf-8") as handle:
        for raw_line in handle:
            line_count += 1
            line = raw_line.strip()
            if not line:
                continue
            try:
                entry = json.loads(line)
            except json.JSONDecodeError:
                parse_errors += 1
                continue

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

            role = message.get("role")
            if role == "assistant":
                assistant_messages += 1

            usage = entry.get("usage") or message.get("usage")
            if role != "assistant" or not isinstance(usage, dict):
                continue

            usage_messages += 1
            bucket_values = {
                "input": usage_value(usage, "input", "input_tokens"),
                "output": usage_value(usage, "output", "output_tokens"),
                "cache_read": usage_value(
                    usage,
                    "cacheRead",
                    "cache_read",
                    "cache_read_input_tokens",
                ),
                "cache_write": usage_value(
                    usage,
                    "cacheWrite",
                    "cache_write",
                    "cache_creation_input_tokens",
                ),
            }
            for bucket, value in bucket_values.items():
                totals[bucket] += value

            billable = sum(bucket_values.values())
            if billable <= 0:
                continue

            model = (
                entry.get("responseModel")
                or entry.get("model")
                or message.get("responseModel")
                or message.get("model")
            )
            provider = entry.get("provider") or message.get("provider")
            api = entry.get("api") or message.get("api")
            if isinstance(model, str):
                models[model] += billable
            if isinstance(provider, str):
                providers[provider] += billable
            if isinstance(api, str):
                apis[api] += billable

    prompt_tokens = totals["input"] + totals["cache_read"] + totals["cache_write"]
    billable_tokens = prompt_tokens + totals["output"]
    cache_activity = totals["cache_read"] + totals["cache_write"]

    shares = {
        bucket: safe_div(totals[bucket], billable_tokens) for bucket in BUCKETS
    }

    return {
        "path": str(path),
        "lines": line_count,
        "parse_errors": parse_errors,
        "assistant_messages": assistant_messages,
        "usage_messages": usage_messages,
        "totals": totals,
        "prompt_tokens": prompt_tokens,
        "billable_tokens": billable_tokens,
        "cache_activity_tokens": cache_activity,
        "primary_model": session_primary(models),
        "primary_provider": session_primary(providers),
        "primary_api": session_primary(apis),
        "metrics": {
            "prompt_to_output_ratio": safe_div(prompt_tokens, totals["output"]),
            "input_to_output_ratio": safe_div(totals["input"], totals["output"]),
            "cache_hit_rate": safe_div(totals["cache_read"], cache_activity),
            "cache_coverage": safe_div(cache_activity, prompt_tokens),
            "prompt_reuse_rate": safe_div(totals["cache_read"], prompt_tokens),
            "output_share": safe_div(totals["output"], billable_tokens),
            "prompt_share": safe_div(prompt_tokens, billable_tokens),
            **{f"{bucket}_share": value for bucket, value in shares.items()},
        },
    }


def build_report(root: Path) -> dict[str, Any]:
    session_files = sorted(root.rglob("*.jsonl"))
    sessions = [analyze_session(path) for path in session_files]
    analyzed_sessions = [s for s in sessions if s["billable_tokens"] > 0]

    pooled_totals = {bucket: 0 for bucket in BUCKETS}
    total_prompt = 0
    total_billable = 0
    total_usage_messages = 0
    total_assistant_messages = 0
    total_parse_errors = 0
    provider_sessions = Counter()
    model_sessions = Counter()

    per_session_metric_values: dict[str, list[float]] = {}
    cache_active_metric_values: dict[str, list[float]] = {}
    cache_active_sessions = 0

    for session in analyzed_sessions:
        for bucket in BUCKETS:
            pooled_totals[bucket] += session["totals"][bucket]
        total_prompt += session["prompt_tokens"]
        total_billable += session["billable_tokens"]
        total_usage_messages += session["usage_messages"]
        total_assistant_messages += session["assistant_messages"]
        total_parse_errors += session["parse_errors"]
        if session["primary_provider"]:
            provider_sessions[session["primary_provider"]] += 1
        if session["primary_model"]:
            model_sessions[session["primary_model"]] += 1
        for name, value in session["metrics"].items():
            if value is None:
                continue
            per_session_metric_values.setdefault(name, []).append(value)
            if session["cache_activity_tokens"] > 0 and name in {
                "cache_hit_rate",
                "cache_coverage",
                "prompt_reuse_rate",
                "cache_read_share",
                "cache_write_share",
            }:
                cache_active_metric_values.setdefault(name, []).append(value)
        per_session_metric_values.setdefault("billable_tokens", []).append(
            session["billable_tokens"]
        )
        per_session_metric_values.setdefault("prompt_tokens", []).append(
            session["prompt_tokens"]
        )
        per_session_metric_values.setdefault("output_tokens", []).append(
            session["totals"]["output"]
        )
        if session["cache_activity_tokens"] > 0:
            cache_active_sessions += 1

    pooled_cache_activity = pooled_totals["cache_read"] + pooled_totals["cache_write"]
    pooled_weights = {
        bucket: safe_div(pooled_totals[bucket], total_billable) for bucket in BUCKETS
    }

    mean_session_weights = {
        bucket: mean(per_session_metric_values.get(f"{bucket}_share", []))
        if per_session_metric_values.get(f"{bucket}_share")
        else None
        for bucket in BUCKETS
    }
    median_session_weights = {
        bucket: median(per_session_metric_values.get(f"{bucket}_share", []))
        if per_session_metric_values.get(f"{bucket}_share")
        else None
        for bucket in BUCKETS
    }

    return {
        "root": str(root),
        "sessions_scanned": len(session_files),
        "sessions_with_usage": len(analyzed_sessions),
        "sessions_without_usage": len(session_files) - len(analyzed_sessions),
        "assistant_messages": total_assistant_messages,
        "assistant_messages_with_usage": total_usage_messages,
        "parse_errors": total_parse_errors,
        "cache_active_sessions": cache_active_sessions,
        "pooled_totals": pooled_totals,
        "pooled_metrics": {
            "prompt_tokens": total_prompt,
            "billable_tokens": total_billable,
            "prompt_to_output_ratio": safe_div(total_prompt, pooled_totals["output"]),
            "input_to_output_ratio": safe_div(
                pooled_totals["input"], pooled_totals["output"]
            ),
            "cache_hit_rate": safe_div(
                pooled_totals["cache_read"], pooled_cache_activity
            ),
            "cache_coverage": safe_div(pooled_cache_activity, total_prompt),
            "prompt_reuse_rate": safe_div(pooled_totals["cache_read"], total_prompt),
            "output_share": safe_div(pooled_totals["output"], total_billable),
            "prompt_share": safe_div(total_prompt, total_billable),
            **{f"{bucket}_share": value for bucket, value in pooled_weights.items()},
        },
        "recommended_weights": {
            "pooled_token_mix": pooled_weights,
            "mean_session_mix": mean_session_weights,
            "median_session_mix": median_session_weights,
        },
        "per_session_distributions": {
            name: summarize(values)
            for name, values in sorted(per_session_metric_values.items())
        },
        "cache_active_session_distributions": {
            name: summarize(values)
            for name, values in sorted(cache_active_metric_values.items())
        },
        "top_primary_providers": provider_sessions.most_common(10),
        "top_primary_models": model_sessions.most_common(10),
    }


def print_report(report: dict[str, Any]) -> None:
    print("Pi Session Cost Analysis")
    print("========================")
    print(f"Root: {report['root']}")
    print(
        f"Sessions scanned: {report['sessions_scanned']} "
        f"({report['sessions_with_usage']} with usage)"
    )
    print(
        f"Assistant messages with usage: {report['assistant_messages_with_usage']} "
        f"of {report['assistant_messages']}"
    )
    print(f"Cache-active sessions: {report['cache_active_sessions']}")
    print()

    pooled = report["pooled_metrics"]
    totals = report["pooled_totals"]
    weights = report["recommended_weights"]

    print("Pooled Totals")
    print("-------------")
    print(f"Input:       {totals['input']:,}")
    print(f"Output:      {totals['output']:,}")
    print(f"Cache read:  {totals['cache_read']:,}")
    print(f"Cache write: {totals['cache_write']:,}")
    print(f"Billable:    {pooled['billable_tokens']:,}")
    print()

    print("Recommended Weight Vectors")
    print("--------------------------")
    for label, vector in (
        ("Pooled token mix", weights["pooled_token_mix"]),
        ("Mean session mix", weights["mean_session_mix"]),
        ("Median session mix", weights["median_session_mix"]),
    ):
        print(
            f"{label:18} "
            f"input={format_pct(vector['input'])}  "
            f"output={format_pct(vector['output'])}  "
            f"cache_read={format_pct(vector['cache_read'])}  "
            f"cache_write={format_pct(vector['cache_write'])}"
        )
    print()

    print("Key Corpus Metrics")
    print("------------------")
    print(
        f"Prompt/output ratio: {format_number(pooled['prompt_to_output_ratio'])}x "
        f"(prompt = input + cache_read + cache_write)"
    )
    print(
        f"Uncached input/output ratio: {format_number(pooled['input_to_output_ratio'])}x"
    )
    print(f"Cache hit rate: {format_pct(pooled['cache_hit_rate'])}")
    print(f"Cache coverage of prompt side: {format_pct(pooled['cache_coverage'])}")
    print(f"Prompt reuse rate: {format_pct(pooled['prompt_reuse_rate'])}")
    print()

    print("Per-Session Distributions")
    print("-------------------------")
    metric_names = [
        "billable_tokens",
        "prompt_to_output_ratio",
        "input_to_output_ratio",
        "cache_hit_rate",
        "cache_coverage",
        "prompt_reuse_rate",
        "input_share",
        "output_share",
        "cache_read_share",
        "cache_write_share",
    ]
    for name in metric_names:
        stats = report["per_session_distributions"].get(name)
        if not stats:
            continue
        unit = "%" if (
            name.endswith("_share")
            or "rate" in name
            or "coverage" in name
        ) else ""
        value_fmt = format_pct if unit == "%" else format_number
        print(
            f"{name:22} "
            f"mean={value_fmt(stats['mean'])}  "
            f"median={value_fmt(stats['median'])}  "
            f"p10={value_fmt(stats['p10'])}  "
            f"p90={value_fmt(stats['p90'])}"
        )
    print()

    if report["cache_active_session_distributions"]:
        print("Cache-Active Session Distributions")
        print("----------------------------------")
        for name in [
            "cache_hit_rate",
            "cache_coverage",
            "prompt_reuse_rate",
            "cache_read_share",
            "cache_write_share",
        ]:
            stats = report["cache_active_session_distributions"].get(name)
            if not stats:
                continue
            print(
                f"{name:22} "
                f"mean={format_pct(stats['mean'])}  "
                f"median={format_pct(stats['median'])}  "
                f"p10={format_pct(stats['p10'])}  "
                f"p90={format_pct(stats['p90'])}"
            )
        print()

    print("Top Primary Providers")
    print("---------------------")
    for provider, count in report["top_primary_providers"]:
        print(f"{provider:20} {count}")
    print()

    print("Top Primary Models")
    print("------------------")
    for model, count in report["top_primary_models"]:
        print(f"{model:35} {count}")


def main() -> int:
    args = parse_args()
    root = Path(args.root).expanduser().resolve()
    report = build_report(root)

    if args.json_out:
        out_path = Path(args.json_out).expanduser().resolve()
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_text(json.dumps(report, indent=2), encoding="utf-8")

    print_report(report)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())