"""
orchestrator.py — AIPA Session Orchestrator
Entry point for Principal → Miranda → Leads → Principal sessions.

Usage:
    python orchestrator.py                  # interactive session
    python orchestrator.py --dry-run        # print routing without calling agents
    python orchestrator.py --agent evelyn   # open a direct session with a named agent

Requirements:
    See requirements.txt. Install with: pip install -r requirements.txt
"""

import asyncio
import argparse
import re
from dataclasses import dataclass, field
from datetime import datetime

_THINK_TAG_RE     = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
_TOOL_CALL_TAG_RE = re.compile(r"<tool_call>\s*(.*?)\s*</tool_call>", re.DOTALL)

import config
import task_store
import tools as _tools
import ui

# ---------------------------------------------------------------------------
# Provider Clients
# Each agent builds its own client at session startup via build_client().
# Adding a new provider: implement the class, add a case to build_client().
# ---------------------------------------------------------------------------

class ProviderClient:
    """
    Base / placeholder client.
    Used when AIPA_PROVIDER=none or in dry-run mode.
    Raises a clear error if actually called.
    """
    def call(self, system: str, messages: list[dict], model: str,
             temperature: float, max_tokens: int) -> str:
        raise NotImplementedError(
            "No provider is configured for this agent. "
            "Set 'provider' on the agent or 'default_provider' in config/agents.yaml."
        )

    async def call_async(self, system: str, messages: list[dict], model: str,
                         temperature: float, max_tokens: int) -> str:
        return self.call(system, messages, model, temperature, max_tokens)

    async def call_async_streaming(self, system: str, messages: list[dict], model: str,
                                   temperature: float, max_tokens: int):
        """Async generator yielding ('thinking'|'content', text) chunks.
        Default: emits the full response as a single 'content' chunk."""
        response = await self.call_async(system, messages, model, temperature, max_tokens)
        yield ('content', response)

    async def call_async_with_tools(self, system: str, messages: list[dict], model: str,
                                    temperature: float, max_tokens: int,
                                    tools: list[dict]) -> "_ToolMessage":
        """Non-streaming call with tool schemas. Returns a _ToolMessage.
        Base implementation ignores tools and wraps the text response."""
        text = await self.call_async(system, messages, model, temperature, max_tokens)
        return _ToolMessage(content=text, tool_calls=None)


# --- OpenAI and any OpenAI-compatible endpoint ---
# Covers AIPA_PROVIDER / {AGENT}_PROVIDER = openai or openai_compatible.
# Compatible servers: LM Studio, llama.cpp, Ollama (/v1 endpoint), Groq,
#                     Together AI, Mistral API, Anyscale, Fireworks, etc.
class OpenAIClient(ProviderClient):
    def __init__(self, api_key: str, base_url: str | None, extra_body: dict | None = None):
        from openai import AsyncOpenAI
        self.client     = AsyncOpenAI(api_key=api_key, base_url=base_url)
        self.extra_body = extra_body or {}

    async def call_async(self, system, messages, model, temperature, max_tokens):
        full_messages = [{"role": "system", "content": system}] + messages
        response = await self.client.chat.completions.create(
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            messages=full_messages,
            extra_body=self.extra_body or None,
        )
        content = response.choices[0].message.content or ""
        return _THINK_TAG_RE.sub("", content).strip()

    def call(self, system, messages, model, temperature, max_tokens):
        return asyncio.run(
            self.call_async(system, messages, model, temperature, max_tokens)
        )

    async def call_async_with_tools(self, system, messages, model, temperature, max_tokens, tools):
        """Non-streaming call with tool schemas. Returns a _ToolMessage."""
        full_messages = [{"role": "system", "content": system}] + messages
        response = await self.client.chat.completions.create(
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            messages=full_messages,
            tools=tools,
            tool_choice="auto",
            extra_body=self.extra_body or None,
        )
        msg = response.choices[0].message
        content = _THINK_TAG_RE.sub("", msg.content or "").strip()
        tool_calls = msg.tool_calls or None

        # Fallback: if the model embedded tool calls as <tool_call> XML in
        # the content (common with Qwen3 when the server lacks --tool-call-parser),
        # parse them out and strip the tags from the displayed content.
        if tool_calls is None and content:
            content, tool_calls = _extract_content_tool_calls(content)

        return _ToolMessage(content=content, tool_calls=tool_calls)

    async def call_async_streaming(self, system, messages, model, temperature, max_tokens):
        """Stream via the OpenAI-compatible API, yielding thinking and content chunks.
        Thinking tokens arrive in delta.reasoning_content (Qwen3, DeepSeek R1 style);
        regular tokens arrive in delta.content."""
        full_messages = [{"role": "system", "content": system}] + messages
        stream = await self.client.chat.completions.create(
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            messages=full_messages,
            extra_body=self.extra_body or None,
            stream=True,
        )
        async for chunk in stream:
            if not chunk.choices:
                continue
            delta = chunk.choices[0].delta
            thinking = getattr(delta, 'reasoning_content', None)
            if thinking:
                yield ('thinking', thinking)
            if delta.content:
                yield ('content', delta.content)


