AIPA/orchestration/ui.py

"""
ui.py — AIPA Terminal Interface
Authored by: Iris, Director of Interface & Experience

Visual language:
  Miranda  — cyan       primary orchestrator, delivers to Principal
  Vera     — yellow     auditor, independent track
  Atlas    — green      research
  Cole     — blue       operations
  Clio     — magenta    analysis
  Evelyn   — white      personnel & systems
  Iris     — bright_cyan interface
  System   — dim white  infrastructure messages
  Error    — red        failures and blockers
  Warning  — orange     flags and cautions

Panels carry the agent name as their title. The Principal never has to read
a header to know who is speaking — color and panel style make it immediate.
"""

from contextlib import contextmanager
from datetime import datetime

import config as _config

from rich.console import Console
from rich.columns import Columns
from rich.markdown import Markdown
from rich.panel import Panel
from rich.progress import (
    BarColumn,
    Progress,
    SpinnerColumn,
    TaskProgressColumn,
    TextColumn,
    TimeElapsedColumn,
)
from rich.rule import Rule
from rich.table import Table
from rich.text import Text
from rich.theme import Theme
from rich import box

# ---------------------------------------------------------------------------
# Console setup
# ---------------------------------------------------------------------------

def _build_theme() -> Theme:
    """Build the Rich theme dynamically from agents.yaml color definitions."""
    theme_dict: dict[str, str] = {
        # System styles — not agent-specific
        "system":  "dim white",
        "error":   "bold red",
        "warning": "bold yellow",
        "task_id": "dim cyan",
        "label":   "bold white",
        "muted":   "dim",
        "pass":    "bold green",
        "flag":    "bold yellow",
        "reject":  "bold red",
    }
    # Agent name aliases — one entry per agent in agents.yaml
    for name in _config.AGENT_CONFIGS:
        theme_dict[name] = f"bold {_config.agent_color(name)}"
    # Role aliases — stable names regardless of which agent holds the role
    for role, agent in (
        ("orchestrator", _config.ORCHESTRATOR_AGENT),
        ("auditor",      _config.AUDITOR_AGENT),
        ("recruiter",    _config.RECRUITER_AGENT),
    ):
        if agent:
            theme_dict[role] = f"bold {_config.agent_color(agent)}"
    return Theme(theme_dict)


console = Console(theme=_build_theme(), highlight=False)

# ---------------------------------------------------------------------------
# Agent visual config
# ---------------------------------------------------------------------------

def _agent_color(name: str) -> str:
    return _config.agent_color(name.lower())


def _agent_display_title(name: str) -> str:
    """Return 'Name — Title' using the title from agents.yaml."""
    title = _config.agent_title(name)
    return f"{name.capitalize()} — {title}" if title else name.capitalize()

# ---------------------------------------------------------------------------
# Session chrome
# ---------------------------------------------------------------------------

def print_session_header():
    """Opening banner for a new Principal session."""
    cos_name   = _config.ORCHESTRATOR_AGENT.capitalize()
    cos_title  = _config.agent_title(_config.ORCHESTRATOR_AGENT)
    aud_name   = _config.AUDITOR_AGENT.capitalize()
    app_name   = _config.APP_NAME

    console.print()
    console.rule(f"[bold cyan]{app_name}[/bold cyan]  [dim]Autonomous Intelligence Principal Architecture[/dim]", style="cyan")
    console.print()

    table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2))
    table.add_column(style="dim")
    table.add_column(style="white")
    table.add_row("Session opened", datetime.now().strftime("%Y-%m-%d  %H:%M:%S"))
    table.add_row("Orchestrator",   f"{cos_name} — {cos_title}")
    table.add_row("Auditor",        f"{aud_name} (independent track)")
    table.add_row("Commands",       "/help for available commands")
    console.print(table)
    console.print()


def print_session_footer(task_count: int):
    """Closing rule when the session ends."""
    console.print()
    console.rule(
        f"[dim]Session closed — {task_count} directive{'s' if task_count != 1 else ''} processed[/dim]",
        style="dim"
    )
    console.print()

# ---------------------------------------------------------------------------
# Input prompt
# ---------------------------------------------------------------------------

def get_directive() -> str:
    """Styled input prompt for the Principal."""
    console.print()
    return console.input("[bold white]Principal ›[/bold white] ").strip()

# ---------------------------------------------------------------------------
# System / status messages
# ---------------------------------------------------------------------------

def print_system(message: str):
    console.print(f"  [system]⬡  {message}[/system]")


def print_error(message: str):
    console.print(f"\n  [error]✗  {message}[/error]\n")


def print_warning(message: str):
    console.print(f"  [warning]⚠  {message}[/warning]")


