# common/codex_chat_model.py
"""
CodexChatModel — Custom LangChain BaseChatModel for ChatGPT Backend API.

Codex models (included with ChatGPT Plus/Pro/Business/Enterprise)
are accessed via a proprietary endpoint at chatgpt.com, NOT api.openai.com.

Endpoint: POST https://chatgpt.com/backend-api/codex/responses
Headers:
  - Authorization: Bearer <access_token>
  - chatgpt-account-id: <account_id>
Body: OpenAI Responses API format

Reverse-engineered from Codex CLI source:
  codex-rs/core/src/model_provider_info.rs (line 148-152)
  codex-rs/core/tests/suite/client.rs (line 468-530)
"""
import asyncio
import json
import logging
import uuid
from typing import Any

import httpx
from langchain_core.callbacks import CallbackManagerForLLMRun, AsyncCallbackManagerForLLMRun
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import (
    AIMessage,
    AIMessageChunk,
    BaseMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult

logger = logging.getLogger(__name__)

CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex"

# Retry config for transient 5xx errors
_MAX_RETRIES = 3
_RETRY_BASE_DELAY = 2  # seconds (exponential: 2, 4, 8)


def _split_instructions_and_input(
    messages: list[BaseMessage],
) -> tuple[str, list[dict]]:
    """Split messages into instructions (system) and input items.

    The Codex backend API requires `instructions` as a top-level field
    (NOT inside the input array). SystemMessages are extracted and
    concatenated into the instructions string.

    IMPORTANT: function_call items must be TOP-LEVEL entries in the
    input array, NOT nested inside assistant message content.
    Supported content types inside messages are:
      input_text, input_image, output_text, refusal, input_file,
      computer_screenshot, summary_text.
    """
    instructions_parts: list[str] = []
    items: list[dict] = []
    for msg in messages:
        if isinstance(msg, SystemMessage):
            # SystemMessage → top-level instructions field
            instructions_parts.append(str(msg.content))
        elif isinstance(msg, HumanMessage):
            items.append({
                "role": "user",
                "content": [{"type": "input_text", "text": str(msg.content)}],
            })
        elif isinstance(msg, AIMessage):
            # Add text content as assistant message (output_text only)
            text_content = str(msg.content) if msg.content else ""
            if text_content:
                items.append({
                    "role": "assistant",
                    "content": [{"type": "output_text", "text": text_content}],
                })
            # Tool calls go as SEPARATE top-level items (NOT inside content)
            if hasattr(msg, "tool_calls") and msg.tool_calls:
                for tc in msg.tool_calls:
                    # Codex API needs TWO different IDs:
                    #   id: item ID, MUST start with "fc_"
                    #   call_id: call ID, starts with "call_" (used by function_call_output)
                    raw_id = tc.get("id", "") or f"fc_{uuid.uuid4().hex[:24]}"
                    # Ensure id starts with fc_
                    if raw_id.startswith("fc_"):
                        fc_id = raw_id
                        call_id = raw_id
                    elif raw_id.startswith("call_"):
                        fc_id = f"fc_{raw_id[5:]}"  # convert call_xxx → fc_xxx
                        call_id = raw_id
                    else:
                        fc_id = f"fc_{raw_id}"
                        call_id = raw_id
                    items.append({
                        "type": "function_call",
                        "id": fc_id,
                        "call_id": call_id,
                        "name": tc["name"],
                        "arguments": json.dumps(tc["args"]) if isinstance(tc["args"], dict) else tc["args"],
                    })
        elif isinstance(msg, ToolMessage):
            # call_id in function_call_output must match function_call's call_id
            tool_call_id = msg.tool_call_id or ""
            # If tool_call_id starts with fc_, convert to call_ format
            if tool_call_id.startswith("fc_"):
                tool_call_id = f"call_{tool_call_id[3:]}"
            items.append({
                "type": "function_call_output",
                "call_id": tool_call_id,
                "output": str(msg.content),
            })
        else:
            items.append({
                "role": "user",
                "content": [{"type": "input_text", "text": str(msg.content)}],
            })
    # Default instructions if none provided
    instructions = "\n\n".join(instructions_parts) if instructions_parts else "You are a helpful assistant."
    return instructions, items


def _parse_sse_line(line: str) -> dict | None:
    """Parse a Server-Sent Events line into data dict."""
    line = line.strip()
    if not line or line.startswith(":"):
        return None
    if line.startswith("data: "):
        data_str = line[6:]
        if data_str == "[DONE]":
            return None
        try:
            return json.loads(data_str)
        except json.JSONDecodeError:
            return None
    return None


def _convert_tool_to_function(tool: Any) -> dict:
    """Convert a LangChain tool definition to Responses API function format.

    Uses LangChain's built-in converter which handles edge cases like
    callable schemas, nested Pydantic models, etc.
    """
    from langchain_core.utils.function_calling import convert_to_openai_function

    # If already in the right format, return as-is
    if isinstance(tool, dict) and tool.get("type") == "function":
        return tool

    try:
        # Use LangChain's robust converter
        openai_fn = convert_to_openai_function(tool)
        return {
            "type": "function",
            "name": openai_fn.get("name", ""),
            "description": openai_fn.get("description", ""),
            "parameters": openai_fn.get("parameters", {}),
        }
    except Exception as e:
        logger.warning(f"Failed to convert tool {tool}: {e}")
        # Fallback: try dict-based conversion
        if isinstance(tool, dict):
            return {
                "type": "function",
                "name": tool.get("name", ""),
                "description": tool.get("description", ""),
                "parameters": tool.get("parameters", tool.get("schema", {})),
            }
        raise ValueError(f"Cannot convert tool of type {type(tool)}") from e


class CodexChatModel(BaseChatModel):
    """LangChain ChatModel that calls ChatGPT's backend API for Codex models.

    This targets https://chatgpt.com/backend-api/codex/responses
    which is the proprietary endpoint used by the Codex CLI.

    Streaming: _astream() yields AIMessageChunk for both text tokens
    AND tool_call_chunks. This enables LangGraph to fire
    on_chat_model_stream events needed for reasoning display.
    _agenerate() is kept as a non-streaming fallback.
    """

    model_name: str = "gpt-5.2-codex"
    access_token: str = ""
    account_id: str = ""
    refresh_token: str = ""  # For auto-refresh on 401
    user_id: str = ""  # For MongoDB token persistence
    base_url: str = CODEX_BASE_URL
    streaming: bool = False
    temperature: float | None = None
    max_tokens: int | None = None
    # Bound tools: stored as regular fields (not private) so Pydantic tracks them
    bound_tools: list[dict] = []
    bound_tool_choice: Any = None

    class Config:
        arbitrary_types_allowed = True

    @property
    def _llm_type(self) -> str:
        return "codex-chatgpt-backend"

    @property
    def _identifying_params(self) -> dict[str, Any]:
        return {
            "model_name": self.model_name,
            "base_url": self.base_url,
        }

    def bind_tools(
        self,
        tools: list[Any],
        *,
        tool_choice: Any | None = None,
        **kwargs: Any,
    ) -> "CodexChatModel":
        """Return a new model with tools bound for function calling.

        This is called by LangGraph agent framework to attach tool
        definitions to the model before making API calls.
        """
        converted_tools = [_convert_tool_to_function(t) for t in tools]
        logger.info(
            f"🔧 bind_tools: binding {len(converted_tools)} tools "
            f"(names: {[t.get('name', '?') for t in converted_tools]})"
        )
        # Create a copy of self with tools bound
        new_model = self.__class__(
            model_name=self.model_name,
            access_token=self.access_token,
            account_id=self.account_id,
            refresh_token=self.refresh_token,
            user_id=self.user_id,
            base_url=self.base_url,
            streaming=self.streaming,
            temperature=self.temperature,
            max_tokens=self.max_tokens,
            bound_tools=converted_tools,
            bound_tool_choice=tool_choice,
        )
        return new_model

    def _build_headers(self) -> dict[str, str]:
        headers = {
            "Authorization": f"Bearer {self.access_token}",
            "Content-Type": "application/json",
            "Accept": "text/event-stream",
        }
        if self.account_id:
            headers["chatgpt-account-id"] = self.account_id
        return headers

    async def _refresh_and_update_token(self) -> bool:
        """Refresh the access token using refresh_token. Returns True on success."""
        if not self.refresh_token:
            logger.warning("⚠️ No refresh_token available for Codex token refresh")
            return False

        logger.info("🔄 Codex token expired (401), refreshing...")
        try:
            async with httpx.AsyncClient(timeout=15) as client:
                resp = await client.post(
                    "https://auth.openai.com/oauth/token",
                    json={
                        "grant_type": "refresh_token",
                        "client_id": "app_EMoamEEZ73f0CkXaXp7hrann",
                        "refresh_token": self.refresh_token,
                    },
                    headers={"Content-Type": "application/json"},
                )
            if resp.status_code == 200:
                data = resp.json()
                new_access = data.get("access_token", "")
                new_refresh = data.get("refresh_token", "") or self.refresh_token
                if new_access:
                    self.access_token = new_access
                    self.refresh_token = new_refresh
                    # Extract new account_id from JWT
                    from common.codex_auth import _extract_account_id
                    new_account_id = _extract_account_id(new_access)
                    if new_account_id:
                        self.account_id = new_account_id
                    # Persist to MongoDB if user_id is available
                    if self.user_id:
                        try:
                            from common.codex_auth import save_codex_token_to_db
                            await save_codex_token_to_db(
                                self.user_id, new_access, new_refresh, self.account_id
                            )
                        except Exception as e:
                            logger.warning(f"⚠️ Failed to persist refreshed token: {e}")
                    # Also update local auth.json
                    try:
                        from common.codex_auth import _update_auth_json
                        _update_auth_json(new_access, new_refresh)
                    except Exception:
                        pass
                    logger.info("✅ Codex token refreshed successfully")
                    return True
            logger.error("❌ Codex token refresh failed: %d %s", resp.status_code, resp.text[:200])
        except Exception as e:
            logger.error("❌ Codex token refresh error: %s", e)
        return False

    def _build_request_body(
        self, messages: list[BaseMessage], **kwargs: Any
    ) -> dict:
        instructions, input_items = _split_instructions_and_input(messages)
        body: dict[str, Any] = {
            "model": self.model_name,
            "instructions": instructions,
            "input": input_items,
            "stream": True,  # Codex API requires stream=true always
            "store": False,
        }
        if self.temperature is not None:
            body["temperature"] = self.temperature
        if self.max_tokens:
            body["max_output_tokens"] = self.max_tokens
        # Include reasoning content (matches Codex CLI behavior)
        body["include"] = ["reasoning.encrypted_content"]
        # Add tools if bound
        if self.bound_tools:
            body["tools"] = self.bound_tools
            logger.info(f"🔧 Request body includes {len(self.bound_tools)} tools")
        if self.bound_tool_choice is not None:
            body["tool_choice"] = self.bound_tool_choice
        return body

    def _generate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: CallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Synchronous generation — internally streams SSE and collects full response."""
        url = f"{self.base_url}/responses"
        headers = self._build_headers()
        body = self._build_request_body(messages, **kwargs)

        # Debug: log request
        self._log_request_debug(body, "sync")

        parts: list[str] = []
        tool_calls: list[dict] = []
        event_types_seen: list[str] = []
        # Track function call metadata across SSE events
        pending_fn_calls: dict[int, dict] = {}  # output_index -> {name, call_id, id}
        with httpx.Client(timeout=120) as client:
            with client.stream("POST", url, headers=headers, json=body) as resp:
                if resp.status_code != 200:
                    resp.read()
                    logger.error(f"❌ Codex API error {resp.status_code}: {resp.text[:500]}")
                    raise httpx.HTTPStatusError(
                        f"HTTP {resp.status_code}", request=resp.request, response=resp
                    )
                for line in resp.iter_lines():
                    event = _parse_sse_line(line)
                    if not event:
                        continue
                    event_type = event.get("type", "")
                    event_types_seen.append(event_type)
                    text, tc = self._process_sse_event(event, event_type, pending_fn_calls)
                    if text:
                        parts.append(text)
                    if tc:
                        tool_calls.append(tc)

        content = "".join(parts)
        self._log_response_debug(content, tool_calls, event_types_seen)

        ai_msg = AIMessage(content=content)
        if tool_calls:
            ai_msg.tool_calls = tool_calls
        return ChatResult(generations=[ChatGeneration(message=ai_msg)])

    async def _agenerate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: AsyncCallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Async generation — internally streams SSE and collects full response.

        This is the PRIMARY method used by LangGraph's async execution.
        Handles both text responses and function/tool calls.
        """
        url = f"{self.base_url}/responses"
        headers = self._build_headers()
        body = self._build_request_body(messages, **kwargs)

        # Debug: log request
        self._log_request_debug(body, "async")

        parts: list[str] = []
        tool_calls: list[dict] = []
        event_types_seen: list[str] = []
        # Track function call metadata across SSE events
        pending_fn_calls: dict[int, dict] = {}  # output_index -> {name, call_id, id}
        async with httpx.AsyncClient(timeout=120) as client:
            async with client.stream("POST", url, headers=headers, json=body) as resp:
                if resp.status_code == 401:
                    await resp.aread()
                    logger.warning("⚠️ Codex API 401 — attempting token refresh...")
                    if await self._refresh_and_update_token():
                        headers = self._build_headers()
                        # Retry with refreshed token
                        async with client.stream("POST", url, headers=headers, json=body) as retry_resp:
                            if retry_resp.status_code != 200:
                                await retry_resp.aread()
                                raise httpx.HTTPStatusError(
                                    f"HTTP {retry_resp.status_code} after refresh",
                                    request=retry_resp.request, response=retry_resp,
                                )
                            async for line in retry_resp.aiter_lines():
                                event = _parse_sse_line(line)
                                if not event:
                                    continue
                                event_type = event.get("type", "")
                                text, tc = self._process_sse_event(event, event_type, pending_fn_calls)
                                if text:
                                    parts.append(text)
                                if tc:
                                    tool_calls.append(tc)
                        content = "".join(parts)
                        self._log_response_debug(content, tool_calls, event_types_seen)
                        ai_msg = AIMessage(content=content)
                        if tool_calls:
                            ai_msg.tool_calls = tool_calls
                        return ChatResult(generations=[ChatGeneration(message=ai_msg)])
                    else:
                        raise httpx.HTTPStatusError(
                            "HTTP 401: Token expired and refresh failed",
                            request=resp.request, response=resp,
                        )
                if resp.status_code != 200:
                    await resp.aread()
                    error_body = resp.text[:1000] if resp.text else 'N/A'
                    logger.error(f"❌ Codex API error {resp.status_code}: {error_body}")
                    logger.error(
                        f"❌ Request body that caused {resp.status_code}:\n"
                        f"{json.dumps(body, default=str, ensure_ascii=False)[:2000]}"
                    )
                    raise httpx.HTTPStatusError(
                        f"HTTP {resp.status_code}: {error_body[:200]}", request=resp.request, response=resp
                    )
                async for line in resp.aiter_lines():
                    event = _parse_sse_line(line)
                    if not event:
                        continue
                    event_type = event.get("type", "")
                    event_types_seen.append(event_type)
                    text, tc = self._process_sse_event(event, event_type, pending_fn_calls)
                    if text:
                        parts.append(text)
                    if tc:
                        tool_calls.append(tc)

        content = "".join(parts)
        self._log_response_debug(content, tool_calls, event_types_seen)

        ai_msg = AIMessage(content=content)
        if tool_calls:
            ai_msg.tool_calls = tool_calls
        return ChatResult(generations=[ChatGeneration(message=ai_msg)])

    # ── Async Streaming (_astream) ────────────────────────────────

    async def _astream(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: AsyncCallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ):
        """Async streaming — yields AIMessageChunk for real-time token display.

        This enables LangGraph to fire on_chat_model_stream events,
        which drives the reasoning/reflect display on the frontend.
        Yields both text tokens (as content) and tool calls (as tool_call_chunks).

        Includes retry logic for transient 5xx errors from the Codex API.
        """
        url = f"{self.base_url}/responses"
        headers = self._build_headers()
        body = self._build_request_body(messages, **kwargs)

        self._log_request_debug(body, "astream")

        last_error: Exception | None = None
        for attempt in range(_MAX_RETRIES):
            # Reset state for each retry attempt
            pending_fn_calls: dict[int, dict] = {}  # output_index → {name, call_id, id}
            fn_call_index_map: dict[int, int] = {}  # output_index → sequential tool_call index
            next_tc_index = 0
            event_types_seen: list[str] = []

            try:
                async with httpx.AsyncClient(timeout=120) as client:
                    async with client.stream("POST", url, headers=headers, json=body) as resp:
                        if resp.status_code >= 500:
                            await resp.aread()
                            error_body = resp.text[:1000] if resp.text else 'N/A'
                            logger.warning(
                                f"⚠️ Codex API {resp.status_code} (attempt {attempt + 1}/{_MAX_RETRIES}): "
                                f"{error_body[:200]}"
                            )
                            last_error = httpx.HTTPStatusError(
                                f"HTTP {resp.status_code}: {error_body[:200]}",
                                request=resp.request, response=resp,
                            )
                            if attempt < _MAX_RETRIES - 1:
                                delay = _RETRY_BASE_DELAY * (2 ** attempt)
                                logger.info(f"🔄 Retrying in {delay}s...")
                                await asyncio.sleep(delay)
                                continue
                            raise last_error

                        # ── 401 Token Expired → auto-refresh and retry ──
                        if resp.status_code == 401:
                            await resp.aread()
                            logger.warning("⚠️ Codex API 401 in _astream — attempting token refresh...")
                            if await self._refresh_and_update_token():
                                headers = self._build_headers()
                                # Don't count this as an attempt — continue will retry
                                continue
                            else:
                                raise httpx.HTTPStatusError(
                                    "HTTP 401: Token expired and refresh failed",
                                    request=resp.request, response=resp,
                                )

                        if resp.status_code != 200:
                            await resp.aread()
                            error_body = resp.text[:1000] if resp.text else 'N/A'
                            logger.error(f"❌ Codex API error {resp.status_code}: {error_body}")
                            raise httpx.HTTPStatusError(
                                f"HTTP {resp.status_code}: {error_body[:200]}",
                                request=resp.request, response=resp,
                            )

                        async for line in resp.aiter_lines():
                            event = _parse_sse_line(line)
                            if not event:
                                continue

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

                            # ── Text delta → yield content chunk ──
                            if event_type == "response.output_text.delta":
                                delta = event.get("delta", "")
                                if delta:
                                    yield ChatGenerationChunk(
                                        message=AIMessageChunk(content=delta)
                                    )

                            # ── Function call started → capture metadata + yield name chunk ──
                            elif event_type == "response.output_item.added":
                                item = event.get("item", {})
                                if item.get("type") == "function_call":
                                    output_index = event.get("output_index", 0)
                                    fn_name = item.get("name", "")
                                    fn_id = item.get("id", "") or f"fc_{uuid.uuid4().hex[:24]}"
                                    fn_call_id = item.get("call_id", "") or fn_id

                                    # Normalize call_id format
                                    if fn_call_id.startswith("fc_"):
                                        call_id = fn_call_id
                                    elif fn_call_id.startswith("call_"):
                                        call_id = fn_call_id
                                    else:
                                        call_id = f"call_{fn_call_id}"

                                    pending_fn_calls[output_index] = {
                                        "name": fn_name,
                                        "call_id": call_id,
                                        "id": fn_id,
                                    }
                                    tc_idx = next_tc_index
                                    fn_call_index_map[output_index] = tc_idx
                                    next_tc_index += 1

                                    logger.info(
                                        f"🔧 [stream] Function call started: {fn_name} "
                                        f"(id={fn_id}, call_id={call_id})"
                                    )

                                    yield ChatGenerationChunk(
                                        message=AIMessageChunk(
                                            content="",
                                            tool_call_chunks=[{
                                                "index": tc_idx,
                                                "id": call_id,
                                                "name": fn_name,
                                                "args": "",
                                            }],
                                        )
                                    )

                            # ── Function call args streaming → yield args delta ──
                            elif event_type == "response.function_call_arguments.delta":
                                output_index = event.get("output_index", 0)
                                tc_idx = fn_call_index_map.get(output_index, 0)
                                args_delta = event.get("delta", "")
                                if args_delta:
                                    yield ChatGenerationChunk(
                                        message=AIMessageChunk(
                                            content="",
                                            tool_call_chunks=[{
                                                "index": tc_idx,
                                                "id": None,
                                                "name": None,
                                                "args": args_delta,
                                            }],
                                        )
                                    )

                            # ── Function call done → log completion ──
                            elif event_type == "response.function_call_arguments.done":
                                output_index = event.get("output_index", 0)
                                fn_info = pending_fn_calls.pop(output_index, {})
                                name = fn_info.get("name") or event.get("name", "")
                                args_str = event.get("arguments", "{}")
                                logger.info(
                                    f"🔧 [stream] Function call done: {name}"
                                    f"({args_str[:100]})"
                                )

                # Success — break out of retry loop
                unique_types = list(dict.fromkeys(event_types_seen))
                logger.info(
                    f"📊 Codex _astream complete: event_types={unique_types}"
                )
                break

            except httpx.HTTPStatusError:
                # Already logged above — will retry or propagate
                if attempt >= _MAX_RETRIES - 1:
                    raise

    # ── Debug helpers ──────────────────────────────────────────────

    def _log_request_debug(self, body: dict, mode: str) -> None:
        """Log request body structure for debugging."""
        debug = {k: v for k, v in body.items() if k != "input"}
        debug["input_count"] = len(body.get("input", []))
        debug["tools_count"] = len(body.get("tools", []))
        if body.get("tools"):
            debug["tool_names"] = [t.get("name", "?") for t in body["tools"]]
        logger.info(
            f"📡 Codex API request ({mode}): "
            f"{json.dumps(debug, default=str, ensure_ascii=False)[:600]}"
        )

    def _log_response_debug(
        self, content: str, tool_calls: list[dict], event_types: list[str]
    ) -> None:
        """Log response summary for debugging."""
        unique_types = list(dict.fromkeys(event_types))
        logger.info(
            f"📊 Codex response: {len(content)} chars, "
            f"{len(tool_calls)} tool_calls, "
            f"event_types={unique_types}"
        )
        if tool_calls:
            logger.info(f"🔧 Tool calls: {[tc.get('name') for tc in tool_calls]}")
        elif not content:
            logger.warning("⚠️ Empty response — no text and no tool calls!")
        else:
            logger.info(f"💬 Text response preview: {content[:150]}...")

    # ── SSE parsing ──────────────────────────────────────────────

    @staticmethod
    def _process_sse_event(
        event: dict,
        event_type: str,
        pending_fn_calls: dict[int, dict],
    ) -> tuple[str | None, dict | None]:
        """Process a parsed SSE event, returning (text_delta, tool_call_or_none).

        Uses pending_fn_calls to track function call metadata across events:
        - response.output_item.added: captures name, call_id, id
        - response.function_call_arguments.done: uses captured metadata

        This is necessary because the Codex API sends function call name/id
        only in output_item.added, NOT in function_call_arguments.done.
        """
        # Text delta
        if event_type == "response.output_text.delta":
            return event.get("delta", ""), None

        # Function call started — capture metadata for later
        if event_type == "response.output_item.added":
            item = event.get("item", {})
            if item.get("type") == "function_call":
                output_index = event.get("output_index", 0)
                fn_name = item.get("name", "")
                # id starts with fc_ (item ID), call_id starts with call_ (call ID)
                fn_id = item.get("id", "") or f"fc_{uuid.uuid4().hex[:24]}"
                fn_call_id = item.get("call_id", "") or fn_id
                pending_fn_calls[output_index] = {
                    "name": fn_name,
                    "call_id": fn_call_id,
                    "id": fn_id,
                }
                logger.info(f"🔧 Function call started: {fn_name} (id={fn_id}, call_id={fn_call_id})")
            return None, None

        # Function call completed — use captured metadata
        if event_type == "response.function_call_arguments.done":
            output_index = event.get("output_index", 0)
            fn_info = pending_fn_calls.pop(output_index, {})

            # Use metadata from output_item.added, fallback to event fields
            name = fn_info.get("name") or event.get("name", "")
            # LangChain tool_call only has one 'id' field — use call_id 
            # (this is what ToolMessage.tool_call_id will match against)
            call_id = fn_info.get("call_id") or event.get("call_id", "") or f"call_{uuid.uuid4().hex[:24]}"
            args_str = event.get("arguments", "{}")

            logger.info(f"🔧 Function call done: {name}({args_str[:100]}) call_id={call_id}")
            try:
                args = json.loads(args_str)
            except json.JSONDecodeError:
                args = {}
            return None, {
                "id": call_id,
                "name": name,
                "args": args,
            }

        return None, None
