"""
Interactive REPL for zwarm session management.

A clean, autocomplete-enabled interface for managing codex sessions.
This is the user's direct REPL over the session primitives.

Topology: interactive → CodexSessionManager (substrate)
"""

from __future__ import annotations

import shlex
import time
from datetime import datetime
from pathlib import Path
from typing import Any

from prompt_toolkit import PromptSession
from prompt_toolkit.completion import Completer, Completion
from prompt_toolkit.history import InMemoryHistory
from prompt_toolkit.styles import Style
from rich.console import Console
from rich.table import Table

console = Console()


# =============================================================================
# Session ID Completer
# =============================================================================


class SessionCompleter(Completer):
    """
    Autocomplete for session IDs.

    Provides completions for commands that take session IDs.
    """

    def __init__(self, get_sessions_fn):
        """
        Args:
            get_sessions_fn: Callable that returns list of sessions
        """
        self.get_sessions_fn = get_sessions_fn

        # Commands that take session ID as first argument
        self.session_commands = {
            "?", "peek", "show", "traj", "trajectory", "watch",
            "c", "continue",
        }
        # Commands that take session ID OR "all"
        self.session_or_all_commands = {"kill", "rm", "delete"}

    def get_completions(self, document, complete_event):
        text = document.text_before_cursor
        words = text.split()

        # If we're completing the first word, suggest commands
        if len(words) == 0 or (len(words) == 1 and not text.endswith(" ")):
            word = words[0] if words else ""
            commands = [
                "spawn", "ls", "peek", "show", "traj", "watch",
                "c", "kill", "rm", "help", "quit",
            ]
            for cmd in commands:
                if cmd.startswith(word.lower()):
                    yield Completion(cmd, start_position=-len(word))
            return

        # If we have a command and need session ID
        cmd = words[0].lower()
        needs_session = cmd in self.session_commands or cmd in self.session_or_all_commands

        if needs_session:
            # Get sessions
            try:
                sessions = self.get_sessions_fn()
            except Exception:
                return

            # What has user typed for session ID?
            if len(words) == 1 and text.endswith(" "):
                # Just typed command + space, show all IDs
                partial = ""
            elif len(words) == 2 and not text.endswith(" "):
                # Typing session ID
                partial = words[1]
            else:
                return

            # For kill/rm, also offer "all" as option
            if cmd in self.session_or_all_commands:
                if "all".startswith(partial.lower()):
                    yield Completion(
                        "all",
                        start_position=-len(partial),
                        display="all",
                        display_meta="all sessions",
                    )

            # Yield matching session IDs
            for s in sessions:
                short_id = s.short_id
                if short_id.lower().startswith(partial.lower()):
                    # Show task as meta info
                    task_preview = s.task[:30] + "..." if len(s.task) > 30 else s.task
                    yield Completion(
                        short_id,
                        start_position=-len(partial),
                        display=short_id,
                        display_meta=f"{s.status.value}: {task_preview}",
                    )


# =============================================================================
# Display Helpers
# =============================================================================


def time_ago(iso_str: str) -> str:
    """Convert ISO timestamp to human-readable 'Xs/Xm/Xh ago' format."""
    try:
        dt = datetime.fromisoformat(iso_str)
        delta = datetime.now() - dt
        secs = delta.total_seconds()
        if secs < 60:
            return f"{int(secs)}s"
        elif secs < 3600:
            return f"{int(secs/60)}m"
        elif secs < 86400:
            return f"{secs/3600:.1f}h"
        else:
            return f"{secs/86400:.1f}d"
    except Exception:
        return "?"


STATUS_ICONS = {
    "running": "[yellow]●[/]",
    "completed": "[green]✓[/]",
    "failed": "[red]✗[/]",
    "killed": "[dim]○[/]",
    "pending": "[dim]◌[/]",
}


# =============================================================================
# Commands
# =============================================================================


