from __future__ import annotations

from typing import TYPE_CHECKING, Any

from openai.types.responses import (
    ResponseReasoningSummaryTextDeltaEvent,
    ResponseTextDeltaEvent,
)

from fast_agent.core.logging.logger import get_logger
from fast_agent.event_progress import ProgressAction
from fast_agent.llm.provider.openai._stream_capture import (
    save_stream_chunk as _save_stream_chunk,
)
from fast_agent.llm.provider.openai.streaming_utils import finalize_stream_response
from fast_agent.llm.provider.openai.tool_notifications import OpenAIToolNotificationMixin
from fast_agent.llm.stream_types import StreamChunk

_logger = get_logger(__name__)

class ResponsesStreamingMixin(OpenAIToolNotificationMixin):
    if TYPE_CHECKING:
        from pathlib import Path

        from fast_agent.core.logging.logger import Logger

        logger: Logger
        name: str | None

        def _notify_stream_listeners(self, chunk: StreamChunk) -> None: ...

        def _notify_tool_stream_listeners(
            self, event_type: str, payload: dict[str, Any] | None = None
        ) -> None: ...

        def _update_streaming_progress(
            self, content: str, model: str, estimated_tokens: int
        ) -> int: ...

        def chat_turn(self) -> int: ...

    async def _process_stream(
        self, stream: Any, model: str, capture_filename: Path | None
    ) -> tuple[Any, list[str]]:
        estimated_tokens = 0
        reasoning_chars = 0
        reasoning_segments: list[str] = []
        tool_streams: dict[int, dict[str, Any]] = {}
        notified_tool_indices: set[int] = set()
        final_response: Any | None = None

        async for event in stream:
            _save_stream_chunk(capture_filename, event)
            if isinstance(event, ResponseReasoningSummaryTextDeltaEvent):
                if event.delta:
                    reasoning_segments.append(event.delta)
                    self._notify_stream_listeners(
                        StreamChunk(text=event.delta, is_reasoning=True)
                    )
                    reasoning_chars += len(event.delta)
                    await self._emit_streaming_progress(
                        model=f"{model} (summary)",
                        new_total=reasoning_chars,
                        type=ProgressAction.THINKING,
                    )
                continue

            if isinstance(event, ResponseTextDeltaEvent):
                if event.delta:
                    self._notify_stream_listeners(
                        StreamChunk(text=event.delta, is_reasoning=False)
                    )
                    estimated_tokens = self._update_streaming_progress(
                        event.delta, model, estimated_tokens
                    )
                    self._notify_tool_stream_listeners(
                        "text",
                        {
                            "chunk": event.delta,
                        },
                    )
                continue

            event_type = getattr(event, "type", None)
            if event_type in {"response.completed", "response.incomplete"}:
                final_response = getattr(event, "response", None) or final_response
                continue
            if event_type == "response.output_item.added":
                item = getattr(event, "item", None)
                if getattr(item, "type", None) == "function_call":
                    index = getattr(event, "output_index", None)
                    if index is None:
                        continue
                    tool_info = {
                        "tool_name": getattr(item, "name", None),
                        "tool_use_id": getattr(item, "call_id", None)
                        or getattr(item, "id", None),
                        "notified": False,
                    }
                    tool_streams[index] = tool_info
                    if tool_info["tool_name"] and tool_info["tool_use_id"]:
                        self._notify_tool_stream_listeners(
                            "start",
                            {
                                "tool_name": tool_info["tool_name"],
                                "tool_use_id": tool_info["tool_use_id"],
                                "index": index,
                            },
                        )
                        self.logger.info(
                            "Model started streaming tool call",
                            data={
                                "progress_action": ProgressAction.CALLING_TOOL,
                                "agent_name": self.name,
                                "model": model,
                                "tool_name": tool_info["tool_name"],
                                "tool_use_id": tool_info["tool_use_id"],
                                "tool_event": "start",
                            },
                        )
                        tool_info["notified"] = True
                        notified_tool_indices.add(index)
                continue

            if event_type == "response.function_call_arguments.delta":
                index = getattr(event, "output_index", None)
                if index is None:
                    continue
                tool_info = tool_streams.get(index, {})
                self._notify_tool_stream_listeners(
                    "delta",
                    {
                        "tool_name": tool_info.get("tool_name"),
                        "tool_use_id": tool_info.get("tool_use_id"),
                        "index": index,
                        "chunk": getattr(event, "delta", None),
                    },
                )
                continue

            if event_type == "response.output_item.done":
                item = getattr(event, "item", None)
                if getattr(item, "type", None) != "function_call":
                    continue
                index = getattr(event, "output_index", None)
                tool_info = tool_streams.pop(index, {}) if index is not None else {}
                tool_name = getattr(item, "name", None) or tool_info.get("tool_name")
                tool_use_id = (
                    getattr(item, "call_id", None)
                    or getattr(item, "id", None)
                    or tool_info.get("tool_use_id")
                )
                if index is None:
                    index = -1
                self._notify_tool_stream_listeners(
                    "stop",
                    {
                        "tool_name": tool_name,
                        "tool_use_id": tool_use_id,
                        "index": index,
                    },
                )
                self.logger.info(
                    "Model finished streaming tool call",
                    data={
                        "progress_action": ProgressAction.CALLING_TOOL,
                        "agent_name": self.name,
                        "model": model,
                        "tool_name": tool_name,
                        "tool_use_id": tool_use_id,
                        "tool_event": "stop",
                    },
                )
                if index >= 0:
                    notified_tool_indices.add(index)
                continue

        if tool_streams:
            incomplete_tools = [
                f"{info.get('tool_name', 'unknown')}:{info.get('tool_use_id', 'unknown')}"
                for info in tool_streams.values()
            ]
            self.logger.error(
                "Tool call streaming incomplete - started but never finished",
                data={
                    "incomplete_tools": incomplete_tools,
                    "tool_count": len(tool_streams),
                },
            )
            raise RuntimeError(
                "Streaming completed but tool call(s) never finished: "
                f"{', '.join(incomplete_tools)}"
            )

        if final_response is None:
            try:
                final_response = await stream.get_final_response()
            except Exception as exc:
                self.logger.warning("Failed to fetch final Responses payload", exc_info=exc)
                raise

        finalize_stream_response(
            final_response=final_response,
            model=model,
            agent_name=self.name,
            chat_turn=self.chat_turn,
            logger=self.logger,
            notified_tool_indices=notified_tool_indices,
            emit_tool_fallback=self._emit_tool_notification_fallback,
        )

        return final_response, reasoning_segments

    def _emit_tool_notification_fallback(
        self,
        output_items: list[Any],
        notified_indices: set[int],
        *,
        model: str,
    ) -> None:
        """Emit start/stop notifications when streaming metadata was missing."""
        if not output_items:
            return

        for index, item in enumerate(output_items):
            if index in notified_indices:
                continue
            if getattr(item, "type", None) != "function_call":
                continue

            tool_name = getattr(item, "name", None) or "tool"
            tool_use_id = (
                getattr(item, "call_id", None)
                or getattr(item, "id", None)
                or f"tool-{index}"
            )

            self._emit_fallback_tool_notification_event(
                tool_name=tool_name,
                tool_use_id=tool_use_id,
                index=index,
                model=model,
            )

    async def _emit_streaming_progress(
        self,
        model: str,
        new_total: int,
        type: ProgressAction = ProgressAction.STREAMING,
    ) -> None:
        """Emit a streaming progress event.

        Args:
            model: The model being used.
            new_total: The new total token count.
        """
        token_str = str(new_total).rjust(5)

        data = {
            "progress_action": type,
            "model": model,
            "agent_name": self.name,
            "chat_turn": self.chat_turn(),
            "details": token_str.strip(),
        }
        self.logger.info("Streaming progress", data=data)
