tui.py

"""Support Intelligence Agent — Interactive TUI.

Run: .venv/bin/python3 tui.py

Toolbox integration (optional):
  Set in .env:
    FOUNDRY_PROJECT_ENDPOINT=https://<account>.services.ai.azure.com/api/projects/<project>
    TOOLBOX_NAME=Support          # default
    TOOLBOX_VERSION=1             # default; omit for latest
  Or set the full URL directly:
    TOOLBOX_ENDPOINT=https://...toolboxes/Support/versions/1/mcp?api-version=v1
"""

import asyncio
import json
import logging
import os
import subprocess
import sys
import time
from pathlib import Path
from types import SimpleNamespace

import httpx
from dotenv import load_dotenv

load_dotenv(override=False)

from textual.app import App, ComposeResult
from textual.containers import Vertical, VerticalScroll
from textual.screen import ModalScreen
from textual.widgets import Footer, Header, LoadingIndicator, Markdown, Static, TextArea
from textual.reactive import reactive
from textual import work

from copilot import CopilotClient
from copilot.generated.session_events import SessionEvent, SessionEventType
from copilot.session import PermissionHandler
from copilot.tools import Tool, ToolInvocation, ToolResult

import tools as _tools_pkg

# When running locally, set HOME to the project dir so $HOME/workspace
# lands next to the source code (matches hosted behavior where $HOME is
# the container's persistent volume). Must run BEFORE importing tool
# modules, which bind WORKSPACE into their own module scope at import time.
if not os.environ.get("FOUNDRY_PROJECT_ENDPOINT"):
    _local_home = str(Path(__file__).parent)
    os.environ["HOME"] = _local_home
    _tools_pkg.WORKSPACE = Path(_local_home) / "workspace"

from tools import WORKSPACE
from tools.support_cases_tool import get_support_cases
from tools.workiq_tool import query_workiq
from tools.trend_forecast_tool import generate_trend_forecast

logger = logging.getLogger("tui")


# ── Foundry Toolbox MCP Bridge ─────────────────────────────────────────────────


def _resolve_toolbox_url() -> str | None:
    """Build the toolbox MCP endpoint URL from env vars, or return None."""
    # Option 1: full URL provided directly
    url = os.environ.get("TOOLBOX_ENDPOINT")
    if url:
        return url

    # Option 2: build from components
    endpoint = os.environ.get("FOUNDRY_PROJECT_ENDPOINT")
    if not endpoint:
        return None

    name = os.environ.get("TOOLBOX_NAME", "Support")
    version = os.environ.get("TOOLBOX_VERSION")
    if version:
        return f"{endpoint.rstrip('/')}/toolboxes/{name}/versions/{version}/mcp?api-version=v1"
    return f"{endpoint.rstrip('/')}/toolboxes/{name}/mcp?api-version=v1"


def _get_toolbox_token() -> str:
    """Get a bearer token for the toolbox MCP endpoint (scope: https://ai.azure.com/.default)."""
    from azure.identity import DefaultAzureCredential
    credential = DefaultAzureCredential()
    return credential.get_token("https://ai.azure.com/.default").token


def _get_toolbox_headers(token: str) -> dict:
    return {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {token}",
        "Foundry-Features": "Toolboxes=V1Preview",
    }