# --- Anthropic ---
# Uncomment when ready to use. Install: pip install anthropic
# class AnthropicClient(ProviderClient):
#     def __init__(self, api_key: str):
#         import anthropic
#         self.client = anthropic.Anthropic(api_key=api_key)
#
#     def call(self, system, messages, model, temperature, max_tokens):
#         response = self.client.messages.create(
#             model=model,
#             max_tokens=max_tokens,
#             system=system,
#             messages=messages,
#             temperature=temperature,
#         )
#         return response.content[0].text
#
#     async def call_async(self, system, messages, model, temperature, max_tokens):
#         loop = asyncio.get_event_loop()
#         return await loop.run_in_executor(
#             None, lambda: self.call(system, messages, model, temperature, max_tokens)
#         )


# --- Ollama native API ---
# Uncomment when ready to use. No extra SDK required.
# class OllamaClient(ProviderClient):
#     def __init__(self, base_url: str):
#         import requests as req
#         self._requests = req
#         self.base_url = base_url
#
#     def call(self, system, messages, model, temperature, max_tokens):
#         full_messages = [{"role": "system", "content": system}] + messages
#         response = self._requests.post(
#             f"{self.base_url}/api/chat",
#             json={
#                 "model": model,
#                 "stream": False,
#                 "options": {"temperature": temperature, "num_predict": max_tokens},
#                 "messages": full_messages,
#             },
#             timeout=config.LEAD_TIMEOUT_SECONDS,
#         )
#         response.raise_for_status()
#         return response.json()["message"]["content"]
#
#     async def call_async(self, system, messages, model, temperature, max_tokens):
#         loop = asyncio.get_event_loop()
#         return await loop.run_in_executor(
#             None, lambda: self.call(system, messages, model, temperature, max_tokens)
#         )


def build_client(agent_name: str) -> ProviderClient:
    """
    Build the correct ProviderClient for a named agent.
    Provider type is resolved from agents.yaml via config.
    API keys and base URLs are resolved from .env via config.
    """
    provider_type = config.agent_provider_type(agent_name)

    if provider_type == "none":
        return ProviderClient()
    elif provider_type in ("openai", "openai_compatible"):
        return OpenAIClient(
            api_key=config.agent_api_key(agent_name),
            base_url=config.agent_base_url(agent_name),
            extra_body=config.agent_extra_body(agent_name),
        )
    # elif provider_type == "anthropic":
    #     return AnthropicClient(api_key=config.agent_api_key(agent_name))
    # elif provider_type == "ollama":
    #     return OllamaClient(base_url=config.agent_base_url(agent_name))
    else:
        raise ValueError(
            f"Agent '{agent_name}' has unrecognised provider type: {provider_type!r}. "
            f"Check the provider definition in config/agents.yaml."
        )


# ---------------------------------------------------------------------------
# Data Classes
# ---------------------------------------------------------------------------

@dataclass
class _ToolMessage:
    """Minimal message returned by call_async_with_tools."""
    content: str
    tool_calls: list | None   # list of openai ToolCall objects, or None


# ---------------------------------------------------------------------------
# Content-based tool call fallback
# Qwen3 (and some other models) emit tool calls as <tool_call>…</tool_call>
# XML blocks in the content when the server isn't configured to return
# structured tool_calls (e.g. vLLM without --tool-call-parser).
# We parse these and synthesise objects that _run_tool_loop can consume.
# ---------------------------------------------------------------------------

import json as _json

@dataclass
class _SyntheticFunction:
    name: str
    arguments: str   # JSON string, same shape as openai ToolCall.function.arguments


@dataclass
class _SyntheticToolCall:
    id: str
    function: _SyntheticFunction


def _extract_content_tool_calls(content: str) -> tuple[str, list[_SyntheticToolCall] | None]:
    """
    Scan *content* for <tool_call>…</tool_call> blocks.
    Returns (cleaned_content, tool_calls_list) where cleaned_content has the
    blocks removed, and tool_calls_list is None if nothing was found.
    """
    matches = _TOOL_CALL_TAG_RE.findall(content)
    if not matches:
        return content, None

    calls = []
    for i, raw in enumerate(matches):
        try:
            parsed = _json.loads(raw)
        except _json.JSONDecodeError:
            continue
        name = parsed.get("name") or parsed.get("function") or ""
        args = parsed.get("arguments") or parsed.get("parameters") or {}
        if isinstance(args, dict):
            args = _json.dumps(args)
        calls.append(_SyntheticToolCall(
            id=f"call_{i}",
            function=_SyntheticFunction(name=name, arguments=args),
        ))

    if not calls:
        return content, None

    cleaned = _TOOL_CALL_TAG_RE.sub("", content).strip()
    return cleaned, calls


@dataclass
class AgentState:
    """Runtime state for one named agent during a session."""
    name: str
    system_prompt: str
    model: str
    temperature: float
    max_tokens: int
    stateful: bool
    client: ProviderClient
    history: list[dict] = field(default_factory=list)
    tools: list[str] = field(default_factory=list)   # tool names from agents.yaml

    def add_user(self, content: str):
        self.history.append({"role": "user", "content": content})

    def add_assistant(self, content: str):
        self.history.append({"role": "assistant", "content": content})

    def messages_for_call(self) -> list[dict]:
        return self.history if self.stateful else self.history[-2:]


