#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "pydantic",
#   "python-dotenv",
#   "click",
#   "rich",
#   "claude-code-sdk",
#   "anyio",
# ]
# ///
"""
Run Claude Code prompts using the official Python SDK.

This ADW demonstrates using the Claude Code Python SDK for both one-shot
and interactive sessions. The SDK provides better type safety, error handling,
and a more Pythonic interface compared to subprocess-based implementations.

Usage:
    # One-shot query (default)
    ./adws/adw_sdk_prompt.py "Hello Claude Code"

    # Interactive session
    ./adws/adw_sdk_prompt.py --interactive

    # Resume a previous session
    ./adws/adw_sdk_prompt.py --interactive --session-id abc123

    # With specific model
    ./adws/adw_sdk_prompt.py "Create a FastAPI app" --model opus

    # From different directory
    ./adws/adw_sdk_prompt.py "List files here" --working-dir /path/to/project

Examples:
    # Simple query
    ./adws/adw_sdk_prompt.py "Explain async/await in Python"

    # Interactive debugging session
    ./adws/adw_sdk_prompt.py --interactive --context "Debugging a memory leak"

    # Resume session with context
    ./adws/adw_sdk_prompt.py --interactive --session-id abc123 --context "Continue debugging"

    # Query with tools
    ./adws/adw_sdk_prompt.py "Create a Python web server" --tools Read,Write,Bash

Key Features:
    - Uses official Claude Code Python SDK
    - Supports both one-shot and interactive modes
    - Better error handling with typed exceptions
    - Native async/await support
    - Clean message type handling
"""

import os
import sys
import json
import asyncio
from pathlib import Path
from typing import Optional, List
import click
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.live import Live
from rich.spinner import Spinner
from rich.text import Text
from rich.prompt import Prompt

# Add the adw_modules directory to the path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "adw_modules"))

# Import SDK functions from our clean module
from agent_sdk import (
    simple_query,
    query_with_tools,
    collect_query_response,
    create_session,
    safe_query,
    stream_with_progress,
    extract_text,
    extract_tool_uses,
)

# Import SDK types
from claude_code_sdk import (
    ClaudeCodeOptions,
    AssistantMessage,
    ResultMessage,
    TextBlock,
    ToolUseBlock,
)


def generate_short_id() -> str:
    """Generate a short ID for tracking."""
    import uuid

    return str(uuid.uuid4())[:8]


async def run_one_shot_query(
    prompt: str,
    model: str,
    working_dir: str,
    allowed_tools: Optional[List[str]] = None,
    session_id: Optional[str] = None,
) -> None:
    """Run a one-shot query using the SDK."""
    console = Console()
    adw_id = generate_short_id()

    # Display execution info
    info_table = Table(show_header=False, box=None, padding=(0, 1))
    info_table.add_column(style="bold cyan")
    info_table.add_column()

    info_table.add_row("ADW ID", adw_id)
    info_table.add_row("Mode", "One-shot Query")
    info_table.add_row("Prompt", prompt)
    info_table.add_row("Model", model)
    info_table.add_row("Working Dir", working_dir)
    if allowed_tools:
        info_table.add_row("Tools", ", ".join(allowed_tools))
    if session_id:
        info_table.add_row("Session ID", session_id)
    info_table.add_row("[bold green]SDK[/bold green]", "Claude Code Python SDK")

    console.print(
        Panel(
            info_table,
            title="[bold blue]🚀 SDK Query Execution[/bold blue]",
            border_style="blue",
        )
    )
    console.print()

    try:
        # Execute query based on whether tools are needed
        with console.status("[bold yellow]Executing via SDK...[/bold yellow]"):
            if allowed_tools:
                # Query with tools
                options = ClaudeCodeOptions(
                    model=model,
                    allowed_tools=allowed_tools,
                    cwd=working_dir,
                    permission_mode="bypassPermissions",
                )
                if session_id:
                    options.resume = session_id
                messages, result = await collect_query_response(prompt, options=options)

                # Extract response text
                response_text = ""
                tool_uses = []

                for msg in messages:
                    if isinstance(msg, AssistantMessage):
                        text = extract_text(msg)
                        if text:
                            response_text += text + "\n"
                        for tool in extract_tool_uses(msg):
                            tool_uses.append(f"{tool.name} ({tool.id[:8]}...)")

                success = result and not result.is_error if result else False

            else:
                # Simple query
                response_text, error = await safe_query(prompt)
                success = error is None
                tool_uses = []

                if error:
                    response_text = error

        # Display result
        if success:
            console.print(
                Panel(
                    response_text.strip(),
                    title="[bold green]✅ SDK Success[/bold green]",
                    border_style="green",
                    padding=(1, 2),
                )
            )

            if tool_uses:
                console.print(
                    f"\n[bold cyan]Tools used:[/bold cyan] {', '.join(tool_uses)}"
                )
        else:
            console.print(
                Panel(
                    response_text,
                    title="[bold red]❌ SDK Error[/bold red]",
                    border_style="red",
                    padding=(1, 2),
                )
            )

        # Show cost and session info if available
        if "result" in locals() and result:
            if result.total_cost_usd:
                console.print(
                    f"\n[bold cyan]Cost:[/bold cyan] ${result.total_cost_usd:.4f}"
                )
            if hasattr(result, 'session_id') and result.session_id:
                console.print(
                    f"[bold cyan]Session ID:[/bold cyan] {result.session_id}"
                )
                console.print(
                    f"[dim]Resume with: --session-id {result.session_id}[/dim]"
                )

    except Exception as e:
        console.print(
            Panel(
                f"[bold red]{str(e)}[/bold red]",
                title="[bold red]❌ Unexpected Error[/bold red]",
                border_style="red",
            )
        )


