#!/usr/bin/env python3
"""Extract last assistant output from current Claude session."""

import argparse
import json
import sys
from pathlib import Path


def encode_project_path(project_dir: str) -> str:
    """Convert absolute path to Claude history encoding format."""
    if not Path(project_dir).is_absolute():
        msg = "project_dir must be an absolute path"
        raise ValueError(msg)
    # Handle root path specially
    if project_dir == "/":
        return "-"
    return project_dir.rstrip("/").replace("/", "-")


def get_project_history_dir(project_dir: str) -> Path:
    """Return Path to ~/.claude/projects/[ENCODED-PATH]/."""
    return Path.home() / ".claude" / "projects" / encode_project_path(project_dir)


def find_latest_session(project_dir: Path) -> Path | None:
    """Find the most recently modified session file."""
    history_dir = get_project_history_dir(str(project_dir))

    if not history_dir.exists():
        return None

    # Find all UUID-named JSONL files
    import re
    uuid_pattern = re.compile(
        r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\.jsonl$"
    )

    session_files = [
        f for f in history_dir.glob("*.jsonl")
        if uuid_pattern.match(f.name)
    ]

    if not session_files:
        return None

    # Return most recently modified
    return max(session_files, key=lambda f: f.stat().st_mtime)


def extract_nth_assistant_output(session_file: Path, n: int = 1) -> str | None:
    """Extract text content from nth-last assistant message (1=most recent)."""
    with session_file.open() as f:
        lines = f.readlines()

    found = 0
    for line in reversed(lines):
        try:
            entry = json.loads(line)

            message = entry.get('message', {})
            if message.get('role') != 'assistant':
                continue

            content = message.get('content', [])
            if not isinstance(content, list):
                continue

            text_blocks = [
                c.get('text', '')
                for c in content
                if isinstance(c, dict) and c.get('type') == 'text'
            ]
            joined = '\n\n'.join(text_blocks).strip()

            if joined:
                found += 1
                if found == n:
                    return joined

        except (json.JSONDecodeError, KeyError, AttributeError):
            continue

    return None


def main():
    parser = argparse.ArgumentParser(
        description="Extract last assistant output from current Claude session"
    )
    parser.add_argument(
        "-o", "--output",
        help="Output file (default: stdout)",
        type=Path,
    )
    parser.add_argument(
        "n",
        nargs="?",
        type=int,
        default=1,
        help="Which output to retrieve: 1=latest (default), 2=previous, 3=before that",
    )
    parser.add_argument(
        "-d", "--directory",
        help="Project directory (default: current directory)",
        type=Path,
        default=Path.cwd(),
    )

    args = parser.parse_args()

    if args.n < 1:
        print("Error: n must be >= 1", file=sys.stderr)
        sys.exit(1)

    # Resolve to absolute path
    project_dir = args.directory.resolve()

    # Find latest session
    session_file = find_latest_session(project_dir)
    if not session_file:
        print("Error: No session files found", file=sys.stderr)
        sys.exit(1)

    # Extract nth output
    output = extract_nth_assistant_output(session_file, args.n)
    if not output:
        print(f"Error: No assistant message found at position {args.n}", file=sys.stderr)
        sys.exit(1)

    # Write to file or stdout
    if args.output:
        args.output.write_text(output)
        print(f"Wrote {len(output)} characters to {args.output}", file=sys.stderr)
    else:
        print(output)


if __name__ == "__main__":
    main()