def cmd_help():
    """Show help."""
    table = Table(show_header=False, box=None, padding=(0, 2))
    table.add_column("Command", style="cyan", width=30)
    table.add_column("Description")

    table.add_row("[bold]Session Lifecycle[/]", "")
    table.add_row('spawn "task" [--dir PATH]', "Start new session")
    table.add_row('c ID "message"', "Continue conversation")
    table.add_row("kill ID | all", "Stop session(s)")
    table.add_row("rm ID | all", "Delete session(s)")
    table.add_row("", "")
    table.add_row("[bold]Viewing[/]", "")
    table.add_row("ls", "Dashboard of all sessions")
    table.add_row("? ID  /  peek ID", "Quick peek (status + latest)")
    table.add_row("show ID", "Full session details")
    table.add_row("traj ID [--full]", "Show trajectory (steps taken)")
    table.add_row("watch ID", "Live follow session output")
    table.add_row("", "")
    table.add_row("[bold]Meta[/]", "")
    table.add_row("help", "Show this help")
    table.add_row("quit", "Exit")

    console.print(table)


def cmd_ls(manager):
    """List all sessions."""
    from zwarm.sessions import SessionStatus
    from zwarm.core.costs import estimate_session_cost, format_cost

    sessions = manager.list_sessions()

    if not sessions:
        console.print("  [dim]No sessions. Use 'spawn \"task\"' to start one.[/]")
        return

    # Summary counts
    running = sum(1 for s in sessions if s.status == SessionStatus.RUNNING)
    completed = sum(1 for s in sessions if s.status == SessionStatus.COMPLETED)
    failed = sum(1 for s in sessions if s.status == SessionStatus.FAILED)
    killed = sum(1 for s in sessions if s.status == SessionStatus.KILLED)

    # Total cost and tokens
    total_cost = 0.0
    total_tokens = 0
    for s in sessions:
        cost_info = estimate_session_cost(s.model, s.token_usage)
        if cost_info["cost"] is not None:
            total_cost += cost_info["cost"]
        total_tokens += s.token_usage.get("total_tokens", 0)

    parts = []
    if running:
        parts.append(f"[yellow]{running} running[/]")
    if completed:
        parts.append(f"[green]{completed} done[/]")
    if failed:
        parts.append(f"[red]{failed} failed[/]")
    if killed:
        parts.append(f"[dim]{killed} killed[/]")
    parts.append(f"[cyan]{total_tokens:,} tokens[/]")
    parts.append(f"[green]{format_cost(total_cost)}[/]")
    if parts:
        console.print(" | ".join(parts))
        console.print()

    # Table
    table = Table(box=None, show_header=True, header_style="bold dim")
    table.add_column("ID", style="cyan", width=10)
    table.add_column("", width=2)
    table.add_column("T", width=2)
    table.add_column("Task", max_width=30)
    table.add_column("Updated", justify="right", width=8)
    table.add_column("Last Message", max_width=40)

    for s in sessions:
        icon = STATUS_ICONS.get(s.status.value, "?")
        task_preview = s.task[:27] + "..." if len(s.task) > 30 else s.task
        updated = time_ago(s.updated_at)

        # Get last assistant message
        messages = manager.get_messages(s.id)
        last_msg = ""
        for msg in reversed(messages):
            if msg.role == "assistant":
                last_msg = msg.content.replace("\n", " ")[:37]
                if len(msg.content) > 37:
                    last_msg += "..."
                break

        # Style based on status
        if s.status == SessionStatus.RUNNING:
            last_msg_styled = f"[yellow]{last_msg or '(working...)'}[/]"
            updated_styled = f"[yellow]{updated}[/]"
        elif s.status == SessionStatus.COMPLETED:
            try:
                dt = datetime.fromisoformat(s.updated_at)
                is_recent = (datetime.now() - dt).total_seconds() < 60
            except Exception:
                is_recent = False
            if is_recent:
                last_msg_styled = f"[green bold]{last_msg or '(done)'}[/]"
                updated_styled = f"[green bold]{updated} ★[/]"
            else:
                last_msg_styled = f"[green]{last_msg or '(done)'}[/]"
                updated_styled = f"[dim]{updated}[/]"
        elif s.status == SessionStatus.FAILED:
            err = s.error[:37] if s.error else "(failed)"
            last_msg_styled = f"[red]{err}...[/]"
            updated_styled = f"[red]{updated}[/]"
        else:
            last_msg_styled = f"[dim]{last_msg or '-'}[/]"
            updated_styled = f"[dim]{updated}[/]"

        table.add_row(s.short_id, icon, str(s.turn), task_preview, updated_styled, last_msg_styled)

    console.print(table)