@dataclass
class Task:
    task_id: str
    directive: str
    assigned_to: str
    brief: str
    status: str = "pending"
    output: str = ""
    error: str = ""


# ---------------------------------------------------------------------------
# Prompt Loader
# ---------------------------------------------------------------------------

def load_system_prompt(prompt_file: str) -> str:
    """
    Read a named agent's prompt file and extract the system prompt section.
    Strips everything from '## Access Configuration' onward.
    """
    path = config.PROMPTS_DIR / prompt_file
    if not path.exists():
        raise FileNotFoundError(f"Prompt file not found: {path}")

    text = path.read_text(encoding="utf-8")

    match = re.search(
        r"## System Prompt\n(.+?)(?=\n## Access Configuration|\Z)",
        text,
        re.DOTALL,
    )
    if match:
        return match.group(1).strip()

    # Fallback: strip header metadata block
    lines = text.splitlines()
    in_header = True
    body_lines = []
    for line in lines:
        if in_header and (line.startswith("**") or line.startswith("#")):
            in_header = False
        if not in_header:
            body_lines.append(line)
    return "\n".join(body_lines).strip()


def load_standing_brief() -> str:
    path = config.STANDING_BRIEF_PATH
    if path.exists() and path.stat().st_size > 64:
        return path.read_text(encoding="utf-8")
    # Fall back to template for a clean starting point
    if config.BRIEF_TEMPLATE_PATH.exists():
        return config.BRIEF_TEMPLATE_PATH.read_text(encoding="utf-8")
    return "(No standing brief found — create docs/standing_brief.md to initialise.)"


def save_standing_brief(content: str):
    config.STANDING_BRIEF_PATH.write_text(content, encoding="utf-8")
    ui.print_system("Standing brief updated.")


# ---------------------------------------------------------------------------
# Agent Builder
# Resolves provider, model, key, and URL for each agent independently.
# ---------------------------------------------------------------------------

def build_agent(name: str) -> AgentState:
    """
    Construct an AgentState for a named agent.
    Each agent gets its own ProviderClient based on its own provider config.
    """
    cfg = config.AGENT_CONFIGS[name]
    return AgentState(
        name=name,
        system_prompt=load_system_prompt(cfg["prompt_file"]),
        model=config.agent_model(name),
        temperature=cfg["temperature"],
        max_tokens=cfg["max_tokens"],
        stateful=cfg["stateful"],
        client=build_client(name),
        tools=cfg.get("tools", []),
    )


# ---------------------------------------------------------------------------
# Task ID Generator
# ---------------------------------------------------------------------------

_task_counter = 0

def new_task_id() -> str:
    global _task_counter
    _task_counter += 1
    return f"T-{datetime.now().strftime('%Y%m%d')}-{_task_counter:03d}"


# ---------------------------------------------------------------------------
# Tool-Calling Loop
# ---------------------------------------------------------------------------

async def _run_tool_loop(
    agent: AgentState,
    stream_mode: str = "deliverable",
    dry_run: bool = False,
) -> str:
    """
    Agentic loop for tool-enabled agents.

    Opens a transient height-capped Live panel (like the thinking panel) that
    accumulates tool call summaries as each tool completes.  When the model
    stops issuing tool_calls the panel closes, a persistent one-liner summary
    is printed, and the final content is handed off to stream_agent_output.

    The agent's history is updated with assistant tool-call messages and
    tool result messages so that multi-turn tool use works correctly.

    Returns the final content string.
    """
    if dry_run:
        return f"[DRY RUN] Would run tool loop for {agent.name}"

    schemas = _tools.get_schemas(agent.tools)
    show_panel = stream_mode != "background"

    if show_panel:
        ui.start_tool_panel(agent.name)

    try:
        # Iterate until the model stops calling tools
        while True:
            msg: _ToolMessage = await agent.client.call_async_with_tools(
                system=agent.system_prompt,
                messages=agent.messages_for_call(),
                model=agent.model,
                temperature=agent.temperature,
                max_tokens=agent.max_tokens,
                tools=schemas,
            )

            if not msg.tool_calls:
                # No more tool calls — close panel, then stream final content
                if show_panel:
                    ui.end_tool_panel(agent.name)

                final_content = msg.content

                async def _emit_final():
                    yield ('content', final_content)

                response = await ui.stream_agent_output(agent.name, _emit_final(), mode=stream_mode)
                agent.add_assistant(response)
                return response

            # Append the assistant message with tool_calls to history
            tool_calls_repr = [
                {
                    "id":       tc.id,
                    "type":     "function",
                    "function": {
                        "name":      tc.function.name,
                        "arguments": tc.function.arguments,
                    },
                }
                for tc in msg.tool_calls
            ]
            agent.history.append({
                "role":       "assistant",
                "content":    msg.content or "",
                "tool_calls": tool_calls_repr,
            })

            # Execute each tool call and update the panel
            for tc in msg.tool_calls:
                tool_name = tc.function.name
                args_json = tc.function.arguments
                result_json = _tools.call(tool_name, args_json)
                if show_panel:
                    ui.record_tool_call(agent.name, tool_name, args_json, result_json)
                agent.history.append({
                    "role":         "tool",
                    "tool_call_id": tc.id,
                    "content":      result_json,
                })

    except Exception:
        if show_panel:
            ui.end_tool_panel(agent.name)
        raise