def print_task_received(task_id: str, directive_preview: str):
    console.print()
    console.print(
        f"  [task_id]{task_id}[/task_id]  [dim]Directive received:[/dim] "
        f"[white]{directive_preview[:100]}{'…' if len(directive_preview) > 100 else ''}[/white]"
    )

# ---------------------------------------------------------------------------
# Lead dispatch progress
# ---------------------------------------------------------------------------

def _brief_preview(brief: str, max_len: int = 72) -> str:
    """Return a short preview of a task brief, skipping the header lines."""
    lines = [l.strip() for l in brief.splitlines() if l.strip()]
    content = [
        l for l in lines
        if not l.upper().startswith("TASK BRIEF") and not l.upper().startswith("TO:")
    ]
    text = " ".join(content)
    return (text[:max_len] + "…") if len(text) > max_len else text


def print_task_dispatch_plan(tasks: list):
    """
    Show the task assignments before leads are dispatched.
    Each row: task_id | lead | brief preview.
    """
    console.print()
    console.rule("[dim]Lead dispatch[/dim]", style="dim")
    table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2), expand=False)
    table.add_column(style="task_id", no_wrap=True)
    table.add_column(no_wrap=True)
    table.add_column(style="dim")
    for t in tasks:
        color = _agent_color(t.assigned_to)
        table.add_row(
            t.task_id,
            f"[{color}]{_agent_display_title(t.assigned_to)}[/{color}]",
            _brief_preview(t.brief),
        )
    console.print(table)


@contextmanager
def lead_dispatch_progress(tasks: list):
    """
    Per-task live progress display while leads are working.
    Yields (progress, rows) where rows is keyed by task_id.
    """
    def _row(task_id: str, lead_name: str, status: str) -> str:
        color = _agent_color(lead_name)
        return (
            f"[task_id]{task_id}[/task_id]  "
            f"[{color}]{lead_name.capitalize()}[/{color}]  "
            f"{status}"
        )

    with Progress(
        SpinnerColumn(style="dim"),
        TextColumn("{task.description}"),
        TimeElapsedColumn(),
        console=console,
        transient=True,
    ) as progress:
        rows = {
            t.task_id: progress.add_task(
                _row(t.task_id, t.assigned_to, "[dim]queued[/dim]"),
                total=1,
            )
            for t in tasks
        }
        yield progress, rows


def mark_task_running(progress, rows: dict, task_id: str, lead_name: str):
    """Update a task row to show it is actively being worked."""
    if task_id not in rows:
        return
    color = _agent_color(lead_name)
    progress.update(
        rows[task_id],
        description=(
            f"[task_id]{task_id}[/task_id]  "
            f"[{color}]{lead_name.capitalize()}[/{color}]  "
            f"[yellow]working…[/yellow]"
        ),
    )


def mark_task_done(progress, rows: dict, task_id: str, lead_name: str, error: bool = False):
    """Update a task row to show it is finished (complete or error)."""
    if task_id not in rows:
        return
    color = _agent_color(lead_name)
    status = "[red]error[/red]" if error else "[green]done[/green]"
    progress.update(
        rows[task_id],
        completed=1,
        description=(
            f"[task_id]{task_id}[/task_id]  "
            f"[{color}]{lead_name.capitalize()}[/{color}]  "
            f"{status}"
        ),
    )


def print_direct_response_notice():
    """Shown when Miranda handles a directive without dispatching leads."""
    console.print(
        "  [dim]Miranda responded directly — no lead dispatch required.[/dim]"
    )

# ---------------------------------------------------------------------------
# Agent output panels
# ---------------------------------------------------------------------------

def print_agent_panel(agent_name: str, content: str, task_id: str = ""):
    """
    Render an agent's output in a styled panel.
    Content is treated as Markdown.
    """
    color      = _agent_color(agent_name)
    title_text = _agent_display_title(agent_name)
    subtitle   = f"[dim]{task_id}[/dim]" if task_id else ""

    panel = Panel(
        Markdown(content),
        title=f"[{color}] {title_text} [/{color}]",
        subtitle=subtitle,
        border_style=color,
        padding=(1, 2),
        box=box.ROUNDED,
    )
    console.print()
    console.print(panel)


def print_deliverable(content: str, task_id: str = ""):
    """Orchestrator's synthesized deliverable."""
    print_agent_panel(_config.ORCHESTRATOR_AGENT, content, task_id)