def cmd_peek(manager, session_id: str):
    """Quick peek at session status."""
    session = manager.get_session(session_id)
    if not session:
        console.print(f"  [red]Session not found:[/] {session_id}")
        return

    icon = STATUS_ICONS.get(session.status.value, "?")
    console.print(f"\n{icon} [cyan]{session.short_id}[/] ({session.status.value})")
    console.print(f"  [dim]Task:[/] {session.task[:60]}...")
    console.print(f"  [dim]Turn:[/] {session.turn} | [dim]Updated:[/] {time_ago(session.updated_at)}")

    # Latest message
    messages = manager.get_messages(session.id)
    for msg in reversed(messages):
        if msg.role == "assistant":
            preview = msg.content.replace("\n", " ")[:100]
            if len(msg.content) > 100:
                preview += "..."
            console.print(f"\n  [bold]Latest:[/] {preview}")
            break
    console.print()


def cmd_show(manager, session_id: str):
    """Full session details with messages."""
    from zwarm.core.costs import estimate_session_cost

    session = manager.get_session(session_id)
    if not session:
        console.print(f"  [red]Session not found:[/] {session_id}")
        return

    # Header
    icon = STATUS_ICONS.get(session.status.value, "?")
    console.print(f"\n{icon} [bold cyan]{session.short_id}[/] - {session.status.value}")
    console.print(f"  [dim]Task:[/] {session.task}")
    console.print(f"  [dim]Turn:[/] {session.turn} | [dim]Runtime:[/] {session.runtime:.1f}s")

    # Token usage with cost estimate
    usage = session.token_usage
    input_tok = usage.get("input_tokens", 0)
    output_tok = usage.get("output_tokens", 0)
    total_tok = usage.get("total_tokens", input_tok + output_tok)

    cost_info = estimate_session_cost(session.model, usage)
    cost_str = f"[green]{cost_info['cost_formatted']}[/]" if cost_info["pricing_known"] else "[dim]?[/]"

    console.print(f"  [dim]Tokens:[/] {total_tok:,} ({input_tok:,} in / {output_tok:,} out) | [dim]Cost:[/] {cost_str}")

    if session.error:
        console.print(f"  [red]Error:[/] {session.error}")

    # Messages
    messages = manager.get_messages(session.id)
    if messages:
        console.print(f"\n[bold]Messages ({len(messages)}):[/]")
        for msg in messages:
            role = msg.role
            content = msg.content[:200]
            if len(msg.content) > 200:
                content += "..."

            if role == "user":
                console.print(f"  [blue]USER:[/] {content}")
            elif role == "assistant":
                console.print(f"  [green]ASSISTANT:[/] {content}")
            else:
                console.print(f"  [dim]{role.upper()}:[/] {content[:100]}")

    console.print()


