"""
ML-Ralph CLI

Commands:
- ml-ralph init: Initialize Ralph in project
- ml-ralph run: Run the autonomous loop

Supports two tools:
- claude: Claude Code CLI with streaming JSON output
- codex: OpenAI Codex CLI with JSONL streaming (via `codex exec --json`)
"""

import json
import subprocess
from pathlib import Path
from typing import Literal

import typer
from rich.console import Console

app = typer.Typer(help="ML-Ralph - Autonomous ML Agent")
console = Console()

# Type alias for supported tools
ToolType = Literal["claude", "codex"]

# Paths
ML_RALPH_DIR = Path(".ml-ralph")
TEMPLATES_DIR = Path(__file__).parent / "templates"


@app.command()
def init(
    force: bool = typer.Option(False, "--force", "-f", help="Overwrite existing files"),
):
    """Initialize Ralph in the current project."""
    import shutil

    # Check for conflicts
    claude_md = Path("CLAUDE.md")
    agents_md = Path("AGENTS.md")

    if not force and (claude_md.exists() or agents_md.exists()):
        console.print(
            "[red]CLAUDE.md or AGENTS.md already exists. Use --force to overwrite.[/red]"
        )
        raise typer.Exit(1)

    # Create .ml-ralph directory
    ML_RALPH_DIR.mkdir(exist_ok=True)

    # Copy templates
    templates = [
        (TEMPLATES_DIR / "RALPH.md", ML_RALPH_DIR / "RALPH.md"),
        (TEMPLATES_DIR / "CLAUDE.md", claude_md),
        (TEMPLATES_DIR / "AGENTS.md", agents_md),
    ]

    for src, dst in templates:
        if src.exists():
            dst.write_text(src.read_text())
            console.print(f"[dim]Created {dst}[/dim]")

    # Install skills
    skills_src = TEMPLATES_DIR / "skills"
    if skills_src.exists():
        for target in [Path(".claude/skills"), Path(".codex/skills")]:
            target.mkdir(parents=True, exist_ok=True)
            for skill in skills_src.iterdir():
                if skill.is_dir():
                    dest = target / skill.name
                    if dest.exists():
                        shutil.rmtree(dest)
                    shutil.copytree(skill, dest)
            console.print(f"[dim]Installed skills to {target}[/dim]")

    console.print("\n[green]Ralph initialized![/green]")
    console.print(
        "Use /ml-ralph in Claude Code to create a PRD, then run: ml-ralph run"
    )


@app.command()
def run(
    tool: ToolType = typer.Option("claude", help="Tool: claude or codex"),
    max_iterations: int = typer.Option(100, help="Max iterations"),
    sandbox: str = typer.Option(
        "workspace-write",
        help="Codex sandbox mode: read-only, workspace-write, danger-full-access",
    ),
):
    """Run the autonomous loop.

    Examples:
        ml-ralph run --tool claude
        ml-ralph run --tool codex
        ml-ralph run --tool codex --sandbox danger-full-access
    """
    ralph_md = ML_RALPH_DIR / "RALPH.md"

    if not ralph_md.exists():
        console.print("[red]Not initialized. Run 'ml-ralph init' first.[/red]")
        raise typer.Exit(1)

    tool_display = f"[cyan]{tool}[/cyan]"
    if tool == "codex":
        tool_display += f" [dim](sandbox: {sandbox})[/dim]"
    console.print(
        f"Starting Ralph loop with {tool_display} ({max_iterations} max iterations)"
    )

    prompt = """Read .ml-ralph/RALPH.md for instructions.

Execute one iteration of the cognitive loop. Update state files as needed.
When done, output exactly: <iteration_complete>

If the project is complete (success criteria met), output: <project_complete>"""

    for i in range(1, max_iterations + 1):
        console.print(f"\n[yellow]{'═' * 50}[/yellow]")
        console.print(f"[yellow]═══ Iteration {i} ═══[/yellow]")
        console.print(f"[yellow]{'═' * 50}[/yellow]\n")

        result = _run_agent_streaming(tool, prompt, sandbox=sandbox)

        if "<project_complete>" in result:
            console.print("\n[green]═══ Project complete! ═══[/green]")
            break
    else:
        console.print(f"\n[yellow]Reached max iterations ({max_iterations})[/yellow]")