async def run_interactive_session(
    model: str,
    working_dir: str,
    context: Optional[str] = None,
    session_id: Optional[str] = None,
) -> None:
    """Run an interactive session using the SDK."""
    console = Console()
    adw_id = generate_short_id()

    # Display session info
    info_table = Table(show_header=False, box=None, padding=(0, 1))
    info_table.add_column(style="bold cyan")
    info_table.add_column()

    info_table.add_row("ADW ID", adw_id)
    info_table.add_row("Mode", "Interactive Session")
    info_table.add_row("Model", model)
    info_table.add_row("Working Dir", working_dir)
    if context:
        info_table.add_row("Context", context)
    if session_id:
        info_table.add_row("Session ID", session_id)
    info_table.add_row("[bold green]SDK[/bold green]", "Claude Code Python SDK")

    console.print(
        Panel(
            info_table,
            title="[bold blue]💬 SDK Interactive Session[/bold blue]",
            border_style="blue",
        )
    )
    console.print()

    # Instructions
    console.print("[bold yellow]Interactive Mode[/bold yellow]")
    console.print("Commands: 'exit' or 'quit' to end session")
    console.print("Just type your questions or requests\n")

    # Start session
    options = ClaudeCodeOptions(
        model=model,
        cwd=working_dir,
        permission_mode="bypassPermissions",
    )
    if session_id:
        options.resume = session_id
    
    from claude_code_sdk import ClaudeSDKClient
    client = ClaudeSDKClient(options=options)
    await client.connect()
    
    # Track session ID from results throughout the session
    session_id_from_result = None
    
    try:
        # Send initial context if provided
        if context:
            console.print(f"[dim]Setting context: {context}[/dim]\n")
            await client.query(f"Context: {context}")

            # Consume the context response
            async for msg in client.receive_response():
                if isinstance(msg, AssistantMessage):
                    text = extract_text(msg)
                    if text:
                        console.print(f"[dim]Claude: {text}[/dim]\n")

        # Interactive loop
        while True:
            # Get user input
            try:
                user_input = Prompt.ask("[bold cyan]You[/bold cyan]")
            except (EOFError, KeyboardInterrupt):
                console.print("\n[yellow]Session interrupted[/yellow]")
                break

            if user_input.lower() in ["exit", "quit"]:
                break

            # Send to Claude
            await client.query(user_input)

            # Show response with progress
            console.print()
            response_parts = []
            tool_uses = []
            cost = None
            session_id_from_result = None

            with Live(
                Spinner("dots", text="Thinking..."),
                console=console,
                refresh_per_second=4,
            ):
                async for msg in client.receive_response():
                    if isinstance(msg, AssistantMessage):
                        text = extract_text(msg)
                        if text:
                            response_parts.append(text)

                        for tool in extract_tool_uses(msg):
                            tool_uses.append(f"{tool.name}")

                    elif isinstance(msg, ResultMessage):
                        if msg.total_cost_usd:
                            cost = msg.total_cost_usd
                        if hasattr(msg, 'session_id') and msg.session_id:
                            session_id_from_result = msg.session_id

            # Display response
            if response_parts:
                console.print("[bold green]Claude:[/bold green]")
                for part in response_parts:
                    console.print(part)

            if tool_uses:
                console.print(f"\n[dim]Tools used: {', '.join(tool_uses)}[/dim]")

            if cost:
                console.print(f"[dim]Cost: ${cost:.4f}[/dim]")
            
            if session_id_from_result:
                console.print(f"[dim]Session ID: {session_id_from_result}[/dim]")

            console.print()
    
    finally:
        await client.disconnect()

    console.print("\n[bold green]Session ended[/bold green]")
    console.print(f"[dim]ADW ID: {adw_id}[/dim]")
    if 'session_id_from_result' in locals() and session_id_from_result:
        console.print(f"[bold cyan]Session ID:[/bold cyan] {session_id_from_result}")
        console.print(f"[dim]Resume with: ./adws/adw_sdk_prompt.py --interactive --session-id {session_id_from_result}[/dim]")


