"""
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 hashlib
import json
import logging
import queue
import subprocess
import threading
import time
from dataclasses import dataclass, field
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__)


# =============================================================================
# MessageCollector: Robust event collection with deduplication
# =============================================================================


@dataclass
class MessageSegment:
    """A segment within an assistant turn (for future segment-aware rendering)."""
    id: str
    kind: Literal["assistant_text", "progress", "tool_call", "tool_result", "error"]
    text: str
    status: Literal["open", "closed"] = "open"
    source_event_ids: set[str] = field(default_factory=set)


class MessageCollector:
    """
    Collects and deduplicates messages from MCP event stream.

    Solves the transcript rendering bugs by:
    1. Deduplicating events by ID
    2. Using priority-based message selection (item_completed > task_complete > streaming)
    3. Tracking message sources for debugging
    4. Never mixing streaming deltas with finalized messages

    Priority order (highest to lowest):
    - item_completed with AgentMessage/agent_message → DEFINITIVE
    - task_complete.last_agent_message → FALLBACK ONLY
    - streaming deltas → ONLY IF NO DEFINITIVE SOURCE
    """

    def __init__(self):
        # Deduplication
        self._seen_event_ids: set[str] = set()
        self._seen_content_hashes: set[str] = set()  # Content-based dedup

        # Message collection (priority-ordered)
        self._definitive_messages: list[str] = []  # From item_completed
        self._fallback_message: str | None = None  # From task_complete
        self._streaming_buffer: list[str] = []  # Streaming deltas

        # Metadata
        self._conversation_id: str | None = None
        self._session_id: str | None = None
        self._token_usage: dict[str, Any] = {}
        self._is_complete: bool = False

        # Debug tracking
        self._message_sources: list[tuple[str, str]] = []  # (source, text_preview)

    def _extract_event_id(self, event: dict) -> str | None:
        """Extract a unique event ID for deduplication."""
        # Try various ID fields that MCP events might have
        for key in ("id", "event_id", "item_id", "message_id"):
            if key in event:
                return str(event[key])

        # For nested events, try params
        params = event.get("params", {})
        msg = params.get("msg", {})
        for key in ("id", "event_id", "item_id"):
            if key in msg:
                return str(msg[key])

        return None

    def _content_hash(self, text: str) -> str:
        """Create a hash of content for deduplication."""
        # Normalize whitespace for comparison
        normalized = " ".join(text.split())
        return hashlib.md5(normalized.encode()).hexdigest()[:16]

    def _is_duplicate_content(self, text: str) -> bool:
        """Check if this content was already collected."""
        if not text or not text.strip():
            return True  # Empty is "duplicate" (skip it)

        content_hash = self._content_hash(text)
        if content_hash in self._seen_content_hashes:
            return True

        self._seen_content_hashes.add(content_hash)
        return False

    def _add_definitive_message(self, text: str, source: str) -> None:
        """Add a definitive message (from item_completed)."""
        if not text or not text.strip():
            return

        if self._is_duplicate_content(text):
            logger.debug(f"Skipping duplicate message from {source}: {text[:50]}...")
            return

        self._definitive_messages.append(text)
        self._message_sources.append((source, text[:50]))
        logger.debug(f"Added definitive message from {source}: {text[:50]}...")

    def _set_fallback_message(self, text: str, source: str) -> None:
        """Set fallback message (from task_complete). Only used if no definitive."""
        if not text or not text.strip():
            return

        # Only set if we don't have definitive messages
        if self._definitive_messages:
            logger.debug(f"Ignoring fallback from {source}: have definitive messages")
            return

        if self._is_duplicate_content(text):
            logger.debug(f"Skipping duplicate fallback from {source}")
            return

        self._fallback_message = text
        self._message_sources.append((source, text[:50]))

    def _add_streaming_delta(self, text: str) -> None:
        """Add streaming delta. Only used if no definitive messages at end."""
        if text:
            self._streaming_buffer.append(text)

    def process_event(self, event: dict) -> bool:
        """
        Process a single MCP event.

        Returns True if processing should continue, False if complete.
        """
        # 1. Check for event ID and dedupe
        event_id = self._extract_event_id(event)
        if event_id and event_id in self._seen_event_ids:
            logger.debug(f"Skipping duplicate event: {event_id}")
            return True
        if event_id:
            self._seen_event_ids.add(event_id)

        # 2. Handle codex/event notifications
        if event.get("method") == "codex/event":
            params = event.get("params", {})
            msg = params.get("msg", {})
            msg_type = msg.get("type")

            self._handle_codex_event(msg, msg_type)

            # Check for completion events
            if msg_type in ("task_complete", "task_completed"):
                self._is_complete = True
                return False

        return True

    def _handle_codex_event(self, msg: dict, msg_type: str | None) -> None:
        """Handle a codex/event notification."""
        if not msg_type:
            return

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

        # Item completed - DEFINITIVE SOURCE
        elif msg_type == "item_completed":
            self._handle_item_completed(msg)

        # Direct agent message - DEFINITIVE SOURCE
        elif msg_type == "agent_message":
            text = msg.get("message", "") or msg.get("text", "") or msg.get("content", "")
            self._add_definitive_message(text, "agent_message_event")

        # Task complete - FALLBACK SOURCE
        elif msg_type in ("task_complete", "task_completed"):
            last_msg = msg.get("last_agent_message")
            if last_msg:
                self._set_fallback_message(last_msg, "task_complete")

        # Token usage
        elif msg_type == "token_count":
            info = msg.get("info") or {}
            if info:
                usage = info.get("total_token_usage", {})
                if usage:
                    self._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),
                    }

        # Streaming deltas - LOWEST PRIORITY
        elif msg_type in ("text_delta", "content_block_delta", "message_delta", "text"):
            delta = msg.get("delta", {})
            text = delta.get("text", "") or msg.get("text", "")
            self._add_streaming_delta(text)

        # Response event - MEDIUM PRIORITY (treat as definitive)
        elif msg_type == "response":
            text = msg.get("response", "") or msg.get("text", "")
            self._add_definitive_message(text, "response_event")

        # Message event - check role
        elif msg_type == "message":
            role = msg.get("role", "").lower()
            if role in ("assistant", "agent", ""):
                text = msg.get("text", "") or msg.get("content", "")
                if text and role != "user":
                    self._add_definitive_message(text, "message_event")

        # Output event
        elif msg_type == "output":
            text = msg.get("output", "") or msg.get("text", "") or msg.get("content", "")
            self._add_definitive_message(text, "output_event")

        # Completion variants
        elif msg_type in ("item.completed", "response.completed"):
            item = msg.get("item", {})
            if item.get("type") == "agent_message":
                text = item.get("text", "")
                self._add_definitive_message(text, f"{msg_type}_event")
            elif "text" in msg:
                self._add_definitive_message(msg["text"], f"{msg_type}_direct")

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

    def _handle_item_completed(self, msg: dict) -> None:
        """Handle item_completed event - the primary source of messages."""
        item = msg.get("item", {})
        item_type = item.get("type")

        # AgentMessage - primary format
        if item_type == "AgentMessage":
            content = item.get("content", [])
            for block in content:
                if isinstance(block, dict) and block.get("text"):
                    self._add_definitive_message(block["text"], "AgentMessage")
                elif isinstance(block, str):
                    self._add_definitive_message(block, "AgentMessage_str")

        # agent_message - variant spelling
        elif item_type == "agent_message":
            text = item.get("text", "") or item.get("message", "")
            if text:
                self._add_definitive_message(text, "agent_message")
            content = item.get("content", [])
            for block in content:
                if isinstance(block, dict) and block.get("text"):
                    self._add_definitive_message(block["text"], "agent_message_content")
                elif isinstance(block, str):
                    self._add_definitive_message(block, "agent_message_content_str")

        # Generic message with assistant role
        elif item_type == "message":
            role = item.get("role", "")
            if role == "assistant":
                content = item.get("content", [])
                for block in content:
                    if isinstance(block, dict) and block.get("text"):
                        self._add_definitive_message(block["text"], "message_assistant")
                    elif isinstance(block, str):
                        self._add_definitive_message(block, "message_assistant_str")
            # Also check text field directly
            text = item.get("text", "")
            if text:
                self._add_definitive_message(text, "message_text")

        # Function call output (for context, truncated)
        elif item_type == "function_call_output":
            output = item.get("output", "")
            if output and len(output) < 1000:
                # Don't add to messages, just log
                logger.debug(f"Tool output: {output[:100]}...")

    def set_conversation_id(self, conv_id: str | None) -> None:
        """Set conversation ID from final result."""
        if conv_id:
            self._conversation_id = conv_id

    @property
    def conversation_id(self) -> str | None:
        """Get the conversation ID."""
        return self._conversation_id or self._session_id

    @property
    def token_usage(self) -> dict[str, Any]:
        """Get token usage stats."""
        return self._token_usage

    @property
    def is_complete(self) -> bool:
        """Check if collection is complete."""
        return self._is_complete

    def get_messages(self) -> list[str]:
        """
        Get the final deduplicated message list.

        Priority:
        1. Definitive messages (from item_completed)
        2. Fallback message (from task_complete)
        3. Streaming buffer (only if no definitive or fallback)
        """
        # Prefer definitive messages
        if self._definitive_messages:
            logger.debug(f"Returning {len(self._definitive_messages)} definitive messages")
            return self._definitive_messages

        # Fall back to task_complete message
        if self._fallback_message:
            logger.debug("Returning fallback message from task_complete")
            return [self._fallback_message]

        # Last resort: streaming buffer
        if self._streaming_buffer:
            full_text = "".join(self._streaming_buffer)
            if full_text.strip():
                logger.debug(f"Returning streaming buffer ({len(self._streaming_buffer)} chunks)")
                return [full_text]

        return []

    def get_response(self) -> str:
        """Get the final response as a single string."""
        messages = self.get_messages()
        return "\n".join(messages) if messages else ""

    def get_debug_info(self) -> dict:
        """Get debug information about message collection."""
        return {
            "seen_event_ids": len(self._seen_event_ids),
            "seen_content_hashes": len(self._seen_content_hashes),
            "definitive_messages": len(self._definitive_messages),
            "has_fallback": self._fallback_message is not None,
            "streaming_chunks": len(self._streaming_buffer),
            "message_sources": self._message_sources,
        }


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.
    """

    # Default config overrides for zwarm-managed codex sessions
    # These override ~/.codex/config.toml to ensure consistent behavior
    # Only used as fallback if no config_path is provided
    DEFAULT_CONFIG_OVERRIDES: dict[str, str] = {
        "model_reasoning_effort": "high",  # Use 'high' for compatibility with all models
    }

    def __init__(
        self,
        config_path: Path | None = None,
        config_overrides: dict[str, str] | None = None,
    ):
        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
        # Config path for full isolation (preferred)
        self._config_path = config_path
        # Fallback: merge default overrides with any custom ones (used if no config_path)
        self._config_overrides = {**self.DEFAULT_CONFIG_OVERRIDES, **(config_overrides or {})}

    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

            # Build command - prefer config file for full isolation, fallback to overrides
            cmd = ["codex", "mcp-server"]
            if self._config_path and self._config_path.exists():
                cmd.extend(["--config", str(self._config_path)])
                logger.info(f"Starting codex mcp-server with config: {self._config_path} (start_count={self._start_count})")
            else:
                # Fallback to individual overrides
                for key, value in self._config_overrides.items():
                    cmd.extend(["-c", f'{key}="{value}"'])
                logger.info(f"Starting codex mcp-server with overrides: {self._config_overrides} (start_count={self._start_count})")
            self._proc = subprocess.Popen(
                cmd,
                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.

        Uses MessageCollector for robust deduplication and priority-based
        message selection. This prevents the transcript rendering bugs:
        - Message duplication
        - Role contamination
        - Turn mis-association

        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")

        # Use MessageCollector for robust event handling
        collector = MessageCollector()
        final_result = None
        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
            line = self._read_line(timeout=30.0)

            if not line:
                # Timeout waiting for event - process is still alive, just slow
                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):
                        conv_id = final_result.get("conversationId") or final_result.get("conversation_id")
                        collector.set_conversation_id(conv_id)
                    logger.debug(f"Got final result after {event_count} events")
                    break
                elif "error" in event:
                    error = event["error"]
                    raise RuntimeError(f"MCP tool error: {error.get('message', error)}")

            # Process event through collector
            try:
                should_continue = collector.process_event(event)
                if not should_continue:
                    logger.debug(f"Collector signaled completion after {event_count} events")
                    break
            except RuntimeError as e:
                # Collector raises RuntimeError for codex errors
                raise

        # Try to extract content from final_result if collector has no messages
        messages = collector.get_messages()
        if final_result and not 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"):
                            messages.append(block["text"])
                        elif isinstance(block, str):
                            messages.append(block)
                elif isinstance(content, str):
                    messages.append(content)
            if not messages and "text" in final_result:
                messages.append(final_result["text"])

        # Build result
        result = {
            "conversationId": collector.conversation_id,
            "messages": messages,
            "output": "\n".join(messages) if messages else "",
            "usage": collector.token_usage,
        }

        # Log detailed debug info if we didn't capture any messages
        if not messages:
            debug_info = collector.get_debug_info()
            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={collector.conversation_id}, "
                f"event_count={len(all_events)}, "
                f"event_types={event_types}, "
                f"collector_debug={debug_info}, "
                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(messages)} messages, conversation_id={collector.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.

    Config isolation: Pass config_path to use a local codex.toml instead of
    the user's global ~/.codex/config.toml. This is the preferred approach.
    Falls back to config_overrides if no config_path is provided.
    """
    DEFAULT_MODEL = "gpt-5.1-codex-mini"  # Default codex model

    def __init__(
        self,
        model: str | None = None,
        config_path: Path | None = None,
        config_overrides: dict[str, str] | None = None,
    ):
        self._model = model or self.DEFAULT_MODEL
        self._config_path = config_path  # Path to local codex.toml for isolation
        self._config_overrides = config_overrides or {}
        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(
                config_path=self._config_path,
                config_overrides=self._config_overrides,
            )

        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,
        reasoning_effort: 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

        # Pass reasoning_effort to override codex config defaults
        # The config key is "model_reasoning_effort"
        if reasoning_effort:
            args["model_reasoning_effort"] = reasoning_effort

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

        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)

        # Filter out the sent message from the response using content hashing
        # The MCP may echo our prompt back, but we use robust content comparison
        raw_messages = result.get("messages", [])

        # Create hash of user message for comparison (normalized)
        def normalize_for_comparison(text: str) -> str:
            """Normalize text for comparison (lowercase, collapsed whitespace)."""
            return " ".join(text.lower().split())

        user_msg_normalized = normalize_for_comparison(message)
        user_msg_hash = hashlib.md5(user_msg_normalized.encode()).hexdigest()

        def is_user_message_echo(text: str) -> bool:
            """Check if text is just an echo of the user message."""
            if not text:
                return True  # Empty is effectively an echo (skip it)

            text_normalized = normalize_for_comparison(text)
            text_hash = hashlib.md5(text_normalized.encode()).hexdigest()

            # Exact match (case-insensitive, whitespace-normalized)
            if text_hash == user_msg_hash:
                return True

            # Check if text IS the user message (not just starts with it)
            # This avoids the bug where "Fix bug by X" gets filtered when user said "Fix bug"
            if text_normalized == user_msg_normalized:
                return True

            return False

        filtered_messages = [m for m in raw_messages if not is_user_message_echo(m)]

        # Build filtered result for extraction
        filtered_result = {
            **result,
            "messages": filtered_messages,
            "output": "\n".join(filtered_messages) if filtered_messages else result.get("output", ""),
        }

        response = self._extract_response(filtered_result)
        filtered_count = len(raw_messages) - len(filtered_messages)
        if filtered_count > 0:
            logger.debug(f"Filtered {filtered_count} user echo messages from response")
        logger.debug(f"codex-reply response length: {len(response)} chars")

        return {
            "response": response,
            "raw_messages": filtered_messages,  # Return filtered 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,
                reasoning_effort=kwargs.get("reasoning_effort"),
            )

            # 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)