def print_audit_memo(content: str, task_id: str = ""):
    """
    Auditor's memo. Verdict line drives the border colour.
    PASS = green, PASS WITH NOTES = yellow, FLAG = yellow, REJECT = red.
    """
    aud_name  = _config.AUDITOR_AGENT
    aud_title = _config.agent_title(aud_name)
    aud_color = _agent_color(aud_name)

    verdict_style = "dim"
    for line in content.splitlines():
        upper = line.upper()
        if "VERDICT:" in upper:
            if "REJECT" in upper:
                verdict_style = "reject"
            elif "FLAG" in upper:
                verdict_style = "flag"
            elif "PASS WITH NOTES" in upper:
                verdict_style = "warning"
            elif "PASS" in upper:
                verdict_style = "pass"
            break

    panel = Panel(
        Markdown(content),
        title=f"[{aud_color}] {aud_name.capitalize()} — {aud_title} [/{aud_color}]",
        subtitle=f"[dim]{task_id}[/dim]" if task_id else "",
        border_style=verdict_style if verdict_style != "dim" else aud_color,
        padding=(1, 2),
        box=box.ROUNDED,
    )
    console.print()
    console.print(panel)


def print_lead_output(agent_name: str, content: str, task_id: str = ""):
    """A lead's raw output (shown only in DEBUG mode)."""
    print_agent_panel(agent_name, content, task_id)

# ---------------------------------------------------------------------------
# Standing Brief renderer
# ---------------------------------------------------------------------------

def print_standing_brief(raw_md: str):
    """
    Render the standing brief as a navigable terminal layout.
    Sections are separated by rules; tables are preserved where present.
    """
    console.print()
    console.rule("[bold cyan]Standing Brief[/bold cyan]", style="cyan")
    console.print()
    console.print(Markdown(raw_md))
    console.print()
    console.rule(style="dim")
    console.print()

# ---------------------------------------------------------------------------
# Synthesis progress (Miranda)
# ---------------------------------------------------------------------------

@contextmanager
def synthesis_progress():
    """Spinner shown while Miranda synthesizes lead outputs."""
    with Progress(
        SpinnerColumn(style="cyan"),
        TextColumn("[cyan]Miranda synthesizing…[/cyan]"),
        TimeElapsedColumn(),
        console=console,
        transient=True,
    ) as progress:
        progress.add_task("", total=None)
        yield


@contextmanager
def audit_progress():
    """Spinner shown while Vera audits."""
    with Progress(
        SpinnerColumn(style="yellow"),
        TextColumn("[yellow]Vera auditing…[/yellow]"),
        TimeElapsedColumn(),
        console=console,
        transient=True,
    ) as progress:
        progress.add_task("", total=None)
        yield


@contextmanager
def brief_update_progress():
    """Spinner shown while Miranda updates the standing brief."""
    with Progress(
        SpinnerColumn(style="dim"),
        TextColumn("[dim]Updating standing brief…[/dim]"),
        console=console,
        transient=True,
    ) as progress:
        progress.add_task("", total=None)
        yield

# ---------------------------------------------------------------------------
# Dry-run callout
# ---------------------------------------------------------------------------

def print_dry_run_banner():
    panel = Panel(
        "[bold yellow]DRY RUN MODE[/bold yellow]\n"
        "[dim]Routing and pipeline logic will execute normally.\n"
        "No provider API calls will be made.[/dim]",
        border_style="yellow",
        box=box.ROUNDED,
        padding=(0, 2),
    )
    console.print()
    console.print(panel)
    console.print()


# ---------------------------------------------------------------------------
# Slash command displays
# ---------------------------------------------------------------------------

def print_help(commands: list[tuple[str, str]]):
    """Render the slash command reference table."""
    console.print()
    console.rule("[bold cyan]Commands[/bold cyan]", style="cyan")
    console.print()

    table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2), expand=False)
    table.add_column(style="bold cyan", no_wrap=True)
    table.add_column(style="white")

    for cmd, description in commands:
        table.add_row(cmd, description)

    console.print(table)
    console.print()