@click.command()
@click.argument("prompt", required=False)
@click.option(
    "--interactive",
    "-i",
    is_flag=True,
    help="Start an interactive session instead of one-shot query",
)
@click.option(
    "--model",
    type=click.Choice(["sonnet", "opus", "haiku"]),
    default="sonnet",
    help="Claude model to use (sonnet=balanced, opus=max intelligence, haiku=fast & economical)",
)
@click.option(
    "--working-dir",
    type=click.Path(exists=True, file_okay=False, dir_okay=True, resolve_path=True),
    help="Working directory (default: current directory)",
)
@click.option(
    "--tools",
    help="Comma-separated list of allowed tools (e.g., Read,Write,Bash)",
)
@click.option(
    "--context",
    help="Context for interactive session (e.g., 'Debugging a memory leak')",
)
@click.option(
    "--session-id",
    help="Resume a previous session by its ID",
)
def main(
    prompt: Optional[str],
    interactive: bool,
    model: str,
    working_dir: Optional[str],
    tools: Optional[str],
    context: Optional[str],
    session_id: Optional[str],
):
    """Run Claude Code prompts using the Python SDK.

    Examples:
        # One-shot query
        adw_sdk_prompt.py "What is 2 + 2?"

        # Interactive session
        adw_sdk_prompt.py --interactive
        
        # Resume session
        adw_sdk_prompt.py --interactive --session-id abc123

        # Query with tools
        adw_sdk_prompt.py "Create hello.py" --tools Write,Read
    """
    if not working_dir:
        working_dir = os.getcwd()

    # Convert model names
    model_map = {
        "sonnet": "claude-sonnet-4-5-20250929",
        "opus": "claude-opus-4-20250514",
        "haiku": "claude-haiku-4-5-20251001"
    }
    full_model = model_map.get(model, model)

    # Parse tools if provided
    allowed_tools = None
    if tools:
        allowed_tools = [t.strip() for t in tools.split(",")]

    # Run appropriate mode
    if interactive:
        if prompt:
            console = Console()
            console.print(
                "[yellow]Warning: Prompt ignored in interactive mode[/yellow]\n"
            )

        asyncio.run(
            run_interactive_session(
                model=full_model,
                working_dir=working_dir,
                context=context,
                session_id=session_id,
            )
        )
    else:
        if not prompt:
            console = Console()
            console.print("[red]Error: Prompt required for one-shot mode[/red]")
            console.print("Use --interactive for interactive session")
            sys.exit(1)

        asyncio.run(
            run_one_shot_query(
                prompt=prompt,
                model=full_model,
                working_dir=working_dir,
                allowed_tools=allowed_tools,
                session_id=session_id,
            )
        )


if __name__ == "__main__":
    main()