def _run_agent_streaming(
    tool: ToolType, prompt: str, sandbox: str = "workspace-write"
) -> str:
    """Run the agent with streaming output.

    Args:
        tool: Which CLI to use (claude or codex)
        prompt: The prompt to send to the agent
        sandbox: Codex sandbox mode (ignored for claude)

    Returns:
        The final result text for completion detection
    """
    if tool == "claude":
        return _stream_claude(prompt)
    elif tool == "codex":
        return _stream_codex(prompt, sandbox)
    else:
        console.print(f"[red]Unknown tool: {tool}[/red]")
        raise typer.Exit(1)


def _stream_claude(prompt: str) -> str:
    """Stream Claude Code CLI output."""
    cmd = [
        "claude",
        "-p",
        prompt,
        "--output-format",
        "stream-json",
        "--allowedTools",
        "Bash,Read,Write,Edit,Glob,Grep",
    ]

    full_result = ""
    try:
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
        )

        for line in process.stdout:
            line = line.strip()
            if not line:
                continue

            try:
                event = json.loads(line)
                _handle_claude_event(event)

                # Capture final result
                if event.get("type") == "result":
                    full_result = event.get("result", "")

            except json.JSONDecodeError:
                # Not JSON, just print it
                print(line)

        process.wait()
        return full_result

    except FileNotFoundError:
        console.print(
            "[red]Claude CLI not found. Install with: npm install -g @anthropic-ai/claude-code[/red]"
        )
        raise typer.Exit(1)


def _stream_codex(prompt: str, sandbox: str) -> str:
    """Stream Codex CLI output using `codex exec --json`.

    Codex event types:
    - thread.started: Session start with thread_id
    - turn.started: New turn begins
    - item.started/item.completed: Items with types:
        - reasoning: Model thinking
        - file_change: File operations
        - command_execution: Shell commands
        - agent_message: Agent output text
    - turn.completed: Turn ends with usage stats
    """
    # Check if we're in a git repo
    git_check = subprocess.run(
        ["git", "rev-parse", "--git-dir"],
        capture_output=True,
        text=True,
    )
    in_git_repo = git_check.returncode == 0

    cmd = [
        "codex",
        "exec",
        "--json",
        "--full-auto",
        "--sandbox",
        sandbox,
    ]

    # Add skip-git-repo-check if not in a git repo
    if not in_git_repo:
        cmd.append("--skip-git-repo-check")

    cmd.append(prompt)

    full_result = ""
    try:
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
        )

        for line in process.stdout:
            line = line.strip()
            if not line:
                continue

            try:
                event = json.loads(line)
                message_text = _handle_codex_event(event)
                if message_text:
                    full_result = message_text

            except json.JSONDecodeError:
                # Not JSON, just print it
                print(line)

        process.wait()
        return full_result

    except FileNotFoundError:
        console.print(
            "[red]Codex CLI not found. Install with: npm install -g @openai/codex[/red]"
        )
        raise typer.Exit(1)


def _handle_claude_event(event: dict):
    """Handle a streaming JSON event from Claude Code.

    Event types:
    - assistant: Contains text or tool_use blocks
    - user: Contains tool_result blocks
    - result: Final result
    """
    event_type = event.get("type", "")

    if event_type == "assistant":
        # Assistant message - can contain text or tool_use
        message = event.get("message", {})
        content = message.get("content", [])

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

            if block_type == "text":
                text = block.get("text", "")
                if text:
                    print(text)

            elif block_type == "tool_use":
                tool_name = block.get("name", "unknown")
                tool_input = block.get("input", {})

                # Show tool being called
                if tool_name == "Bash":
                    cmd = tool_input.get("command", "")
                    desc = tool_input.get("description", "")
                    if desc:
                        console.print(f"\n[dim]► {tool_name}: {desc}[/dim]")
                    else:
                        console.print(
                            f"\n[dim]► {tool_name}: {cmd[:50]}...[/dim]"
                            if len(cmd) > 50
                            else f"\n[dim]► {tool_name}: {cmd}[/dim]"
                        )
                elif tool_name in ("Read", "Write", "Edit"):
                    file_path = tool_input.get("file_path", "")
                    console.print(f"\n[dim]► {tool_name}: {file_path}[/dim]")
                elif tool_name == "Glob":
                    pattern = tool_input.get("pattern", "")
                    console.print(f"\n[dim]► {tool_name}: {pattern}[/dim]")
                elif tool_name == "Grep":
                    pattern = tool_input.get("pattern", "")
                    console.print(f"\n[dim]► {tool_name}: {pattern}[/dim]")
                else:
                    console.print(f"\n[dim]► {tool_name}[/dim]")

    elif event_type == "user":
        # Tool result
        message = event.get("message", {})
        content = message.get("content", [])

        for block in content:
            if block.get("type") == "tool_result":
                is_error = block.get("is_error", False)
                if is_error:
                    console.print(" [red]✗[/red]")
                else:
                    console.print(" [green]✓[/green]")

    elif event_type == "result":
        # Final result - nothing to show
        pass