# ---------------------------------------------------------------------------
# Core Agent Caller
# Uses the agent's own client — no shared client passed in.
# ---------------------------------------------------------------------------

async def call_agent_async(
    agent: AgentState,
    user_message: str,
    dry_run: bool = False,
    stream: bool = True,
    stream_mode: str = "deliverable",
    use_tools: bool = True,
) -> str:
    """
    Call an agent with a user message and return the response.

    use_tools controls whether the tool-calling loop is engaged.
    Pass use_tools=False to force a plain streaming/non-streaming call
    even when the agent has tools assigned — used for synthesis and
    standing-brief updates where tool dispatch must not occur.
    """
    if dry_run:
        return f"[DRY RUN] Would call {agent.name} ({config.agent_provider_name(agent.name)}) with: {user_message[:100]}..."

    agent.add_user(user_message)

    # Route through the tool loop only when the agent has tools AND the caller
    # has not explicitly suppressed tool use for this call.
    if agent.tools and use_tools:
        try:
            return await _run_tool_loop(agent, stream_mode=stream_mode, dry_run=dry_run)
        except NotImplementedError as e:
            raise RuntimeError(str(e)) from e

    try:
        if stream:
            response = await ui.stream_agent_output(
                agent.name,
                agent.client.call_async_streaming(
                    system=agent.system_prompt,
                    messages=agent.messages_for_call(),
                    model=agent.model,
                    temperature=agent.temperature,
                    max_tokens=agent.max_tokens,
                ),
                mode=stream_mode,
            )
        else:
            response = await agent.client.call_async(
                system=agent.system_prompt,
                messages=agent.messages_for_call(),
                model=agent.model,
                temperature=agent.temperature,
                max_tokens=agent.max_tokens,
            )
    except NotImplementedError as e:
        raise RuntimeError(str(e)) from e

    agent.add_assistant(response)

    if config.DEBUG_PRINT_AGENT_OUTPUTS and not stream:
        ui.print_agent_panel(agent.name, response)

    return response


# ---------------------------------------------------------------------------
# Lead Output Parser
# Extracts standard STATUS / SUMMARY / FINDINGS / OPEN ITEMS sections from
# a lead's response.  Falls back to the raw text if no sections are found.
# ---------------------------------------------------------------------------

def _parse_lead_output(text: str) -> dict:
    """
    Parse a lead's response for the four standard sections.
    Returns a dict with keys: status, summary, findings, open_items, raw.
    Any section not found is an empty string; raw always holds the full text.
    """
    result = {"status": "", "summary": "", "findings": "", "open_items": "", "raw": text}

    # STATUS: single line
    m = re.search(r"^STATUS:\s*(.+)$", text, re.MULTILINE | re.IGNORECASE)
    if m:
        result["status"] = m.group(1).strip()

    # SUMMARY: everything after the label until the next known section or EOF
    m = re.search(
        r"^SUMMARY:\s*\n?(.*?)(?=\n^(?:FINDINGS:|OPEN ITEMS:)|\Z)",
        text, re.MULTILINE | re.DOTALL | re.IGNORECASE,
    )
    if m:
        result["summary"] = m.group(1).strip()

    # FINDINGS: everything after the label until OPEN ITEMS: or EOF
    m = re.search(
        r"^FINDINGS:\s*\n(.*?)(?=\n^OPEN ITEMS:|\Z)",
        text, re.MULTILINE | re.DOTALL | re.IGNORECASE,
    )
    if m:
        result["findings"] = m.group(1).strip()

    # OPEN ITEMS: everything after the label to EOF
    m = re.search(
        r"^OPEN ITEMS:\s*\n?(.*?)$",
        text, re.MULTILINE | re.DOTALL | re.IGNORECASE,
    )
    if m:
        result["open_items"] = m.group(1).strip()

    return result


def _format_lead_output_for_synthesis(task: "Task") -> str:
    """
    Format a completed task for inclusion in Miranda's synthesis prompt.
    Uses parsed sections when present; falls back to raw output.
    """
    lead_title = config.agent_title(task.assigned_to)
    header = f"=== {lead_title} [{task.task_id}] ==="

    if task.status != "complete":
        return f"{header}\nSTATUS: Failed\nERROR: {task.error}"

    p = _parse_lead_output(task.output)
    structured = any([p["status"], p["summary"], p["findings"]])

    if not structured:
        # Lead didn't follow the format — pass raw output with a label
        return f"{header}\n{task.output}"

    parts = [header]
    if p["status"]:
        parts.append(f"STATUS: {p['status']}")
    if p["summary"]:
        parts.append(f"\nSUMMARY:\n{p['summary']}")
    if p["findings"]:
        parts.append(f"\nFINDINGS:\n{p['findings']}")
    if p["open_items"]:
        parts.append(f"\nOPEN ITEMS:\n{p['open_items']}")
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# Lead Brief Formatter
# Converts a dispatch_task spec dict into the structured brief sent to a lead.
# ---------------------------------------------------------------------------

