#!/usr/bin/env python3
"""
Utilities for parsing and analyzing Claude Code JSONL conversation files.
"""

import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple


def parse_jsonl(file_path: Path) -> List[Dict[str, Any]]:
    """Parse JSONL file and return list of message objects."""
    messages = []
    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                messages.append(json.loads(line))
    return messages


def get_session_metadata(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Extract session metadata from messages.

    Returns:
        Dict with keys: session_id, cwd, git_branch, version, first_timestamp
    """
    if not messages:
        return {}

    first_msg = messages[0]
    return {
        "session_id": first_msg.get("sessionId", "unknown"),
        "cwd": first_msg.get("cwd", ""),
        "git_branch": first_msg.get("gitBranch", ""),
        "version": first_msg.get("version", ""),
        "first_timestamp": first_msg.get("timestamp", ""),
    }


def extract_token_usage(message: Dict[str, Any]) -> Optional[Dict[str, int]]:
    """
    Extract token usage from an assistant message.

    Returns:
        Dict with keys: input_tokens, output_tokens, cache_creation_input_tokens,
        cache_read_input_tokens, or None if not available
    """
    if message.get("type") != "assistant":
        return None

    usage = message.get("message", {}).get("usage", {})
    if not usage:
        return None

    cache_creation = usage.get("cache_creation", {})

    return {
        "input_tokens": usage.get("input_tokens", 0),
        "output_tokens": usage.get("output_tokens", 0),
        "cache_creation_input_tokens": usage.get("cache_creation_input_tokens", 0),
        "cache_read_input_tokens": usage.get("cache_read_input_tokens", 0),
        "ephemeral_5m_input_tokens": cache_creation.get("ephemeral_5m_input_tokens", 0),
        "ephemeral_1h_input_tokens": cache_creation.get("ephemeral_1h_input_tokens", 0),
    }


def extract_model_info(message: Dict[str, Any]) -> Optional[str]:
    """Extract model name from an assistant message."""
    if message.get("type") != "assistant":
        return None
    return message.get("message", {}).get("model")


def get_tool_uses(message: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Extract all tool_use blocks from an assistant message."""
    if message.get("type") != "assistant":
        return []

    content = message.get("message", {}).get("content", [])
    return [block for block in content if block.get("type") == "tool_use"]


def get_tool_results(message: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Extract all tool_result blocks from a user message."""
    if message.get("type") != "user":
        return []

    content = message.get("message", {}).get("content", [])
    if isinstance(content, str):
        return []

    return [block for block in content if block.get("type") == "tool_result"]


def get_text_content(message: Dict[str, Any]) -> str:
    """Extract text content from a message."""
    msg_data = message.get("message", {})
    content = msg_data.get("content", "")

    if isinstance(content, str):
        return content

    if isinstance(content, list):
        text_blocks = [block.get("text", "") for block in content if block.get("type") == "text"]
        return "\n\n".join(text_blocks)

    return ""


def calculate_relative_time(timestamp_str: str, base_timestamp_str: str) -> str:
    """
    Calculate relative time between two ISO timestamps.

    Args:
        timestamp_str: The timestamp to compare (later time)
        base_timestamp_str: The base timestamp (earlier time)

    Returns:
        Human-readable relative time string like "5m ago", "2h ago", "3d ago"
    """
    try:
        timestamp = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00"))
        base = datetime.fromisoformat(base_timestamp_str.replace("Z", "+00:00"))

        delta = timestamp - base
        total_seconds = int(delta.total_seconds())

        if total_seconds < 60:
            return "just now"
        elif total_seconds < 3600:
            minutes = total_seconds // 60
            return f"{minutes}m"
        elif total_seconds < 86400:
            hours = total_seconds // 3600
            return f"{hours}h"
        else:
            days = total_seconds // 86400
            return f"{days}d"
    except (ValueError, AttributeError):
        return ""


def get_relative_file_path(file_path: str, cwd: str) -> str:
    """
    Convert absolute file path to relative path from cwd.

    Args:
        file_path: Absolute file path
        cwd: Current working directory from session

    Returns:
        Relative path, or original path if conversion fails
    """
    try:
        abs_path = Path(file_path)
        cwd_path = Path(cwd)

        if abs_path.is_relative_to(cwd_path):
            return str(abs_path.relative_to(cwd_path))
        return file_path
    except (ValueError, AttributeError):
        return file_path


def detect_tool_success(tool_result: Dict[str, Any]) -> bool:
    """
    Detect if a tool execution was successful.

    Args:
        tool_result: Tool result block

    Returns:
        True if successful, False if error
    """
    return not tool_result.get("is_error", False)


def find_tool_use_by_id(messages: List[Dict[str, Any]], tool_use_id: str) -> Optional[Dict[str, Any]]:
    """
    Find the tool_use block that corresponds to a tool_result.

    Args:
        messages: List of all messages
        tool_use_id: The tool_use_id to search for

    Returns:
        The tool_use block, or None if not found
    """
    for message in messages:
        if message.get("type") == "assistant":
            tool_uses = get_tool_uses(message)
            for tool_use in tool_uses:
                if tool_use.get("id") == tool_use_id:
                    return tool_use
    return None


def calculate_total_tokens(messages: List[Dict[str, Any]]) -> Dict[str, int]:
    """
    Calculate total token usage across all messages.

    Returns:
        Dict with keys: total_input, total_output, total_cache_creation, total_cache_read
    """
    totals = {
        "total_input": 0,
        "total_output": 0,
        "total_cache_creation": 0,
        "total_cache_read": 0,
    }

    for message in messages:
        usage = extract_token_usage(message)
        if usage:
            totals["total_input"] += usage["input_tokens"]
            totals["total_output"] += usage["output_tokens"]
            totals["total_cache_creation"] += usage["cache_creation_input_tokens"]
            totals["total_cache_read"] += usage["cache_read_input_tokens"]

    return totals