def _handle_codex_event(event: dict) -> str | None:
    """Handle a streaming JSONL event from Codex CLI.

    Event types:
    - thread.started: Session start with thread_id
    - turn.started: New turn begins
    - item.started: Item in progress (e.g., command running)
    - item.completed: Item finished with results
    - turn.completed: Turn ends with usage stats

    Item types (in item.started/item.completed):
    - reasoning: Model thinking process
    - file_change: File add/modify/delete
    - command_execution: Shell command with output
    - agent_message: Final agent response text

    Returns:
        The agent_message text if this is an agent_message item, else None
    """
    event_type = event.get("type", "")

    if event_type == "thread.started":
        thread_id = event.get("thread_id", "unknown")
        console.print(f"[dim]Session: {thread_id[:8]}...[/dim]")

    elif event_type == "turn.started":
        pass  # Nothing to show

    elif event_type in ("item.started", "item.completed"):
        item = event.get("item", {})
        item_type = item.get("type", "")
        status = item.get("status", "")

        if item_type == "reasoning":
            # Show model thinking (only on completed to avoid duplicates)
            if event_type == "item.completed":
                text = item.get("text", "")
                if text:
                    # Show abbreviated reasoning
                    lines = text.strip().split("\n")
                    preview = (
                        lines[0][:100] + "..." if len(lines[0]) > 100 else lines[0]
                    )
                    console.print(f"[dim italic]💭 {preview}[/dim italic]")

        elif item_type == "file_change":
            # Show file changes
            if event_type == "item.completed":
                changes = item.get("changes", [])
                for change in changes:
                    path = change.get("path", "unknown")
                    kind = change.get("kind", "modify")
                    icon = {"add": "✚", "modify": "✎", "delete": "✖"}.get(kind, "•")
                    console.print(f"[dim]► File {icon}: {path}[/dim]")
                if status == "completed":
                    console.print(" [green]✓[/green]")

        elif item_type == "command_execution":
            command = item.get("command", "")

            if event_type == "item.started":
                # Show command being run
                # Strip shell wrapper if present
                if command.startswith("/bin/"):
                    # Extract actual command from wrapper like "/bin/zsh -lc 'actual command'"
                    if "'" in command:
                        actual = command.split("'", 1)[1].rsplit("'", 1)[0]
                        command = actual

                cmd_preview = command[:60] + "..." if len(command) > 60 else command
                console.print(f"\n[dim]► Bash: {cmd_preview}[/dim]")

            elif event_type == "item.completed":
                exit_code = item.get("exit_code")
                if exit_code == 0:
                    console.print(" [green]✓[/green]")
                else:
                    console.print(f" [red]✗ (exit {exit_code})[/red]")

        elif item_type == "agent_message":
            # Final agent message - print it
            if event_type == "item.completed":
                text = item.get("text", "")
                if text:
                    print(text)
                    return text

    elif event_type == "turn.completed":
        usage = event.get("usage", {})
        input_tokens = usage.get("input_tokens", 0)
        output_tokens = usage.get("output_tokens", 0)
        cached = usage.get("cached_input_tokens", 0)
        console.print(
            f"\n[dim]Tokens: {input_tokens:,} in ({cached:,} cached) / {output_tokens:,} out[/dim]"
        )

    return None


if __name__ == "__main__":
    app()