class McpBridge:
    """HTTP-based MCP client that connects to a Foundry toolbox MCP endpoint."""

    def __init__(self, endpoint: str, token: str):
        self.endpoint = endpoint
        self.headers = _get_toolbox_headers(token)
        self._session_id: str | None = None
        self._client = httpx.AsyncClient(timeout=60.0)
        self._req_id = 0

    def _next_id(self) -> int:
        self._req_id += 1
        return self._req_id

    def _request_headers(self) -> dict:
        headers = dict(self.headers)
        if self._session_id:
            headers["mcp-session-id"] = self._session_id
        return headers

    async def initialize(self) -> str:
        resp = await self._client.post(
            self.endpoint, headers=self.headers,
            json={
                "jsonrpc": "2.0", "id": self._next_id(), "method": "initialize",
                "params": {
                    "protocolVersion": "2025-03-26",
                    "capabilities": {},
                    "clientInfo": {"name": "copilot-toolbox-bridge", "version": "1.0.0"},
                },
            },
        )
        resp.raise_for_status()
        data = resp.json()
        self._session_id = resp.headers.get("mcp-session-id")

        await self._client.post(
            self.endpoint, headers=self._request_headers(),
            json={"jsonrpc": "2.0", "method": "notifications/initialized"},
        )
        return data.get("result", {}).get("serverInfo", {}).get("name", "unknown")

    async def list_tools(self) -> list[dict]:
        resp = await self._client.post(
            self.endpoint, headers=self._request_headers(),
            json={"jsonrpc": "2.0", "id": self._next_id(), "method": "tools/list", "params": {}},
        )
        resp.raise_for_status()
        return resp.json().get("result", {}).get("tools", [])

    async def call_tool(self, name: str, arguments: dict) -> str:
        resp = await self._client.post(
            self.endpoint, headers=self._request_headers(),
            json={
                "jsonrpc": "2.0", "id": self._next_id(), "method": "tools/call",
                "params": {"name": name, "arguments": arguments},
            },
        )
        resp.raise_for_status()
        result = resp.json().get("result", {})
        content = result.get("content", [])
        texts = [c.get("text", "") for c in content if isinstance(c, dict) and c.get("type") == "text"]
        return "\n".join(t for t in texts if t) or json.dumps(result)

    async def close(self):
        await self._client.aclose()


def _make_copilot_tools(bridge: McpBridge, mcp_tools: list[dict]) -> list[Tool]:
    """Convert MCP tool definitions into Copilot SDK Tool objects."""
    tools = []
    for mcp_tool in mcp_tools:
        mcp_name = mcp_tool["name"]
        sdk_name = mcp_name.replace(".", "_").replace("-", "_")
        desc = mcp_tool.get("description", f"MCP tool: {mcp_name}")
        schema = mcp_tool.get("inputSchema", {"type": "object", "properties": {}})

        def _make_handler(original_name):
            async def handler(invocation: ToolInvocation) -> ToolResult:
                args = invocation.arguments if isinstance(invocation.arguments, dict) else {}
                try:
                    result_text = await bridge.call_tool(original_name, args)
                    return ToolResult(text_result_for_llm=result_text)
                except Exception as e:
                    logger.warning("Toolbox tool %s failed: %s", original_name, e)
                    return ToolResult(text_result_for_llm="", result_type="error", error=str(e))
            return handler

        tools.append(Tool(
            name=sdk_name,
            description=desc,
            parameters=schema,
            handler=_make_handler(mcp_name),
            skip_permission=True,
        ))
    return tools


SYSTEM_MESSAGE = """You are a Support Intelligence Agent for an enterprise SaaS platform. Your job is to analyze support case data, identify patterns and trends, correlate with internal team knowledge, and generate actionable forecasts.

When asked to analyze support cases, follow this workflow:
1. Pull the support case data using the get_support_cases tool
2. Analyze the data for patterns — look for spikes, clusters by category/region/customer, severity distributions
3. Search the web for any known issues, CVEs, or outage reports related to the top categories you identified
4. Check WorkIQ for internal team discussions about recurring themes (use specific queries like "authentication failures" or "API timeout")
5. Generate a comprehensive trend forecast with all artifacts using the generate_trend_forecast tool

Always announce what you're doing at each step. When writing files, explicitly state:
"Writing [description] to ~/workspace/[filename]..."

Be thorough but concise in your analysis. Highlight the most actionable insights first."""



# ── Custom Widgets ─────────────────────────────────────────────────────────────


class PromptInput(TextArea):
    """TextArea that submits on Enter, inserts newline on Shift+Enter."""

    DEFAULT_CSS = """
    PromptInput {
        height: 3;
    }
    """

    async def _on_key(self, event) -> None:
        if event.key == "enter":
            event.prevent_default()
            event.stop()
            await self.app._submit_prompt()
        elif event.key == "shift+enter":
            event.prevent_default()
            event.stop()
            self.insert("\n")
        else:
            await super()._on_key(event)