def _format_lead_brief(task_id: str, spec: dict) -> str:
    """Build the task brief message sent to a lead from a dispatch_task spec."""
    lead_name  = spec["to"]
    lead_title = config.agent_title(lead_name)
    lines = [
        f"TASK {task_id}",
        f"From: Miranda, Chief of Staff",
        f"To: {lead_title}",
        "",
        f"DIRECTIVE: {spec['directive']}",
        "",
        f"SCOPE: {spec['scope']}",
    ]
    if spec.get("constraints") and spec["constraints"].upper() != "NONE":
        lines += ["", f"CONSTRAINTS: {spec['constraints']}"]
    if spec.get("dependencies") and spec["dependencies"].upper() != "NONE":
        lines += ["", f"DEPENDENCIES: {spec['dependencies']}"]
    if spec.get("return_format"):
        lines += ["", f"RETURN FORMAT: {spec['return_format']}"]
    lines += [
        "",
        "---",
        "Return your output using your standard STATUS / SUMMARY / FINDINGS / OPEN ITEMS format.",
    ]
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Lead Dispatcher
# ---------------------------------------------------------------------------

async def dispatch_to_lead(
    lead: AgentState,
    task: Task,
    dry_run: bool = False,
    session_id: str = "",
) -> Task:
    task.status = "in_progress"
    if session_id:
        task_store.update_task(task, session_id)
    ui.print_lead_dispatch(lead.name, task.task_id)

    try:
        output = await call_agent_async(
            lead, task.brief, dry_run=dry_run,
            stream=True, stream_mode="lead",
        )
        task.output = output
        task.status = "complete"
        ui.record_lead_output(lead.name, task.task_id, output)
        ui.print_lead_completion(lead.name, task.task_id, len(output))
    except Exception as e:
        task.error = str(e)
        task.status = "error"
        ui.print_error(f"{lead.name} failed on {task.task_id}: {e}")
        ui.print_lead_completion(lead.name, task.task_id, 0, error=True)
    finally:
        if session_id:
            task_store.update_task(task, session_id)
    return task


async def dispatch_all_leads(
    leads: dict[str, AgentState],
    tasks: list[Task],
    dry_run: bool = False,
    session_id: str = "",
) -> list[Task]:
    # Run leads sequentially so each can display its own transient Live panel
    # (Rich does not support nested Live displays).
    results = []
    for task in tasks:
        lead = leads[task.assigned_to]
        result = await dispatch_to_lead(
            lead, task,
            dry_run=dry_run,
            session_id=session_id,
        )
        results.append(result)
    return results


# ---------------------------------------------------------------------------
# Vera Audit
# ---------------------------------------------------------------------------

async def run_auditor(
    auditor: AgentState,
    deliverable: str,
    task_id: str,
    dry_run: bool = False,
) -> str:
    audit_request = (
        f"Please review the following deliverable and return your Audit Memo.\n\n"
        f"Task ID: {task_id}\n\n"
        f"{deliverable}"
    )
    auditor.history = []   # Auditor is stateless — fresh call each time
    return await call_agent_async(auditor, audit_request, dry_run=dry_run, stream_mode="audit")


# ---------------------------------------------------------------------------
# Standing Brief Update
# ---------------------------------------------------------------------------

async def update_standing_brief(
    orchestrator: AgentState,
    session_summary: str,
    dry_run: bool = False,
) -> str:
    current_brief = load_standing_brief()
    update_request = (
        f"The session is closing. Please produce a fully updated Standing Brief.\n\n"
        f"Current brief:\n{current_brief}\n\n"
        f"Session summary:\n{session_summary}\n\n"
        f"Return the complete updated Standing Brief document, ready to save. "
        f"Increment the version number, update the date, and write the new session "
        f"log entry at the top of Section 9. Archive anything older than "
        f"{config.SESSION_LOG_RETENTION} sessions."
    )
    return await call_agent_async(orchestrator, update_request, dry_run=dry_run, stream_mode="background", use_tools=False)


# ---------------------------------------------------------------------------
# Session
# ---------------------------------------------------------------------------