def cmd_traj(manager, session_id: str, full: bool = False):
    """Show session trajectory."""
    session = manager.get_session(session_id)
    if not session:
        console.print(f"  [red]Session not found:[/] {session_id}")
        return

    trajectory = manager.get_trajectory(session_id, full=full)

    console.print(f"\n[bold]Trajectory for {session.short_id}[/] ({len(trajectory)} steps)")
    console.print(f"  [dim]Task:[/] {session.task[:60]}...")
    console.print()

    for i, step in enumerate(trajectory):
        step_type = step.get("type", "unknown")

        if step_type == "reasoning":
            text = step.get("full_text") if full else step.get("summary", "")
            console.print(f"  [dim]{i+1}.[/] [magenta]💭 thinking[/]")
            if text:
                console.print(f"     {text[:150]}{'...' if len(text) > 150 else ''}")

        elif step_type == "command":
            cmd = step.get("command", "")
            output = step.get("output", "")
            exit_code = step.get("exit_code", 0)
            console.print(f"  [dim]{i+1}.[/] [yellow]$ {cmd}[/]")
            if output and (full or len(output) < 100):
                console.print(f"     {output[:200]}")
            if exit_code and exit_code != 0:
                console.print(f"     [red](exit: {exit_code})[/]")

        elif step_type == "tool_call":
            tool = step.get("tool", "unknown")
            args_preview = step.get("args_preview", "")
            console.print(f"  [dim]{i+1}.[/] [cyan]🔧 {tool}[/]({args_preview})")

        elif step_type == "tool_output":
            output = step.get("output", "")
            preview = output[:100] if not full else output[:300]
            console.print(f"     [dim]→ {preview}[/]")

        elif step_type == "message":
            text = step.get("full_text") if full else step.get("summary", "")
            console.print(f"  [dim]{i+1}.[/] [green]💬 response[/]")
            if text:
                console.print(f"     {text[:150]}{'...' if len(text) > 150 else ''}")

    console.print()


def cmd_watch(manager, session_id: str):
    """
    Watch session output live.

    Polls trajectory and displays new steps as they appear.
    """
    from zwarm.sessions import SessionStatus

    session = manager.get_session(session_id)
    if not session:
        console.print(f"  [red]Session not found:[/] {session_id}")
        return

    console.print(f"\n[bold]Watching {session.short_id}[/]...")
    console.print(f"  [dim]Task:[/] {session.task[:60]}...")
    console.print(f"  [dim]Press Ctrl+C to stop watching[/]\n")

    seen_steps = 0
    last_status = None

    try:
        while True:
            # Refresh session
            session = manager.get_session(session_id)
            if not session:
                console.print("[red]Session disappeared![/]")
                break

            # Status change
            if session.status.value != last_status:
                icon = STATUS_ICONS.get(session.status.value, "?")
                console.print(f"\n{icon} Status: [bold]{session.status.value}[/]")
                last_status = session.status.value

            # Get trajectory
            trajectory = manager.get_trajectory(session_id, full=False)

            # Show new steps
            for i, step in enumerate(trajectory[seen_steps:], start=seen_steps + 1):
                step_type = step.get("type", "unknown")

                if step_type == "reasoning":
                    text = step.get("summary", "")[:80]
                    console.print(f"  [magenta]💭[/] {text}...")

                elif step_type == "command":
                    cmd = step.get("command", "")
                    console.print(f"  [yellow]$[/] {cmd}")

                elif step_type == "tool_call":
                    tool = step.get("tool", "unknown")
                    console.print(f"  [cyan]🔧[/] {tool}(...)")

                elif step_type == "message":
                    text = step.get("summary", "")[:80]
                    console.print(f"  [green]💬[/] {text}...")

            seen_steps = len(trajectory)

            # Check if done
            if session.status in (SessionStatus.COMPLETED, SessionStatus.FAILED, SessionStatus.KILLED):
                console.print(f"\n[dim]Session {session.status.value}. Final message:[/]")
                messages = manager.get_messages(session.id)
                for msg in reversed(messages):
                    if msg.role == "assistant":
                        console.print(f"  {msg.content[:200]}...")
                        break
                break

            time.sleep(1.0)

    except KeyboardInterrupt:
        console.print("\n[dim]Stopped watching.[/]")

    console.print()


