"""
Codex MCP adapter for sync conversations.

Uses codex mcp-server for true iterative conversations:
- codex() to start a session with conversationId
- codex-reply() to continue the conversation
"""

from __future__ import annotations

import json
import logging
import queue
import subprocess
import threading
import time
from pathlib import Path
from typing import Any, Literal

import weave

from zwarm.adapters.base import ExecutorAdapter
from zwarm.adapters.registry import register_adapter
from zwarm.core.models import (
    ConversationSession,
    SessionMode,
    SessionStatus,
)

logger = logging.getLogger(__name__)


class MCPClient:
    """
    Robust MCP client for communicating with codex mcp-server.

    Uses subprocess.Popen (NOT asyncio.subprocess) to avoid being tied to
    any specific event loop. This allows the MCP server to stay alive across
    multiple asyncio.run() calls, preserving conversation state.

    Uses dedicated reader threads that queue lines, avoiding the race condition
    of spawning new reader threads on timeout.
    """

    def __init__(self):
        self._proc: subprocess.Popen | None = None
        self._proc_pid: int | None = None  # Track PID to detect restarts
        self._request_id = 0
        self._initialized = False
        self._stderr_thread: threading.Thread | None = None
        self._stdout_thread: threading.Thread | None = None
        self._stderr_lines: list[str] = []
        self._stdout_queue: queue.Queue[str | None] = queue.Queue()
        self._lock = threading.Lock()  # Protect writes only
        self._start_count = 0  # Track how many times we've started

    def start(self) -> None:
        """Start the MCP server process."""
        with self._lock:
            if self._proc is not None and self._proc.poll() is None:
                logger.debug(f"MCP server already running (pid={self._proc.pid}, start_count={self._start_count})")
                return  # Already running

            # Check if this is a restart (previous server died)
            if self._proc_pid is not None:
                logger.warning(
                    f"MCP server restart detected! Previous pid={self._proc_pid}, "
                    f"start_count={self._start_count}. All conversation state will be lost."
                )

            self._start_count += 1
            logger.info(f"Starting codex mcp-server... (start_count={self._start_count})")
            self._proc = subprocess.Popen(
                ["codex", "mcp-server"],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=False,  # Binary mode for explicit encoding control
            )
            self._proc_pid = self._proc.pid
            self._initialized = False
            self._stderr_lines = []
            self._stdout_queue = queue.Queue()  # Fresh queue

            # Start background thread to read stderr
            self._stderr_thread = threading.Thread(
                target=self._read_stderr_loop,
                daemon=True,
                name="mcp-stderr-reader",
            )
            self._stderr_thread.start()

            # Start background thread to read stdout into queue
            self._stdout_thread = threading.Thread(
                target=self._read_stdout_loop,
                daemon=True,
                name="mcp-stdout-reader",
            )
            self._stdout_thread.start()

            logger.info(f"MCP server started (pid={self._proc.pid})")

    def _read_stderr_loop(self) -> None:
        """Background thread to read stderr and log errors."""
        if not self._proc or not self._proc.stderr:
            return
        try:
            while True:
                line = self._proc.stderr.readline()
                if not line:
                    break
                decoded = line.decode().strip()
                if decoded:
                    self._stderr_lines.append(decoded)
                    # Keep only last 100 lines
                    if len(self._stderr_lines) > 100:
                        self._stderr_lines = self._stderr_lines[-100:]
                    # Log errors prominently
                    if "error" in decoded.lower() or "ERROR" in decoded:
                        logger.error(f"[MCP stderr] {decoded}")
                    else:
                        logger.debug(f"[MCP stderr] {decoded}")
        except Exception as e:
            logger.warning(f"stderr reader stopped: {e}")

    def _read_stdout_loop(self) -> None:
        """Background thread to read stdout and queue lines."""
        if not self._proc or not self._proc.stdout:
            return
        try:
            while True:
                line = self._proc.stdout.readline()
                if not line:
                    # EOF - signal end
                    self._stdout_queue.put(None)
                    break
                decoded = line.decode()
                self._stdout_queue.put(decoded)
        except Exception as e:
            logger.warning(f"stdout reader stopped: {e}")
            self._stdout_queue.put(None)  # Signal error

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

    def _write(self, data: str) -> None:
        """Write to stdin with error handling."""
        if not self._proc or not self._proc.stdin:
            raise RuntimeError("MCP server not running")
        if self._proc.poll() is not None:
            raise RuntimeError(f"MCP server died (exit code {self._proc.returncode})")

        self._proc.stdin.write(data.encode())
        self._proc.stdin.flush()

    def _read_line(self, timeout: float = 120.0) -> str:
        """
        Read a line from the stdout queue with timeout.

        Uses a dedicated reader thread that queues lines, so we never
        lose data on timeout - we just haven't received it yet.
        """
        if not self._proc:
            raise RuntimeError("MCP server not running")

        try:
            line = self._stdout_queue.get(timeout=timeout)
        except queue.Empty:
            # Timeout - check process health
            if self._proc.poll() is not None:
                stderr_context = "\n".join(self._stderr_lines[-10:]) if self._stderr_lines else "(no stderr)"
                raise RuntimeError(
                    f"MCP server died (exit code {self._proc.returncode}).\n"
                    f"Recent stderr:\n{stderr_context}"
                )
            # Process still alive, just slow - return empty to let caller decide
            return ""

        if line is None:
            # EOF or error from reader thread
            stderr_context = "\n".join(self._stderr_lines[-10:]) if self._stderr_lines else "(no stderr)"
            if self._proc.poll() is not None:
                raise RuntimeError(
                    f"MCP server exited (code {self._proc.returncode}).\n"
                    f"Recent stderr:\n{stderr_context}"
                )
            raise RuntimeError(f"MCP stdout closed unexpectedly.\nRecent stderr:\n{stderr_context}")

        return line

    def _check_alive(self) -> None:
        """Check if the MCP server is still alive, raise if not."""
        if not self._proc:
            raise RuntimeError("MCP server not started")
        if self._proc.poll() is not None:
            stderr_context = "\n".join(self._stderr_lines[-10:]) if self._stderr_lines else "(no stderr)"
            raise RuntimeError(
                f"MCP server died (exit code {self._proc.returncode}).\n"
                f"Recent stderr:\n{stderr_context}"
            )

    def initialize(self) -> dict:
        """Initialize MCP connection."""
        self._check_alive()

        request = {
            "jsonrpc": "2.0",
            "id": self._next_id(),
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "zwarm", "version": "0.1.0"},
            },
        }
        with self._lock:
            self._write(json.dumps(request) + "\n")

        response_line = self._read_line(timeout=30.0)
        if not response_line:
            raise RuntimeError("No response from MCP server during init")

        response = json.loads(response_line)
        if "error" in response:
            raise RuntimeError(f"MCP init error: {response['error']}")

        # Send initialized notification
        notif = {"jsonrpc": "2.0", "method": "notifications/initialized"}
        with self._lock:
            self._write(json.dumps(notif) + "\n")

        self._initialized = True
        logger.info("MCP connection initialized")
        return response

    def call_tool(self, name: str, arguments: dict, timeout: float = 300.0) -> dict:
        """
        Call an MCP tool and collect streaming events.

        Args:
            name: Tool name (codex, codex-reply)
            arguments: Tool arguments
            timeout: Overall timeout for the call (default 5 min)
        """
        self._check_alive()

        if not self._initialized:
            self.initialize()

        request_id = self._next_id()
        request = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": "tools/call",
            "params": {"name": name, "arguments": arguments},
        }

        logger.debug(f"Calling MCP tool: {name} with args: {list(arguments.keys())}")
        with self._lock:
            self._write(json.dumps(request) + "\n")

        # Collect streaming events until final result
        # Reader thread queues lines, we pull from queue with timeout
        session_id = None
        conversation_id = None  # Track conversation ID separately
        agent_messages: list[str] = []
        streaming_text: list[str] = []  # Accumulate streaming delta text
        final_result = None
        token_usage: dict[str, Any] = {}  # Track token usage
        start_time = time.time()
        all_events: list[dict] = []  # Keep ALL events for debugging

        for event_count in range(1000):  # Safety limit on events
            self._check_alive()

            # Check overall timeout
            elapsed = time.time() - start_time
            if elapsed > timeout:
                raise RuntimeError(f"MCP call timed out after {timeout}s ({event_count} events received)")

            # Read from queue with per-event timeout
            # Empty string = timeout (process still alive, just waiting)
            # None sentinel is handled inside _read_line (raises RuntimeError)
            line = self._read_line(timeout=30.0)

            if not line:
                # Timeout waiting for event - process is still alive, just slow
                # This is normal during long codex operations
                logger.debug(f"Waiting for MCP event... (elapsed: {elapsed:.0f}s, events: {event_count})")
                continue

            try:
                event = json.loads(line)
                all_events.append(event)  # Keep for debugging
            except json.JSONDecodeError as e:
                logger.warning(f"Invalid JSON from MCP: {line[:100]}... - {e}")
                continue

            # Check for final result (has matching id)
            if event.get("id") == request_id:
                if "result" in event:
                    final_result = event["result"]
                    # Extract conversation ID from final result
                    if isinstance(final_result, dict):
                        conversation_id = final_result.get("conversationId") or final_result.get("conversation_id")
                    logger.debug(f"Got final result after {event_count} events, conversation_id={conversation_id}")
                    break
                elif "error" in event:
                    error = event["error"]
                    raise RuntimeError(f"MCP tool error: {error.get('message', error)}")

            # Process streaming events
            if event.get("method") == "codex/event":
                params = event.get("params", {})
                msg = params.get("msg", {})
                msg_type = msg.get("type")

                # Log ALL event types to help debug missing messages
                logger.debug(f"MCP event: type={msg_type}, keys={list(msg.keys())}")

                if msg_type == "session_configured":
                    session_id = msg.get("session_id")
                    logger.debug(f"Session configured: {session_id}")

                elif msg_type == "item_completed":
                    item = msg.get("item", {})
                    item_type = item.get("type")

                    # Log ALL item_completed events to help debug
                    logger.debug(f"item_completed: type={item_type}, keys={list(item.keys())}")

                    # Agent text responses - codex uses "AgentMessage" type
                    if item_type == "AgentMessage":
                        content = item.get("content", [])
                        for block in content:
                            if isinstance(block, dict) and block.get("text"):
                                agent_messages.append(block["text"])
                            elif isinstance(block, str):
                                agent_messages.append(block)

                    # Also check for "agent_message" (lowercase) variant
                    elif item_type == "agent_message":
                        text = item.get("text", "") or item.get("message", "")
                        if text:
                            agent_messages.append(text)
                        # Also check content array
                        content = item.get("content", [])
                        for block in content:
                            if isinstance(block, dict) and block.get("text"):
                                agent_messages.append(block["text"])
                            elif isinstance(block, str):
                                agent_messages.append(block)

                    # Legacy format check
                    elif item_type == "message" and item.get("role") == "assistant":
                        content = item.get("content", [])
                        for block in content:
                            if isinstance(block, dict) and block.get("text"):
                                agent_messages.append(block["text"])
                            elif isinstance(block, str):
                                agent_messages.append(block)

                    # Generic message type - check for text/content
                    elif item_type == "message":
                        text = item.get("text", "")
                        if text:
                            agent_messages.append(text)
                        content = item.get("content", [])
                        if isinstance(content, str):
                            agent_messages.append(content)
                        elif isinstance(content, list):
                            for block in content:
                                if isinstance(block, dict) and block.get("text"):
                                    agent_messages.append(block["text"])
                                elif isinstance(block, str):
                                    agent_messages.append(block)

                    # Function call outputs (for context)
                    elif item_type == "function_call_output":
                        output = item.get("output", "")
                        if output and len(output) < 1000:
                            agent_messages.append(f"[Tool output]: {output[:500]}")

                    # Log other item types we're not handling
                    elif item_type not in ("function_call", "tool_call", "UserMessage", "user_message"):
                        logger.debug(f"Unhandled item_completed type: {item_type}, item={item}")

                elif msg_type == "agent_message":
                    # Direct agent message event
                    message = msg.get("message", "")
                    if message:
                        agent_messages.append(message)

                elif msg_type in ("task_complete", "task_completed"):
                    # Task is done - capture last_agent_message as fallback
                    last_msg = msg.get("last_agent_message")
                    if last_msg and last_msg not in agent_messages:
                        agent_messages.append(last_msg)
                    logger.debug(f"Task complete after {event_count} events")
                    break

                elif msg_type == "token_count":
                    # Capture token usage for cost tracking
                    info = msg.get("info") or {}
                    if info:
                        usage = info.get("total_token_usage", {})
                        if usage:
                            token_usage = {
                                "input_tokens": usage.get("input_tokens", 0),
                                "output_tokens": usage.get("output_tokens", 0),
                                "cached_input_tokens": usage.get("cached_input_tokens", 0),
                                "reasoning_tokens": usage.get("reasoning_output_tokens", 0),
                                "total_tokens": usage.get("total_tokens", 0),
                            }
                            logger.debug(f"Token usage: {token_usage}")

                elif msg_type == "error":
                    error_msg = msg.get("error", msg.get("message", str(msg)))
                    raise RuntimeError(f"Codex error: {error_msg}")

                # Handle streaming text events (various formats)
                elif msg_type in ("text_delta", "content_block_delta", "message_delta"):
                    delta = msg.get("delta", {})
                    text = delta.get("text", "") or msg.get("text", "")
                    if text:
                        streaming_text.append(text)

                elif msg_type == "text":
                    text = msg.get("text", "")
                    if text:
                        streaming_text.append(text)

                elif msg_type == "response":
                    # Some versions send the full response this way
                    response_text = msg.get("response", "") or msg.get("text", "")
                    if response_text:
                        agent_messages.append(response_text)

                elif msg_type == "message":
                    # Direct message event
                    text = msg.get("text", "") or msg.get("content", "")
                    if text:
                        agent_messages.append(text)

                else:
                    # Log unknown event types at debug level to help diagnose
                    if msg_type and msg_type not in ("session_started", "thinking", "tool_call", "function_call"):
                        logger.debug(f"Unhandled MCP event type: {msg_type}, msg keys: {list(msg.keys())}")

        # Merge streaming text into messages if we got any
        if streaming_text:
            full_streaming = "".join(streaming_text)
            if full_streaming.strip():
                agent_messages.append(full_streaming)
                logger.debug(f"Captured {len(streaming_text)} streaming chunks ({len(full_streaming)} chars)")

        # Try to extract content from final_result if we have no messages
        if final_result and not agent_messages:
            if "content" in final_result:
                content = final_result["content"]
                if isinstance(content, list):
                    for block in content:
                        if isinstance(block, dict) and block.get("text"):
                            agent_messages.append(block["text"])
                        elif isinstance(block, str):
                            agent_messages.append(block)
                elif isinstance(content, str):
                    agent_messages.append(content)
            # Also check for text field
            if not agent_messages and "text" in final_result:
                agent_messages.append(final_result["text"])

        # Build result - prefer conversation_id from final result, fallback to session_id from events
        effective_conversation_id = conversation_id or session_id
        result = {
            "conversationId": effective_conversation_id,
            "messages": agent_messages,
            "output": "\n".join(agent_messages) if agent_messages else "",
            "usage": token_usage,  # Token usage for cost tracking
        }

        # Log detailed debug info if we didn't capture any messages
        if not agent_messages:
            event_types = [e.get("method") or f"id:{e.get('id')}" for e in all_events[:20]]
            logger.warning(
                f"MCP call returned no messages. "
                f"conversation_id={effective_conversation_id}, "
                f"session_id={session_id}, "
                f"event_count={len(all_events)}, "
                f"event_types={event_types}, "
                f"final_result_keys={list(final_result.keys()) if final_result else 'None'}"
            )
            # Log codex/event details for debugging
            codex_events = [e for e in all_events if e.get("method") == "codex/event"]
            if codex_events:
                for ce in codex_events[-5:]:  # Last 5 codex events
                    msg = ce.get("params", {}).get("msg", {})
                    logger.debug(f"  codex/event: type={msg.get('type')}, keys={list(msg.keys())}")

        logger.debug(f"MCP call complete: {len(agent_messages)} messages, conversation_id={effective_conversation_id}")
        return result

    def close(self) -> None:
        """Close the MCP connection gracefully."""
        if self._proc and self._proc.poll() is None:
            logger.info("Terminating MCP server...")
            self._proc.terminate()
            try:
                self._proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                logger.warning("MCP server didn't terminate, killing...")
                self._proc.kill()
                self._proc.wait()

        self._proc = None
        self._initialized = False

    @property
    def is_alive(self) -> bool:
        """Check if the MCP server is running."""
        return self._proc is not None and self._proc.poll() is None