class Session:
    def __init__(self, dry_run: bool = False):
        self.dry_run = dry_run
        self.session_log: list[str] = []
        self.directive_count = 0
        self.session_start = datetime.now()
        self.session_id = self.session_start.strftime("%Y%m%d_%H%M%S")
        self.last_deliverable: str = ""
        self.last_task_id: str = ""
        self.all_tasks: list[Task] = []

        # Runtime flags (togglable via slash commands)
        self.auto_audit = config.VERA_AUTO_AUDIT
        self.debug = config.DEBUG_PRINT_AGENT_OUTPUTS

        config.ensure_runtime_dirs()
        ui.print_system("Loading agents…")
        self.orchestrator = build_agent(config.ORCHESTRATOR_AGENT)
        self.auditor      = build_agent(config.AUDITOR_AGENT)
        self.summarizer   = build_agent(config.SUMMARIZER_AGENT) if config.SUMMARIZER_AGENT else None
        self.leads        = {
            name: build_agent(name) for name in config.LEAD_AGENT_NAMES
        }

        # Inject standing brief into orchestrator's context
        brief = load_standing_brief()
        self.orchestrator.system_prompt = (
            self.orchestrator.system_prompt
            + "\n\n---\n\n## Standing Brief (current)\n\n"
            + brief
        )
        ui.print_system("Session ready.")

    async def run_directive(self, directive: str) -> tuple[str, str]:
        task_id = new_task_id()
        self.directive_count += 1
        self.last_task_id = task_id
        self.session_log.append(f"Directive received: {directive[:120]} [{task_id}]")

        ui.print_task_received(task_id, directive)

        # Safety: clear any leftover queued tasks from a previous aborted directive
        _tools.get_pending_tasks()

        # ── Step 1: Miranda analyses the directive, dispatches via tool calls ─
        cos_prompt = (
            f"The Principal has issued the following directive. "
            f"Analyse it, then either respond directly or use the dispatch_task tool "
            f"to assign work to the appropriate Leads.\n\n"
            f"Directive: {directive}"
        )
        cos_response = await call_agent_async(
            self.orchestrator, cos_prompt, dry_run=self.dry_run, stream=True
        )

        # ── Step 2: Drain task specs queued during Miranda's tool loop ───────
        task_specs = _tools.get_pending_tasks()

        if not task_specs:
            ui.print_direct_response_notice()
            deliverable = cos_response
        else:
            # ── Step 3: Build Task objects, persist, and dispatch ───────────
            tasks = [
                Task(
                    task_id=f"{task_id}-{chr(65 + i)}",
                    directive=directive,
                    assigned_to=spec["to"],
                    brief=_format_lead_brief(f"{task_id}-{chr(65 + i)}", spec),
                )
                for i, spec in enumerate(task_specs)
            ]
            ui.print_task_dispatch_plan(tasks)
            task_store.record_tasks(tasks, self.session_id, directive)
            self.all_tasks.extend(tasks)
            completed_tasks = await dispatch_all_leads(
                self.leads, tasks, dry_run=self.dry_run, session_id=self.session_id,
            )

            # ── Step 4: Miranda synthesizes lead outputs ────────────────────
            lead_outputs_text = "\n\n".join(
                _format_lead_output_for_synthesis(t) for t in completed_tasks
            )
            synthesis_prompt = (
                f"All Lead outputs for {task_id} are in. "
                f"Synthesize them into your SUMMARY / DETAIL / OPEN ITEMS deliverable for the Principal.\n\n"
                f"{lead_outputs_text}"
            )
            deliverable = await call_agent_async(
                self.orchestrator, synthesis_prompt, dry_run=self.dry_run,
                stream=True, use_tools=False,
            )

            dispatched_leads = sorted({t.assigned_to for t in completed_tasks})
            self.session_log.append(
                f"Leads dispatched: {', '.join(dispatched_leads)} | "
                f"Tasks: {', '.join(t.task_id for t in completed_tasks)}"
            )

        self.last_deliverable = deliverable

        # ── Step 5: Auditor + Miranda summary in parallel ───────────────────
        # Vera streams (Rich Live — UI only); Tillie calls the API silently
        # (stream=False, no Rich output), so they do not conflict visually.
        # We use FIRST_COMPLETED so Tillie's panel appears as soon as she
        # resolves, even if Vera is still streaming.
        audit_memo = ""
        if self.auto_audit:
            vera_task   = asyncio.create_task(
                run_auditor(self.auditor, deliverable, task_id, dry_run=self.dry_run)
            )
            tillie_task = asyncio.create_task(
                self._call_summarizer(deliverable)
            )

            done, pending = await asyncio.wait(
                {vera_task, tillie_task},
                return_when=asyncio.FIRST_COMPLETED,
            )

            # If Tillie resolved first, display immediately (Vera may still be streaming)
            miranda_summary = ""
            if tillie_task in done:
                try:
                    miranda_summary = tillie_task.result()
                    if miranda_summary:
                        ui.print_summary(miranda_summary, source_agent=config.ORCHESTRATOR_AGENT)
                except Exception:
                    pass

            # Ensure Vera finishes
            try:
                audit_memo = await vera_task
                self.session_log.append(f"Audit complete for {task_id}.")
            except RuntimeError as e:
                ui.print_warning(f"Audit skipped — {e}")

            # If Vera resolved first, wait for Tillie now (with spinner) then display
            if tillie_task in pending:
                with ui.summary_progress():
                    try:
                        miranda_summary = await tillie_task
                    except Exception:
                        pass
                if miranda_summary:
                    ui.print_summary(miranda_summary, source_agent=config.ORCHESTRATOR_AGENT)
        else:
            await self._summarize(deliverable, source_agent=config.ORCHESTRATOR_AGENT)

        # ── Tillie: summarize Vera's audit memo (sequential — depends on Vera) ─
        if audit_memo:
            await self._summarize(audit_memo, source_agent=config.AUDITOR_AGENT)

        return deliverable, audit_memo

    async def _call_summarizer(self, content: str) -> str:
        """Call Tillie silently and return the summary text. No UI output.
        Used when running in parallel with another agent that owns the display."""
        if not self.summarizer or not content:
            return ""
        self.summarizer.history = []
        return await call_agent_async(
            self.summarizer,
            f"Summarize the following:\n\n{content}",
            dry_run=self.dry_run,
            stream=False,
        ) or ""

    async def _summarize(self, content: str, source_agent: str = "") -> None:
        """Call Tillie with a progress spinner and display the summary panel."""
        if not self.summarizer or not content:
            return
        try:
            with ui.summary_progress():
                summary = await self._call_summarizer(content)
            if summary:
                ui.print_summary(summary, source_agent=source_agent)
        except Exception as e:
            ui.print_warning(f"Summary skipped — {e}")

    async def close(self):
        summary = "\n".join(self.session_log)
        updated_brief = await update_standing_brief(
            self.orchestrator, summary, dry_run=self.dry_run
        )
        if not self.dry_run:
            save_standing_brief(updated_brief)
        ui.print_session_footer(self.directive_count)