def print_status(
    session_start: datetime,
    directive_count: int,
    session_log: list[str],
    auto_audit: bool,
    debug: bool,
    dry_run: bool,
):
    """Render current session status."""
    console.print()
    console.rule("[bold cyan]Session Status[/bold cyan]", style="cyan")
    console.print()

    elapsed = datetime.now() - session_start
    hours, remainder = divmod(int(elapsed.total_seconds()), 3600)
    minutes, seconds = divmod(remainder, 60)
    elapsed_str = f"{hours}h {minutes}m {seconds}s" if hours else f"{minutes}m {seconds}s"

    global_provider = _config.ACTIVE_PROVIDER
    provider_display = (
        global_provider if global_provider != "none"
        else "[dim]none — set default_provider in config/agents.yaml[/dim]"
    )

    table = Table(box=box.SIMPLE, show_header=False, padding=(0, 2))
    table.add_column(style="dim", no_wrap=True)
    table.add_column(style="white")
    table.add_row("Session started",    session_start.strftime("%Y-%m-%d  %H:%M:%S"))
    table.add_row("Elapsed",            elapsed_str)
    table.add_row("Directives",         str(directive_count))
    table.add_row("Global provider",    provider_display)
    table.add_row("Auto-audit",         "[green]on[/green]" if auto_audit else "[dim]off[/dim]")
    table.add_row("Debug output",       "[yellow]on[/yellow]" if debug else "[dim]off[/dim]")
    table.add_row("Dry run",            "[yellow]on[/yellow]" if dry_run else "[dim]off[/dim]")
    console.print(table)

    if session_log:
        console.print()
        console.rule("[dim]Session log[/dim]", style="dim")
        console.print()
        for entry in session_log:
            console.print(f"  [dim]·[/dim] [muted]{entry}[/muted]")

    console.print()


def print_agents(agent_configs: dict, active_provider: str):
    """Render the named agent roster with per-agent provider resolution."""
    console.print()
    console.rule("[bold cyan]Agent Roster[/bold cyan]", style="cyan")
    console.print()

    table = Table(box=box.SIMPLE, show_header=True, padding=(0, 2), expand=False)
    table.add_column("Agent",    style="bold white",  no_wrap=True)
    table.add_column("Role",     style="dim",         no_wrap=True)
    table.add_column("Provider", style="white",       no_wrap=True)
    table.add_column("Model",    style="cyan")
    table.add_column("Temp",     style="dim",         justify="right")
    table.add_column("Tokens",   style="dim",         justify="right")

    for name, cfg in agent_configs.items():
        color    = _agent_color(name)
        label    = _config.agent_title(name)
        provider = _config.agent_provider_name(name)
        model    = _config.agent_model(name) or "[dim]unset[/dim]"

        provider_display = (
            f"[dim]{provider}[/dim]" if provider == "none"
            else provider
        )

        table.add_row(
            f"[{color}]{name.capitalize()}[/{color}]",
            label,
            provider_display,
            model,
            str(cfg["temperature"]),
            str(cfg["max_tokens"]),
        )

    console.print(table)
    if active_provider == "none":
        console.print("  [dim]Global fallback: none — set default_provider in config/agents.yaml[/dim]")
    else:
        console.print(f"  [dim]Global fallback provider:[/dim] [white]{active_provider}[/white]")
    console.print()


def print_history(session_log: list[str]):
    """Render the session log for this session."""
    console.print()
    console.rule("[bold cyan]Session History[/bold cyan]", style="cyan")
    console.print()

    if not session_log:
        console.print("  [dim]No directives processed yet.[/dim]")
    else:
        for i, entry in enumerate(session_log, 1):
            console.print(f"  [dim]{i:>2}.[/dim]  [muted]{entry}[/muted]")

    console.print()


def print_tasks(tasks: list, session_id: str = ""):
    """Render the task roster for the current session."""
    console.print()
    title = f"[bold cyan]Tasks[/bold cyan]"
    if session_id:
        title += f"  [dim]{session_id}[/dim]"
    console.rule(title, style="cyan")
    console.print()

    if not tasks:
        console.print("  [dim]No tasks dispatched this session.[/dim]")
        console.print()
        return

    STATUS_STYLE = {
        "pending":     "[dim]pending[/dim]",
        "in_progress": "[yellow]in progress[/yellow]",
        "complete":    "[green]complete[/green]",
        "error":       "[red]error[/red]",
    }

    table = Table(box=box.SIMPLE, show_header=True, padding=(0, 2), expand=False)
    table.add_column("Task",   style="task_id", no_wrap=True)
    table.add_column("Lead",   no_wrap=True)
    table.add_column("Status", no_wrap=True)
    table.add_column("Brief",  style="dim")

    for t in tasks:
        color  = _agent_color(t.assigned_to)
        status = STATUS_STYLE.get(t.status, t.status)
        table.add_row(
            t.task_id,
            f"[{color}]{t.assigned_to.capitalize()}[/{color}]",
            status,
            _brief_preview(t.brief, max_len=60),
        )

    console.print(table)
    console.print()


def print_debug_toggle(enabled: bool):
    state = "[yellow]ON[/yellow]" if enabled else "[dim]OFF[/dim]"
    console.print(f"  [dim]Debug output:[/dim] {state}")


def print_audit_toggle(enabled: bool):
    state = "[green]ON[/green]" if enabled else "[dim]OFF[/dim]"
    console.print(f"  [dim]Auto-audit (Vera):[/dim] {state}")