#!/usr/bin/env python3

from __future__ import annotations

import argparse
from dataclasses import dataclass, field
from datetime import datetime, timezone
import json
from pathlib import Path
from typing import Any


def parse_ts(value: str) -> datetime:
    if value.endswith("Z"):
        value = value[:-1] + "+00:00"
    return datetime.fromisoformat(value).astimezone(timezone.utc)


def format_seconds(seconds: float) -> str:
    if seconds < 1:
        return f"{seconds * 1000:.0f}ms"
    if seconds < 60:
        return f"{seconds:.2f}s"
    minutes = int(seconds // 60)
    rest = seconds - minutes * 60
    return f"{minutes}m {rest:.2f}s"


@dataclass
class ToolSpan:
    tool_call_id: str
    name: str
    started_at: datetime
    finished_at: datetime | None = None

    @property
    def duration_seconds(self) -> float | None:
        if self.finished_at is None:
            return None
        return (self.finished_at - self.started_at).total_seconds()


@dataclass
class Turn:
    index: int
    prompt_at: datetime
    prompt_preview: str
    first_assistant_at: datetime | None = None
    final_assistant_at: datetime | None = None
    tool_spans: list[ToolSpan] = field(default_factory=list)

    @property
    def prompt_to_first_assistant_seconds(self) -> float | None:
        if self.first_assistant_at is None:
            return None
        return (self.first_assistant_at - self.prompt_at).total_seconds()

    @property
    def total_turn_seconds(self) -> float | None:
        if self.final_assistant_at is None:
            return None
        return (self.final_assistant_at - self.prompt_at).total_seconds()


def preview_text(text: str, limit: int = 90) -> str:
    single = " ".join(text.strip().split())
    if len(single) <= limit:
        return single
    return single[: limit - 1] + "…"


def load_events(path: Path) -> list[dict[str, Any]]:
    events: list[dict[str, Any]] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        events.append(json.loads(line))
    return events


def build_turns(events: list[dict[str, Any]]) -> tuple[datetime | None, list[Turn], datetime | None]:
    session_started: datetime | None = None
    last_ts: datetime | None = None
    turns: list[Turn] = []
    current_turn: Turn | None = None
    open_tools: dict[str, ToolSpan] = {}

    for event in events:
        ts_raw = event.get("timestamp")
        if not isinstance(ts_raw, str):
            continue
        ts = parse_ts(ts_raw)
        last_ts = ts

        if event.get("type") == "session":
            session_started = ts
            continue

        if event.get("type") != "message":
            continue

        message = event.get("message") or {}
        role = message.get("role")
        content = message.get("content") or []

        if role == "user":
            text_parts = [item.get("text", "") for item in content if isinstance(item, dict) and item.get("type") == "text"]
            current_turn = Turn(
                index=len(turns) + 1,
                prompt_at=ts,
                prompt_preview=preview_text(" ".join(text_parts)),
            )
            turns.append(current_turn)
            open_tools = {}
            continue

        if current_turn is None:
            continue

        if role == "assistant":
            if current_turn.first_assistant_at is None:
                current_turn.first_assistant_at = ts

            has_text = any(isinstance(item, dict) and item.get("type") == "text" for item in content)
            for item in content:
                if not isinstance(item, dict) or item.get("type") != "toolCall":
                    continue
                tool_call_id = str(item.get("id", ""))
                name = str(item.get("name", "tool"))
                span = ToolSpan(tool_call_id=tool_call_id, name=name, started_at=ts)
                current_turn.tool_spans.append(span)
                open_tools[tool_call_id] = span

            if has_text:
                current_turn.final_assistant_at = ts
            continue

        if role == "toolResult":
            tool_call_id = str(message.get("toolCallId", ""))
            span = open_tools.get(tool_call_id)
            if span is not None and span.finished_at is None:
                span.finished_at = ts
            continue

    return session_started, turns, last_ts


def print_turn(turn: Turn) -> None:
    print(f"Turn {turn.index}: {turn.prompt_preview}")
    if turn.prompt_to_first_assistant_seconds is not None:
        print(f"  prompt -> first assistant: {format_seconds(turn.prompt_to_first_assistant_seconds)}")
    else:
        print("  prompt -> first assistant: n/a")

    if not turn.tool_spans:
        print("  tools: none")
    else:
        print("  tools:")
        last_tool_finished: datetime | None = None
        for span in turn.tool_spans:
            duration = span.duration_seconds
            status = format_seconds(duration) if duration is not None else "open"
            print(f"    - {span.name}: {status}")
            if span.finished_at is not None:
                if last_tool_finished is None or span.finished_at > last_tool_finished:
                    last_tool_finished = span.finished_at
        if last_tool_finished is not None and turn.final_assistant_at is not None:
            post_tools = (turn.final_assistant_at - last_tool_finished).total_seconds()
            print(f"  last tool result -> final assistant: {format_seconds(post_tools)}")

    if turn.total_turn_seconds is not None:
        print(f"  total turn: {format_seconds(turn.total_turn_seconds)}")
    else:
        print("  total turn: n/a")


def main() -> int:
    parser = argparse.ArgumentParser(prog="pi-session-timing.py")
    parser.add_argument("session", type=Path, nargs="+", help="Path(s) to Pi session.jsonl files")
    args = parser.parse_args()

    for path in args.session:
        events = load_events(path)
        session_started, turns, last_ts = build_turns(events)
        print(f"== {path} ==")
        if session_started is not None:
            print(f"session start: {session_started.isoformat()}")
        if turns:
            first_prompt = turns[0].prompt_at
            if session_started is not None:
                startup = (first_prompt - session_started).total_seconds()
                print(f"session start -> first prompt: {format_seconds(startup)}")
        if session_started is not None and last_ts is not None:
            print(f"session total: {format_seconds((last_ts - session_started).total_seconds())}")
        print(f"turns: {len(turns)}")
        print()
        for turn in turns:
            print_turn(turn)
            print()

    return 0


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