@register_adapter("codex_mcp")
class CodexMCPAdapter(ExecutorAdapter):
    """
    Codex adapter using MCP server for sync conversations.

    This is the recommended way to have iterative conversations with Codex.
    The MCP client uses subprocess.Popen (not asyncio) so it persists across
    multiple asyncio.run() calls, preserving conversation state.
    """
    DEFAULT_MODEL = "gpt-5.1-codex-mini"  # Default codex model

    def __init__(self, model: str | None = None):
        self._model = model or self.DEFAULT_MODEL
        self._mcp_client: MCPClient | None = None
        self._sessions: dict[str, str] = {}  # session_id -> conversationId
        # Cumulative token usage for cost tracking
        self._total_usage: dict[str, int] = {
            "input_tokens": 0,
            "output_tokens": 0,
            "cached_input_tokens": 0,
            "reasoning_tokens": 0,
            "total_tokens": 0,
        }

    def _accumulate_usage(self, usage: dict[str, Any]) -> None:
        """Add usage to cumulative totals."""
        if not usage:
            return
        for key in self._total_usage:
            self._total_usage[key] += usage.get(key, 0)

    @property
    def total_usage(self) -> dict[str, int]:
        """Get cumulative token usage across all calls."""
        return self._total_usage.copy()

    def _ensure_client(self) -> MCPClient:
        """Ensure MCP client is running and return it."""
        if self._mcp_client is None:
            self._mcp_client = MCPClient()

        if not self._mcp_client.is_alive:
            self._mcp_client.start()

        return self._mcp_client

    @weave.op()
    def _call_codex(
        self,
        task: str,
        cwd: str,
        sandbox: str,
        model: str | None = None,
    ) -> dict[str, Any]:
        """
        Call codex MCP tool - traced by Weave.

        This is synchronous (uses subprocess.Popen, not asyncio) so the MCP
        server persists across calls.
        """
        client = self._ensure_client()

        args: dict[str, Any] = {
            "prompt": task,
            "cwd": cwd,
            "sandbox": sandbox,
        }
        if model:
            args["model"] = model

        logger.info(f"Calling codex with task_len={len(task)}, cwd={cwd}, model={model or 'default'}")

        result = client.call_tool("codex", args)

        # Log the result structure
        conversation_id = result.get("conversationId")
        messages_count = len(result.get("messages", []))
        output_len = len(result.get("output", ""))
        usage = result.get("usage", {})

        logger.info(
            f"codex result: conversation_id={conversation_id}, "
            f"messages_count={messages_count}, output_len={output_len}, "
            f"usage={usage.get('total_tokens', 0)} tokens"
        )

        # Warn if we got a conversation ID but no messages (agent did work but we lost output)
        if conversation_id and not messages_count and not output_len:
            logger.warning(
                f"codex returned conversation_id={conversation_id} but NO messages/output! "
                f"The agent processed {usage.get('total_tokens', 0)} tokens but we didn't capture the response. "
                f"This may indicate an issue with event parsing."
            )

        # Track usage
        self._accumulate_usage(usage)

        return {
            "conversation_id": conversation_id,
            "response": self._extract_response(result),
            "raw_messages": result.get("messages", []),
            "usage": usage,
            "total_usage": self.total_usage,
        }

    @weave.op()
    def _call_codex_reply(
        self,
        conversation_id: str,
        message: str,
    ) -> dict[str, Any]:
        """
        Call codex-reply MCP tool - traced by Weave.

        This is synchronous (uses subprocess.Popen, not asyncio) so the MCP
        server persists across calls.
        """
        client = self._ensure_client()

        logger.info(f"Calling codex-reply with conversation_id={conversation_id}, message_len={len(message)}")
        logger.debug(f"MCP client alive: {client.is_alive}, initialized: {client._initialized}")

        result = client.call_tool("codex-reply", {
            "conversationId": conversation_id,
            "prompt": message,
        })

        # Log the full result structure for debugging
        logger.info(
            f"codex-reply result: conversationId={result.get('conversationId')}, "
            f"messages_count={len(result.get('messages', []))}, "
            f"output_len={len(result.get('output', ''))}, "
            f"usage={result.get('usage', {}).get('total_tokens', 0)} tokens"
        )

        # Check for conversation loss - MCP returns empty result when session not found
        if not result.get("messages") and not result.get("output"):
            logger.error(
                f"codex-reply returned empty result for conversation_id={conversation_id}. "
                f"The MCP server may have lost the conversation state. Result: {result}"
            )

        # Track usage
        usage = result.get("usage", {})
        self._accumulate_usage(usage)

        response = self._extract_response(result)
        logger.debug(f"codex-reply response length: {len(response)} chars")

        return {
            "response": response,
            "raw_messages": result.get("messages", []),
            "usage": usage,
            "total_usage": self.total_usage,
            "conversation_lost": not result.get("messages") and not result.get("output"),
        }

    @weave.op()
    async def start_session(
        self,
        task: str,
        working_dir: Path,
        mode: Literal["sync", "async"] = "sync",
        model: str | None = None,
        sandbox: str = "workspace-write",
        **kwargs,
    ) -> ConversationSession:
        """Start a Codex session (sync or async mode)."""
        effective_model = model or self._model
        session = ConversationSession(
            adapter=self.name,
            mode=SessionMode(mode),
            working_dir=working_dir,
            task_description=task,
            model=effective_model,
        )

        if mode == "sync":
            # Use traced codex call (synchronous - MCP client persists across calls)
            result = self._call_codex(
                task=task,
                cwd=str(working_dir.absolute()),
                sandbox=sandbox,
                model=effective_model,
            )

            # Extract conversation ID and response
            session.conversation_id = result["conversation_id"]
            if session.conversation_id:
                self._sessions[session.id] = session.conversation_id
                logger.debug(f"Session {session.id[:8]} mapped to conversation {session.conversation_id}")
            else:
                # This is bad - we won't be able to continue this conversation
                logger.warning(
                    f"Session {session.id[:8]} started but MCP didn't return a conversation ID. "
                    "Further converse() calls will fail."
                )

            session.add_message("user", task)
            session.add_message("assistant", result["response"])

            # Track token usage on the session
            session.add_usage(result.get("usage", {}))

        else:
            # Async mode: use codex exec (fire-and-forget)
            # This runs in a subprocess without MCP, outputs JSONL events
            cmd = [
                "codex", "exec",
                "--dangerously-bypass-approvals-and-sandbox",
                "--skip-git-repo-check",
                "--json",
                "--model", effective_model,
                "-C", str(working_dir.absolute()),  # Explicit working directory
                "--", task,
            ]

            logger.info(f"Starting async codex: {' '.join(cmd[:8])}...")

            proc = subprocess.Popen(
                cmd,
                cwd=working_dir,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            session.process = proc
            session.add_message("user", task)

        return session

    async def send_message(
        self,
        session: ConversationSession,
        message: str,
    ) -> str:
        """Send a message to continue a sync conversation."""
        if session.mode != SessionMode.SYNC:
            raise ValueError("Cannot send message to async session")
        if session.status != SessionStatus.ACTIVE:
            raise ValueError(f"Session is not active: {session.status}")
        if not session.conversation_id:
            raise ValueError("Session has no conversation ID")

        # Use traced codex-reply call (synchronous - MCP client persists across calls)
        result = self._call_codex_reply(
            conversation_id=session.conversation_id,
            message=message,
        )

        response_text = result["response"]

        # Check if conversation was lost
        if result.get("conversation_lost"):
            logger.warning(
                f"Conversation {session.conversation_id} was lost. "
                f"Session {session.id} will be marked as needing re-delegation."
            )
            # Mark the session as having a lost conversation so orchestrator can handle it
            session.conversation_id = None  # Clear the stale ID

        session.add_message("user", message)
        session.add_message("assistant", response_text)

        # Track token usage on the session
        session.add_usage(result.get("usage", {}))

        return response_text

    @weave.op()
    def _parse_jsonl_output(self, stdout: str) -> dict[str, Any]:
        """
        Parse JSONL output from codex exec --json.

        Returns dict with:
        - response: The agent's message text
        - usage: Token usage stats
        - thread_id: The conversation thread ID
        - events: All parsed events (for debugging)
        """
        response_parts = []
        usage = {}
        thread_id = None
        events = []

        for line in stdout.strip().split("\n"):
            if not line.strip():
                continue
            try:
                event = json.loads(line)
                events.append(event)

                event_type = event.get("type", "")

                if event_type == "thread.started":
                    thread_id = event.get("thread_id")

                elif event_type == "item.completed":
                    item = event.get("item", {})
                    if item.get("type") == "agent_message":
                        response_parts.append(item.get("text", ""))

                elif event_type == "turn.completed":
                    usage = event.get("usage", {})

            except json.JSONDecodeError:
                logger.warning(f"Failed to parse JSONL line: {line[:100]}")
                continue

        return {
            "response": "\n".join(response_parts),
            "usage": usage,
            "thread_id": thread_id,
            "events": events,
        }

    @weave.op()
    async def check_status(
        self,
        session: ConversationSession,
    ) -> dict:
        """Check status of an async session."""
        if session.mode != SessionMode.ASYNC:
            return {"status": session.status.value}

        if session.process is None:
            return {"status": "unknown", "error": "No process handle"}

        # Check if process is still running
        poll = session.process.poll()
        if poll is None:
            return {"status": "running"}

        # Process finished - parse the JSONL output
        stdout, stderr = session.process.communicate()

        if poll == 0:
            # Parse JSONL to extract actual response
            parsed = self._parse_jsonl_output(stdout)
            response_text = parsed["response"] or "(no response captured)"

            # Add the response as a message
            session.add_message("assistant", response_text)

            # Track token usage
            if parsed["usage"]:
                session.add_usage({
                    "input_tokens": parsed["usage"].get("input_tokens", 0),
                    "output_tokens": parsed["usage"].get("output_tokens", 0),
                    "total_tokens": (
                        parsed["usage"].get("input_tokens", 0) +
                        parsed["usage"].get("output_tokens", 0)
                    ),
                })

            session.complete(response_text[:500])
            return {
                "status": "completed",
                "response": response_text,
                "usage": parsed["usage"],
                "thread_id": parsed["thread_id"],
            }
        else:
            # Try to parse stderr or stdout for error info
            error_msg = stderr.strip() if stderr else f"Exit code: {poll}"

            # Sometimes errors come through stdout as JSONL too
            if stdout and not stderr:
                try:
                    parsed = self._parse_jsonl_output(stdout)
                    if not parsed["response"]:
                        error_msg = f"Process failed with no response. Exit code: {poll}"
                except Exception:
                    error_msg = stdout[:500] if stdout else f"Exit code: {poll}"

            session.fail(error_msg[:500])
            return {"status": "failed", "error": error_msg, "exit_code": poll}

    async def stop(
        self,
        session: ConversationSession,
    ) -> None:
        """Stop a session."""
        import subprocess

        if session.process and session.process.poll() is None:
            session.process.terminate()
            try:
                session.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                session.process.kill()

        session.fail("Stopped by user")

        # Remove from tracking
        if session.id in self._sessions:
            del self._sessions[session.id]

    async def cleanup(self) -> None:
        """Clean up MCP server."""
        if self._mcp_client:
            self._mcp_client.close()
            self._mcp_client = None

    def _extract_response(self, result: dict) -> str:
        """Extract response text from MCP result."""
        # Check for error indicators - empty result suggests lost conversation
        if (
            result.get("conversationId") is None
            and not result.get("messages")
            and not result.get("output")
        ):
            logger.warning(f"MCP returned empty result - conversation may be lost: {result}")
            return "[ERROR] Conversation lost - the MCP server no longer has this session. Please re-delegate the task."

        # First check for our collected output
        if result.get("output"):
            return result["output"]

        # Check for messages list
        if result.get("messages"):
            return "\n".join(result["messages"])

        # Result may have different structures depending on codex version
        if "content" in result:
            content = result["content"]
            if isinstance(content, list):
                texts = []
                for block in content:
                    if isinstance(block, dict) and "text" in block:
                        texts.append(block["text"])
                    elif isinstance(block, str):
                        texts.append(block)
                if texts:
                    return "\n".join(texts)
            elif isinstance(content, str):
                return content

        if "text" in result:
            return result["text"]

        # Fallback: stringify the result (but log it as unexpected)
        logger.warning(f"Unexpected MCP result format, returning raw: {list(result.keys())}")
        return json.dumps(result, indent=2)