class UserMessage(Static):
    DEFAULT_CSS = """
    UserMessage {
        background: $primary-darken-2;
        color: $text;
        padding: 0 1;
        margin: 1 0 0 0;
    }
    """

    def __init__(self, text: str) -> None:
        super().__init__(f"[bold]You:[/bold] {text}")


class ToolCallNotice(Static):
    DEFAULT_CSS = """
    ToolCallNotice {
        color: $text-muted;
        text-style: italic;
        padding: 0 1;
        margin: 0;
    }
    """


class ArtifactButton(Static):
    DEFAULT_CSS = """
    ArtifactButton {
        color: $accent;
        text-style: underline;
        padding: 0 1;
    }
    ArtifactButton:hover {
        background: $accent-darken-2;
        color: $text;
    }
    """

    def __init__(self, label: str, artifact_path: str) -> None:
        super().__init__(label)
        self.artifact_path = artifact_path

    def on_click(self) -> None:
        path = Path(self.artifact_path)
        if path.suffix.lower() == ".md":
            self.app.push_screen(ArtifactViewScreen(path))
        else:
            try:
                subprocess.Popen(
                    ["xdg-open", str(path)],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )
            except FileNotFoundError:
                self.app.push_screen(ArtifactViewScreen(path))


class ArtifactViewScreen(ModalScreen[None]):
    BINDINGS = [("escape", "close", "Close")]

    def __init__(self, path: Path) -> None:
        super().__init__()
        self.artifact_path = path

    def compose(self) -> ComposeResult:
        content = self.artifact_path.read_text(errors="replace")
        yield Vertical(
            Static(f" {self.artifact_path.name} ", id="artifact-viewer-title"),
            VerticalScroll(Markdown(content)),
            id="artifact-viewer",
        )

    def action_close(self) -> None:
        self.dismiss()


# ── Main Application ───────────────────────────────────────────────────────────