# ---------------------------------------------------------------------------
# Slash Command Registry
# ---------------------------------------------------------------------------

COMMANDS: list[tuple[list[str], str, object]] = []

def command(names: list[str], description: str):
    def decorator(fn):
        COMMANDS.append((names, description, fn))
        return fn
    return decorator


@command(["/help", "/?"], "Show this command list")
async def cmd_help(session: Session, args: str) -> bool:
    rows = [(", ".join(names), desc) for names, desc, _ in COMMANDS]
    ui.print_help(rows)
    return True


@command(["/status"], "Session stats, settings, and log")
async def cmd_status(session: Session, args: str) -> bool:
    ui.print_status(
        session_start=session.session_start,
        directive_count=session.directive_count,
        session_log=session.session_log,
        auto_audit=session.auto_audit,
        debug=session.debug,
        dry_run=session.dry_run,
    )

    return True


@command(["/brief"], "Show the current standing brief")
async def cmd_brief(session: Session, args: str) -> bool:
    ui.print_standing_brief(load_standing_brief())
    return True


@command(["/reset-brief"], "Replace the standing brief with the blank template")
async def cmd_reset_brief(session: Session, args: str) -> bool:
    if not config.BRIEF_TEMPLATE_PATH.exists():
        ui.print_error(f"Template not found: {config.BRIEF_TEMPLATE_PATH}")
        return True
    confirm = ui.console.input(
        "  [warning]This will overwrite the standing brief. Type YES to confirm:[/warning] "
    ).strip()
    if confirm != "YES":
        ui.print_system("Reset cancelled.")
        return True
    template = config.BRIEF_TEMPLATE_PATH.read_text(encoding="utf-8")
    save_standing_brief(template)
    # Re-inject the fresh brief into the orchestrator's context for this session
    session.orchestrator.system_prompt = (
        load_system_prompt(config.AGENT_CONFIGS[config.ORCHESTRATOR_AGENT]["prompt_file"])
        + "\n\n---\n\n## Standing Brief (current)\n\n"
        + template
    )
    ui.print_system("Standing brief reset to template.")
    return True


@command(["/agents"], "Show agent roster with provider and model assignments")
async def cmd_agents(session: Session, args: str) -> bool:
    ui.print_agents(config.AGENT_CONFIGS, config.ACTIVE_PROVIDER)
    return True


@command(["/history"], "Show directives processed this session")
async def cmd_history(session: Session, args: str) -> bool:
    ui.print_history(session.session_log)
    return True


@command(["/reload"], "Reload agents.yaml and rebuild the agent roster for this session")
async def cmd_reload(session: Session, args: str) -> bool:
    ui.print_system("Reloading agents configuration…")
    config.reload_agents_config()
    # Rebuild all agents with fresh config
    session.orchestrator = build_agent(config.ORCHESTRATOR_AGENT)
    session.auditor      = build_agent(config.AUDITOR_AGENT)
    session.summarizer   = build_agent(config.SUMMARIZER_AGENT) if config.SUMMARIZER_AGENT else None
    session.leads        = {
        name: build_agent(name) for name in config.LEAD_AGENT_NAMES
    }
    # Re-inject standing brief into orchestrator
    brief = load_standing_brief()
    session.orchestrator.system_prompt = (
        session.orchestrator.system_prompt
        + "\n\n---\n\n## Standing Brief (current)\n\n"
        + brief
    )
    ui.print_system("Agents reloaded.")
    return True


@command(["/thinking"], "Browse thinking output history for this session")
async def cmd_thinking(session: Session, args: str) -> bool:
    await ui.show_thinking_picker()
    return True


@command(["/leads"], "Browse completed lead outputs from this session")
async def cmd_leads(session: Session, args: str) -> bool:
    await ui.show_lead_output_picker()
    return True