def cmd_spawn(manager, task: str, working_dir: Path, model: str):
    """Spawn a new session."""
    console.print(f"\n[dim]Spawning session...[/]")
    console.print(f"  [dim]Dir:[/] {working_dir}")
    console.print(f"  [dim]Model:[/] {model}")

    try:
        session = manager.start_session(
            task=task,
            working_dir=working_dir,
            model=model,
            sandbox="workspace-write",
            source="user",
            adapter="codex",
        )

        console.print(f"\n[green]✓[/] Session: [cyan]{session.short_id}[/]")
        console.print(f"  [dim]Use 'watch {session.short_id}' to follow progress[/]")
        console.print(f"  [dim]Use 'show {session.short_id}' when complete[/]")

    except Exception as e:
        console.print(f"  [red]Error:[/] {e}")


def cmd_continue(manager, session_id: str, message: str):
    """Continue a conversation."""
    from zwarm.sessions import SessionStatus

    session = manager.get_session(session_id)
    if not session:
        console.print(f"  [red]Session not found:[/] {session_id}")
        return

    if session.status == SessionStatus.RUNNING:
        console.print(f"  [yellow]Session still running - wait for it to complete[/]")
        return

    if session.status == SessionStatus.KILLED:
        console.print(f"  [red]Session was killed - start a new one[/]")
        return

    console.print(f"\n[dim]Injecting message into {session.short_id}...[/]")

    updated = manager.inject_message(session_id, message)
    if updated:
        console.print(f"[green]✓[/] Message sent (turn {updated.turn})")
        console.print(f"  [dim]Use 'watch {session.short_id}' to follow response[/]")
    else:
        console.print(f"  [red]Failed to inject message[/]")


def cmd_kill(manager, target: str):
    """
    Kill session(s).

    Args:
        target: Session ID or "all" to kill all running
    """
    from zwarm.sessions import SessionStatus

    if target.lower() == "all":
        # Kill all running
        sessions = manager.list_sessions(status=SessionStatus.RUNNING)
        if not sessions:
            console.print("  [dim]No running sessions[/]")
            return

        killed = 0
        for s in sessions:
            if manager.kill_session(s.id):
                killed += 1
                console.print(f"  [green]✓[/] Killed {s.short_id}")

        console.print(f"\n[green]Killed {killed} session(s)[/]")
    else:
        # Kill single session
        session = manager.get_session(target)
        if not session:
            console.print(f"  [red]Session not found:[/] {target}")
            return

        if manager.kill_session(session.id):
            console.print(f"[green]✓[/] Killed {session.short_id}")
        else:
            console.print(f"  [yellow]Session not running or already stopped[/]")


def cmd_rm(manager, target: str):
    """
    Delete session(s).

    Args:
        target: Session ID or "all" to delete all non-running
    """
    from zwarm.sessions import SessionStatus

    if target.lower() == "all":
        # Delete all non-running (completed, failed, killed)
        sessions = manager.list_sessions()
        to_delete = [s for s in sessions if s.status != SessionStatus.RUNNING]

        if not to_delete:
            console.print("  [dim]Nothing to delete[/]")
            return

        deleted = 0
        for s in to_delete:
            if manager.delete_session(s.id):
                deleted += 1

        console.print(f"[green]✓[/] Deleted {deleted} session(s)")
    else:
        # Delete single session
        session = manager.get_session(target)
        if not session:
            console.print(f"  [red]Session not found:[/] {target}")
            return

        if manager.delete_session(session.id):
            console.print(f"[green]✓[/] Deleted {session.short_id}")
        else:
            console.print(f"  [red]Failed to delete[/]")




# =============================================================================
# Main REPL
# =============================================================================