class SupportIntelligenceTUI(App):
    TITLE = "Support Intelligence Agent"
    CSS_PATH = "tui.tcss"
    BINDINGS = [
        ("ctrl+c", "quit", "Quit"),
        ("ctrl+l", "clear_conversation", "Clear"),
    ]

    state: reactive[str] = reactive("idle")

    def __init__(self) -> None:
        super().__init__()
        self._client: CopilotClient | None = None
        self._session = None
        self._bridge: McpBridge | None = None
        self._stream_queue: asyncio.Queue | None = None
        self._accumulated_text: str = ""
        self._current_assistant_widget: Markdown | None = None

    def compose(self) -> ComposeResult:
        yield Header()
        yield VerticalScroll(id="conversation")
        yield Vertical(
            LoadingIndicator(id="spinner"),
            PromptInput(id="prompt-input", soft_wrap=True),
            id="input-area",
        )
        yield Footer()

    async def on_mount(self) -> None:
        self._ensure_github_token()
        self._clean_workspace()

        # Connect to Foundry toolbox if configured
        toolbox_tools: list[Tool] = []
        toolbox_tool_names: set[str] = set()
        toolbox_url = _resolve_toolbox_url()
        if toolbox_url:
            try:
                token = _get_toolbox_token()
                self._bridge = McpBridge(toolbox_url, token)
                server_name = await self._bridge.initialize()
                mcp_tools = await self._bridge.list_tools()
                toolbox_tools = _make_copilot_tools(self._bridge, mcp_tools)
                toolbox_tool_names = {t.name for t in toolbox_tools}
                logger.info("Toolbox '%s' connected: %d tool(s)", server_name, len(toolbox_tools))
            except Exception as e:
                logger.warning("Toolbox connection failed (%s), using local tools only: %s", toolbox_url, e)

        # Local tools — skip mock query_workiq if the toolbox provides a real version
        local_tools = [get_support_cases, generate_trend_forecast]
        if not toolbox_tool_names:
            local_tools.append(query_workiq)

        all_tools = toolbox_tools + local_tools

        self._client = CopilotClient()
        await self._client.__aenter__()

        model = os.environ.get("GITHUB_COPILOT_MODEL", "gpt-4.1")
        self._session = await self._client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            model=model,
            streaming=True,
            tools=all_tools,
            system_message={"mode": "replace", "content": SYSTEM_MESSAGE},
        )

        self.query_one("#prompt-input", PromptInput).focus()

    async def on_unmount(self) -> None:
        if self._session:
            try:
                await self._session.disconnect()
            except Exception:
                pass
        if self._client:
            try:
                await self._client.__aexit__(None, None, None)
            except Exception:
                pass
        if self._bridge:
            try:
                await self._bridge.close()
            except Exception:
                pass

    def _ensure_github_token(self) -> None:
        if os.environ.get("GITHUB_TOKEN"):
            return
        try:
            result = subprocess.run(
                ["gh", "auth", "token"], capture_output=True, text=True, timeout=10
            )
            if result.returncode == 0 and result.stdout.strip():
                os.environ["GITHUB_TOKEN"] = result.stdout.strip()
            else:
                self.exit(message="Missing GITHUB_TOKEN. Set it in .env or run 'gh auth login'.")
        except Exception:
            self.exit(message="Missing GITHUB_TOKEN. Set it in .env or run 'gh auth login'.")

    def _clean_workspace(self) -> None:
        WORKSPACE.mkdir(parents=True, exist_ok=True)
        for f in WORKSPACE.glob("*"):
            if f.name != ".gitkeep":
                f.unlink()

    # ── State Watcher ──────────────────────────────────────────────────────────

    def watch_state(self, new_state: str) -> None:
        spinner = self.query_one("#spinner", LoadingIndicator)
        prompt_input = self.query_one("#prompt-input", PromptInput)

        if new_state == "thinking":
            spinner.display = True
            prompt_input.disabled = True
        elif new_state == "streaming":
            spinner.display = False
            prompt_input.disabled = True
        else:  # idle
            spinner.display = False
            prompt_input.disabled = False
            prompt_input.focus()

    # ── Input Handling ─────────────────────────────────────────────────────────

    def on_text_area_changed(self, event: TextArea.Changed) -> None:
        """Auto-resize the text area height based on content."""
        ta = event.text_area
        if ta.id == "prompt-input":
            line_count = ta.document.line_count
            ta.styles.height = min(max(3, line_count + 1), 10)

    async def _submit_prompt(self) -> None:
        prompt_input = self.query_one("#prompt-input", PromptInput)
        prompt = prompt_input.text.strip()
        if not prompt or self.state != "idle":
            return
        prompt_input.clear()
        prompt_input.styles.height = 3

        conversation = self.query_one("#conversation", VerticalScroll)
        await conversation.mount(UserMessage(prompt))

        self._current_assistant_widget = None
        self._accumulated_text = ""

        self.state = "thinking"
        self._run_agent(prompt)

    # ── Agent Worker ───────────────────────────────────────────────────────────

    def _create_event_handler(self) -> callable:
        queue = self._stream_queue

        def handler(event: SessionEvent) -> None:
            etype = event.type

            if etype == SessionEventType.ASSISTANT_MESSAGE_DELTA:
                delta = getattr(event.data, "delta_content", None) or ""
                if delta:
                    queue.put_nowait(SimpleNamespace(kind="delta", text=delta))

            elif etype == SessionEventType.TOOL_EXECUTION_START:
                name = (
                    getattr(event.data, "tool_name", None)
                    or getattr(event.data, "mcp_tool_name", None)
                    or "tool"
                )
                queue.put_nowait(SimpleNamespace(kind="tool_start", text=name))

            elif etype == SessionEventType.TOOL_EXECUTION_COMPLETE:
                name = getattr(event.data, "tool_name", None) or "tool"
                queue.put_nowait(SimpleNamespace(kind="tool_complete", text=name))

            elif etype == SessionEventType.SESSION_IDLE:
                queue.put_nowait(SimpleNamespace(kind="done", text=""))

            elif etype == SessionEventType.SESSION_ERROR:
                msg = getattr(event.data, "message", None) or "Unknown error"
                queue.put_nowait(SimpleNamespace(kind="error", text=msg))

        return handler

    @work(exclusive=True)
    async def _run_agent(self, prompt: str) -> None:
        self._stream_queue = asyncio.Queue()
        unsubscribe = self._session.on(self._create_event_handler())

        last_render = 0.0
        render_interval = 0.05  # 50ms throttle
        conversation = self.query_one("#conversation", VerticalScroll)

        async def _ensure_markdown_widget():
            """Create a new Markdown widget if we don't have one for the current text segment."""
            if self._current_assistant_widget is None:
                widget = Markdown("")
                await conversation.mount(widget)
                self._current_assistant_widget = widget
                self._accumulated_text = ""

        async def _flush_text():
            """Render any pending accumulated text."""
            if self._current_assistant_widget and self._accumulated_text:
                self._current_assistant_widget.update(self._accumulated_text)
                conversation.scroll_end(animate=False)

        try:
            await self._session.send(prompt)

            while True:
                try:
                    item = await asyncio.wait_for(self._stream_queue.get(), timeout=0.1)
                except asyncio.TimeoutError:
                    if self._accumulated_text:
                        now = time.monotonic()
                        if now - last_render > render_interval:
                            await _flush_text()
                            last_render = now
                    continue

                if item.kind == "delta":
                    await _ensure_markdown_widget()
                    self._accumulated_text += item.text
                    self.state = "streaming"
                    now = time.monotonic()
                    if now - last_render > render_interval:
                        await _flush_text()
                        last_render = now

                elif item.kind == "tool_start":
                    # Flush any text before this tool call
                    await _flush_text()
                    # Close current text segment — next deltas get a fresh widget
                    self._current_assistant_widget = None

                    notice = ToolCallNotice(f"  > Calling {item.text}...")
                    await conversation.mount(notice)
                    conversation.scroll_end(animate=False)

                elif item.kind == "tool_complete":
                    notice = ToolCallNotice(f"  > {item.text} done")
                    await conversation.mount(notice)
                    conversation.scroll_end(animate=False)

                elif item.kind == "error":
                    await _ensure_markdown_widget()
                    self._accumulated_text += f"\n\n**Error:** {item.text}"
                    await _flush_text()
                    break

                elif item.kind == "done":
                    await _flush_text()
                    break

        finally:
            unsubscribe()

        await self._show_artifacts()
        self.state = "idle"

    # ── Artifact Display ───────────────────────────────────────────────────────

    async def _show_artifacts(self) -> None:
        files = sorted(f for f in WORKSPACE.iterdir() if f.name != ".gitkeep" and f.is_file())
        if not files:
            return

        conversation = self.query_one("#conversation", VerticalScroll)
        await conversation.mount(Static("───────────────────────────────────"))
        await conversation.mount(Static("[bold green]Generated Artifacts:[/bold green]"))

        for f in files:
            size = f.stat().st_size
            ext = f.suffix.lower().lstrip(".")
            icon = {"md": "\U0001f4c4", "csv": "\U0001f4ca", "png": "\U0001f5bc\ufe0f"}.get(ext, "\U0001f4ce")
            label = f"  {icon} {f.name} ({size:,} bytes)"
            btn = ArtifactButton(label, artifact_path=str(f))
            await conversation.mount(btn)

        conversation.scroll_end(animate=False)

    # ── Actions ────────────────────────────────────────────────────────────────

    def action_clear_conversation(self) -> None:
        conversation = self.query_one("#conversation", VerticalScroll)
        conversation.remove_children()
        self._clean_workspace()


if __name__ == "__main__":
    app = SupportIntelligenceTUI()
    app.run()