@command(["/audit"], "Re-run Vera's audit on the last deliverable")
async def cmd_audit(session: Session, args: str) -> bool:
    if not session.last_deliverable:
        ui.print_warning("No deliverable yet — issue a directive first.")
        return True
    ui.print_system("Running audit on last deliverable…")
    memo = await run_auditor(
        session.auditor, session.last_deliverable,
        session.last_task_id, dry_run=session.dry_run,
    )
    ui.print_audit_memo(memo, task_id=session.last_task_id)
    if memo:
        await session._summarize(memo, source_agent=config.AUDITOR_AGENT)
    return True


@command(["/debug"], "Toggle raw agent output")
async def cmd_debug(session: Session, args: str) -> bool:
    session.debug = not session.debug
    config.DEBUG_PRINT_AGENT_OUTPUTS = session.debug
    ui.print_debug_toggle(session.debug)
    return True


@command(["/tasks"], "Show tasks dispatched this session and their status")
async def cmd_tasks(session: Session, args: str) -> bool:
    ui.print_tasks(session.all_tasks, session_id=session.session_id)
    return True


@command(["/autoaudit"], "Toggle automatic Vera audit after each directive")
async def cmd_autoaudit(session: Session, args: str) -> bool:
    session.auto_audit = not session.auto_audit
    ui.print_audit_toggle(session.auto_audit)
    return True


@command(["/clear"], "Clear the terminal")
async def cmd_clear(session: Session, args: str) -> bool:
    ui.console.clear()
    ui.print_session_header()
    return True


@command(["/quit", "/exit", "/q"], "Close the session and update the standing brief")
async def cmd_quit(session: Session, args: str) -> bool:
    return False


def _dispatch_command(name: str):
    name_lower = name.lower()
    for names, _, handler in COMMANDS:
        if name_lower in names:
            return handler
    return None


# Register all slash command names for tab completion.
# This runs once at import time, after all @command decorators have executed.
ui.set_completions([(name, desc) for names, desc, _ in COMMANDS for name in names])


# ---------------------------------------------------------------------------
# Direct Agent Session
# Opens an interactive REPL with one named agent.
# ---------------------------------------------------------------------------

async def direct_agent_session(agent_name: str, dry_run: bool = False):
    if agent_name not in config.AGENT_CONFIGS:
        ui.print_error(f"Unknown agent: {agent_name!r}")
        ui.print_system(f"Available agents: {', '.join(config.AGENT_CONFIGS)}")
        return

    agent = build_agent(agent_name)
    provider = config.agent_provider_name(agent_name)

    ui.print_session_header()
    ui.print_system(
        f"Direct session with {agent_name.capitalize()} "
        f"[dim]({provider} / {agent.model or 'no model set'})[/dim]. "
        f"Type /quit to end."
    )

    while True:
        try:
            user_input = await ui.get_directive()
        except (EOFError, KeyboardInterrupt):
            break

        if not user_input or user_input.lower() in ("/quit", "/exit", "/q"):
            break

        if user_input.lower() == "/thinking":
            await ui.show_thinking_picker()
            continue

        response = await call_agent_async(agent, user_input, dry_run=dry_run, stream=True)
        ui.print_agent_panel(agent_name, response)

    ui.print_session_footer(0)


# ---------------------------------------------------------------------------
# Principal REPL
# ---------------------------------------------------------------------------

async def principal_session(dry_run: bool = False):
    ui.print_session_header()

    if dry_run:
        ui.print_dry_run_banner()

    session = Session(dry_run=dry_run)

    while True:
        try:
            raw = await ui.get_directive()
        except (EOFError, KeyboardInterrupt):
            ui.console.print()
            break

        if not raw:
            continue

        # ── Slash command dispatch ──────────────────────────────────────────
        if raw.startswith("/"):
            parts = raw.split(None, 1)
            cmd_name = parts[0]
            cmd_args = parts[1] if len(parts) > 1 else ""
            handler = _dispatch_command(cmd_name)
            if handler is None:
                ui.print_error(f"Unknown command: {cmd_name}  (try /help)")
                continue
            should_continue = await handler(session, cmd_args)
            if not should_continue:
                break
            continue

        # ── Directive ───────────────────────────────────────────────────────
        try:
            deliverable, audit_memo = await session.run_directive(raw)
        except RuntimeError as e:
            ui.print_error(str(e))
            break

        ui.print_deliverable(deliverable, task_id=session.last_task_id)

        if audit_memo:
            ui.print_audit_memo(audit_memo, task_id=session.last_task_id)

    await session.close()


# ---------------------------------------------------------------------------
# Entry Point
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="AIPA Orchestrator — run a Principal session."
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="Trace routing without making provider API calls.",
    )
    parser.add_argument(
        "--agent", type=str, metavar="NAME",
        help="Open a direct session with a named agent (e.g. evelyn, atlas).",
    )
    args = parser.parse_args()

    if args.agent:
        asyncio.run(direct_agent_session(args.agent, dry_run=args.dry_run))
    else:
        asyncio.run(principal_session(dry_run=args.dry_run))


if __name__ == "__main__":
    main()