def run_interactive(
    working_dir: Path,
    model: str = "gpt-5.1-codex-mini",
):
    """
    Run the interactive REPL.

    Args:
        working_dir: Default working directory for sessions
        model: Default model for sessions
    """
    from zwarm.sessions import CodexSessionManager

    manager = CodexSessionManager(working_dir / ".zwarm")

    # Setup prompt with autocomplete
    def get_sessions():
        return manager.list_sessions()

    completer = SessionCompleter(get_sessions)
    style = Style.from_dict({
        "prompt": "cyan bold",
    })

    session = PromptSession(
        completer=completer,
        history=InMemoryHistory(),
        style=style,
        complete_while_typing=True,
    )

    # Welcome
    console.print("\n[bold cyan]zwarm interactive[/] - Session Manager\n")
    console.print(f"  [dim]Dir:[/] {working_dir.absolute()}")
    console.print(f"  [dim]Model:[/] {model}")
    console.print(f"\n  Type [cyan]help[/] for commands, [cyan]quit[/] to exit.")
    console.print(f"  [dim]Tab to autocomplete session IDs[/]\n")

    # REPL
    while True:
        try:
            raw = session.prompt("> ").strip()
            if not raw:
                continue

            try:
                parts = shlex.split(raw)
            except ValueError:
                parts = raw.split()

            cmd = parts[0].lower()
            args = parts[1:]

            # Dispatch
            if cmd in ("q", "quit", "exit"):
                console.print("\n[dim]Goodbye![/]\n")
                break

            elif cmd in ("h", "help"):
                cmd_help()

            elif cmd in ("ls", "list"):
                cmd_ls(manager)

            elif cmd in ("?", "peek"):
                if not args:
                    console.print("  [red]Usage:[/] peek ID")
                else:
                    cmd_peek(manager, args[0])

            elif cmd == "show":
                if not args:
                    console.print("  [red]Usage:[/] show ID")
                else:
                    cmd_show(manager, args[0])

            elif cmd in ("traj", "trajectory"):
                if not args:
                    console.print("  [red]Usage:[/] traj ID [--full]")
                else:
                    full = "--full" in args
                    sid = [a for a in args if not a.startswith("-")][0]
                    cmd_traj(manager, sid, full=full)

            elif cmd == "watch":
                if not args:
                    console.print("  [red]Usage:[/] watch ID")
                else:
                    cmd_watch(manager, args[0])

            elif cmd == "spawn":
                if not args:
                    console.print("  [red]Usage:[/] spawn \"task\" [--dir PATH]")
                else:
                    # Parse spawn args
                    task_parts = []
                    spawn_dir = working_dir
                    spawn_model = model
                    i = 0
                    while i < len(args):
                        if args[i] in ("--dir", "-d") and i + 1 < len(args):
                            spawn_dir = Path(args[i + 1])
                            i += 2
                        elif args[i] in ("--model", "-m") and i + 1 < len(args):
                            spawn_model = args[i + 1]
                            i += 2
                        else:
                            task_parts.append(args[i])
                            i += 1

                    task = " ".join(task_parts)
                    if task:
                        cmd_spawn(manager, task, spawn_dir, spawn_model)
                    else:
                        console.print("  [red]Task required[/]")

            elif cmd in ("c", "continue"):
                if len(args) < 2:
                    console.print("  [red]Usage:[/] c ID \"message\"")
                else:
                    cmd_continue(manager, args[0], " ".join(args[1:]))

            elif cmd == "kill":
                if not args:
                    console.print("  [red]Usage:[/] kill ID | all")
                else:
                    cmd_kill(manager, args[0])

            elif cmd in ("rm", "delete"):
                if not args:
                    console.print("  [red]Usage:[/] rm ID | all")
                else:
                    cmd_rm(manager, args[0])

            else:
                console.print(f"  [yellow]Unknown command:[/] {cmd}")
                console.print("  [dim]Type 'help' for commands[/]")

        except KeyboardInterrupt:
            console.print("\n[dim](Ctrl+C again or 'quit' to exit)[/]")
        except EOFError:
            console.print("\n[dim]Goodbye![/]\n")
            break
        except Exception as e:
            console.print(f"  [red]Error:[/] {e}")
