"""
JSONL Parser for Claude Code session files.

Parses JSONL files from ~/.claude/projects/ and extracts
task chunks for triage.
"""

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


def parse_jsonl(jsonl_path: str) -> List[Dict]:
    """
    Parse JSONL file into list of entries.

    Args:
        jsonl_path: Path to JSONL file

    Returns:
        List of parsed JSON entries
    """
    entries = []
    path = Path(jsonl_path)

    if not path.exists():
        return entries

    with open(path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue

            try:
                entries.append(json.loads(line))
            except json.JSONDecodeError as e:
                print(f"[JSONL] Parse error on line {line_num}: {e}")
                continue

    return entries


def extract_chunk_after_timestamp(
    jsonl_path: str,
    start_timestamp: str
) -> List[Dict]:
    """
    Extract messages from start_timestamp to end of file.

    Args:
        jsonl_path: Path to JSONL file
        start_timestamp: ISO format timestamp to start from

    Returns:
        List of entries from start_timestamp onwards
    """
    entries = parse_jsonl(jsonl_path)
    chunk = []

    try:
        # Parse start time (handle various formats)
        start_time = parse_timestamp(start_timestamp)
    except ValueError as e:
        print(f"[JSONL] Invalid start timestamp: {e}")
        return []

    for entry in entries:
        # Skip queue operations
        if entry.get("type") == "queue-operation":
            continue

        # Get entry timestamp
        entry_ts = entry.get("timestamp")
        if not entry_ts:
            continue

        try:
            entry_time = parse_timestamp(entry_ts)
        except ValueError:
            continue

        # Include if at or after start time
        if entry_time >= start_time:
            chunk.append(entry)

    return chunk


def extract_chunks_by_user_prompts(entries: List[Dict]) -> List[List[Dict]]:
    """
    Split entries into task-based chunks.

    A new chunk starts with each user prompt (string content,
    not tool results which have array content).

    Args:
        entries: List of parsed JSONL entries

    Returns:
        List of chunks, each chunk is a list of entries
    """
    chunks = []
    current_chunk = []

    for entry in entries:
        # Skip queue operations
        if entry.get("type") == "queue-operation":
            continue

        # Check if this is a new user prompt
        if is_user_prompt(entry) and current_chunk:
            # Save current chunk and start new one
            chunks.append(current_chunk)
            current_chunk = []

        current_chunk.append(entry)

    # Don't forget last chunk
    if current_chunk:
        chunks.append(current_chunk)

    return chunks


def is_user_prompt(entry: Dict) -> bool:
    """
    Check if entry is a user prompt (not tool result).

    User prompts have string content.
    Tool results have array content.
    """
    if entry.get("type") != "user":
        return False

    content = entry.get("message", {}).get("content")
    return isinstance(content, str)


def is_task_complete(entry: Dict) -> bool:
    """Check if entry indicates task completion."""
    if entry.get("type") != "assistant":
        return False

    stop_reason = entry.get("message", {}).get("stop_reason")
    return stop_reason == "end_turn"


def chunk_to_transcript(chunk: List[Dict], task_name: str = "Unknown Task") -> str:
    """
    Convert a chunk to readable transcript for triage.

    Args:
        chunk: List of JSONL entries
        task_name: Name/description of the task

    Returns:
        Formatted transcript string
    """
    lines = [
        f"# Task: {task_name}",
        f"# Messages: {len(chunk)}",
        "",
        "---",
        ""
    ]

    for entry in chunk:
        formatted = format_entry(entry)
        if formatted:
            lines.append(formatted)

    return "\n".join(lines)


def format_entry(entry: Dict) -> Optional[str]:
    """
    Format a single JSONL entry for transcript.

    Args:
        entry: Single JSONL entry

    Returns:
        Formatted string, or None if entry should be skipped
    """
    entry_type = entry.get("type")
    message = entry.get("message", {})

    if entry_type == "user":
        return format_user_entry(message)

    elif entry_type == "assistant":
        return format_assistant_entry(message)

    return None


def format_user_entry(message: Dict) -> Optional[str]:
    """Format user message or tool result."""
    content = message.get("content")

    if isinstance(content, str):
        # User prompt
        # Truncate very long prompts
        if len(content) > 500:
            content = content[:500] + "\n... [truncated]"
        return f"**User**: {content}\n"

    elif isinstance(content, list):
        # Tool result(s)
        parts = []
        for block in content:
            if block.get("type") == "tool_result":
                is_error = block.get("is_error", False)
                result = block.get("content", "")

                # Truncate long results
                if len(result) > 1000:
                    result = result[:1000] + "\n... [truncated]"

                prefix = "❌ ERROR" if is_error else "✓ Result"
                parts.append(f"**{prefix}**:\n```\n{result}\n```\n")

        return "\n".join(parts) if parts else None

    return None


def format_assistant_entry(message: Dict) -> Optional[str]:
    """Format assistant message (text or tool use)."""
    content = message.get("content", [])
    parts = []

    for block in content:
        block_type = block.get("type")

        if block_type == "text":
            text = block.get("text", "")
            # Truncate very long text
            if len(text) > 1500:
                text = text[:1500] + "\n... [truncated]"
            parts.append(f"**Claude**: {text}\n")

        elif block_type == "tool_use":
            name = block.get("name", "Unknown")
            input_data = block.get("input", {})

            # Format input based on tool type
            input_str = format_tool_input(name, input_data)
            parts.append(f"**Tool [{name}]**:\n```\n{input_str}\n```\n")

    return "\n".join(parts) if parts else None


def format_tool_input(tool_name: str, input_data: Dict) -> str:
    """
    Format tool input for readability.

    Different tools show different relevant fields.
    """
    if tool_name == "Read":
        return input_data.get("file_path", str(input_data))

    elif tool_name == "Write":
        path = input_data.get("file_path", "unknown")
        content = input_data.get("content", "")
        preview = content[:200] + "..." if len(content) > 200 else content
        return f"File: {path}\nContent:\n{preview}"

    elif tool_name == "Edit":
        path = input_data.get("file_path", "unknown")
        old = input_data.get("old_string", "")[:100]
        new = input_data.get("new_string", "")[:100]
        return f"File: {path}\nOld: {old}\nNew: {new}"

    elif tool_name == "Bash":
        cmd = input_data.get("command", "")
        desc = input_data.get("description", "")
        return f"{desc}\n$ {cmd}" if desc else f"$ {cmd}"

    elif tool_name == "Glob":
        return input_data.get("pattern", str(input_data))

    elif tool_name == "Grep":
        pattern = input_data.get("pattern", "")
        path = input_data.get("path", ".")
        return f"Pattern: {pattern}\nPath: {path}"

    else:
        # Generic: show as compact JSON
        try:
            formatted = json.dumps(input_data, indent=2)
            if len(formatted) > 300:
                formatted = formatted[:300] + "\n..."
            return formatted
        except (TypeError, ValueError):
            return str(input_data)[:300]


def parse_timestamp(ts: str, make_naive: bool = True) -> datetime:
    """
    Parse various timestamp formats to datetime.

    Handles:
    - ISO format: 2026-01-13T21:52:45.524000
    - ISO with Z: 2026-01-13T21:52:45.524Z
    - ISO with timezone: 2026-01-13T21:52:45.524+00:00

    Args:
        ts: Timestamp string
        make_naive: If True, strip timezone info for consistent comparison
    """
    # Remove Z and replace with +00:00
    ts = ts.replace("Z", "+00:00")

    dt = None

    # Try parsing with timezone
    try:
        dt = datetime.fromisoformat(ts)
    except ValueError:
        pass

    # Try without microseconds
    if dt is None:
        try:
            if "." in ts:
                ts = ts.split(".")[0]
            dt = datetime.fromisoformat(ts)
        except ValueError:
            pass

    if dt is None:
        raise ValueError(f"Cannot parse timestamp: {ts}")

    # Strip timezone for consistent comparison (fixes naive vs aware comparison)
    if make_naive and dt.tzinfo is not None:
        dt = dt.replace(tzinfo=None)

    return dt


# CLI for testing
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="JSONL Parser")
    parser.add_argument("jsonl_path", help="Path to JSONL file")
    parser.add_argument("--start", help="Start timestamp (ISO format)")
    parser.add_argument("--chunks", action="store_true", help="Split into chunks")

    args = parser.parse_args()

    entries = parse_jsonl(args.jsonl_path)
    print(f"Parsed {len(entries)} entries")

    if args.start:
        chunk = extract_chunk_after_timestamp(args.jsonl_path, args.start)
        print(f"\nChunk from {args.start}: {len(chunk)} entries")
        transcript = chunk_to_transcript(chunk, "Test Task")
        print("\n--- TRANSCRIPT ---")
        print(transcript[:2000])

    elif args.chunks:
        chunks = extract_chunks_by_user_prompts(entries)
        print(f"\nFound {len(chunks)} chunks:")
        for i, chunk in enumerate(chunks):
            # Get first user prompt
            prompt = "Unknown"
            for entry in chunk:
                if is_user_prompt(entry):
                    prompt = entry.get("message", {}).get("content", "")[:50]
                    break
            print(f"  Chunk {i+1}: {len(chunk)} entries - {prompt}...")
