#
# Copyright (c) 2024–2025, Daily
#
# SPDX-License-Identifier: BSD 2-Clause License
#

"""OpenAI Responses API LLM service implementation.

This module provides OpenAIResponsesLLMService — a parallel LLM service that uses
the OpenAI Responses API (client.responses.create) instead of Chat Completions
(client.chat.completions.create).

Adapted from upstream pipecat-ai/pipecat PR #4074 to work with dv-pipecat's
Dict-based settings architecture.
"""

import asyncio
import json
from contextlib import asynccontextmanager
from typing import Any, Dict, List, Mapping, Optional

import httpx
from loguru import logger
from openai import NOT_GIVEN, APITimeoutError, AsyncOpenAI, AsyncStream, DefaultAsyncHttpxClient
from openai.types.responses import (
    ResponseCompletedEvent,
    ResponseFunctionCallArgumentsDeltaEvent,
    ResponseFunctionCallArgumentsDoneEvent,
    ResponseFunctionToolCall,
    ResponseOutputItemAddedEvent,
    ResponseOutputItemDoneEvent,
    ResponseStreamEvent,
    ResponseTextDeltaEvent,
)
from pydantic import BaseModel, Field

from pipecat.adapters.services.open_ai_responses_adapter import (
    OpenAIResponsesLLMAdapter,
    OpenAIResponsesLLMInvocationParams,
)
from pipecat.frames.frames import (
    Frame,
    LLMContextFrame,
    LLMFullResponseEndFrame,
    LLMFullResponseStartFrame,
    LLMTextFrame,
    LLMUpdateSettingsFrame,
    WarmupLLMFrame,
)
from pipecat.metrics.metrics import LLMTokenUsage
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.aggregators.llm_response import (
    LLMAssistantAggregatorParams,
    LLMUserAggregatorParams,
)
from pipecat.processors.aggregators.llm_response_universal import LLMContextAggregatorPair
from pipecat.processors.frame_processor import FrameDirection
from pipecat.services.llm_service import FunctionCallFromLLM, LLMService
from pipecat.utils.tracing.service_decorators import traced_llm


class OpenAIResponsesLLMService(LLMService):
    """OpenAI Responses API LLM service.

    This service works with the universal LLMContext and LLMContextAggregatorPair.
    It uses client.responses.create() instead of client.chat.completions.create().

    Example::

        llm = OpenAIResponsesLLMService(
            api_key=os.getenv("OPENAI_API_KEY"),
            model="gpt-4.1",
            system_instruction="You are a helpful assistant.",
        )
    """

    class InputParams(BaseModel):
        """Input parameters for OpenAI Responses API model configuration.

        Parameters:
            temperature: Sampling temperature (0.0 to 2.0).
            top_p: Top-p (nucleus) sampling parameter (0.0 to 1.0).
            max_completion_tokens: Maximum completion tokens to generate.
            frequency_penalty: Penalty for frequent tokens (-2.0 to 2.0).
            presence_penalty: Penalty for new tokens (-2.0 to 2.0).
            seed: Random seed for deterministic outputs.
            extra: Additional model-specific parameters (e.g., reasoning_effort).
        """

        temperature: Optional[float] = Field(default=None, ge=0.0, le=2.0)
        top_p: Optional[float] = Field(default=None, ge=0.0, le=1.0)
        max_completion_tokens: Optional[int] = Field(default=None, ge=1)
        frequency_penalty: Optional[float] = Field(default=None, ge=-2.0, le=2.0)
        presence_penalty: Optional[float] = Field(default=None, ge=-2.0, le=2.0)
        seed: Optional[int] = Field(default=None, ge=0)
        extra: Optional[Dict[str, Any]] = Field(default_factory=dict)

    adapter_class = OpenAIResponsesLLMAdapter

    def __init__(
        self,
        *,
        model: str = "gpt-4.1",
        api_key=None,
        base_url=None,
        organization=None,
        project=None,
        default_headers: Optional[Mapping[str, str]] = None,
        system_instruction: Optional[str] = None,
        service_tier: Optional[str] = None,
        params: Optional[InputParams] = None,
        retry_timeout_secs: Optional[float] = 5.0,
        retry_on_timeout: Optional[bool] = False,
        enable_warmup: bool = False,
        warmup_max_tokens: int = 16,
        **kwargs,
    ):
        """Initialize the OpenAI Responses API LLM service.

        Args:
            model: The OpenAI model name to use (e.g., "gpt-4.1", "gpt-4o").
            api_key: OpenAI API key. If None, uses environment variable.
            base_url: Custom base URL for OpenAI API. If None, uses default.
                For Azure: "https://{resource}.openai.azure.com/openai/v1/"
            organization: OpenAI organization ID.
            project: OpenAI project ID.
            default_headers: Additional HTTP headers to include in requests.
            system_instruction: System instruction/prompt for the model.
                Maps to the 'instructions' parameter in the Responses API.
            service_tier: Service tier to use (e.g., "auto", "flex", "priority").
            params: Input parameters for model configuration and behavior.
            retry_timeout_secs: Request timeout in seconds. Defaults to 5.0 seconds.
            retry_on_timeout: Whether to retry the request once if it times out.
            enable_warmup: Whether to enable LLM cache warmup. Defaults to False.
            warmup_max_tokens: Max tokens for warmup request. Responses API minimum
                is 16 (Chat Completions allows 10). Defaults to 16.
            **kwargs: Additional arguments passed to the parent LLMService.
        """
        super().__init__(**kwargs)

        params = params or OpenAIResponsesLLMService.InputParams()

        self._settings = {
            "temperature": params.temperature,
            "top_p": params.top_p,
            "max_completion_tokens": params.max_completion_tokens,
            "frequency_penalty": params.frequency_penalty,
            "presence_penalty": params.presence_penalty,
            "seed": params.seed,
            "extra": params.extra if isinstance(params.extra, dict) else {},
        }
        self._system_instruction = system_instruction
        self._service_tier = service_tier
        self._retry_timeout_secs = retry_timeout_secs
        self._retry_on_timeout = retry_on_timeout
        self._enable_warmup = enable_warmup
        self._warmup_max_tokens = warmup_max_tokens

        self._full_model_name = None
        self.set_model_name(model)
        self._client = self._create_client(
            api_key=api_key,
            base_url=base_url,
            organization=organization,
            project=project,
            default_headers=default_headers,
        )

        if self._system_instruction:
            logger.debug(f"{self}: Using system instruction: {self._system_instruction}")

    def _create_client(
        self,
        api_key=None,
        base_url=None,
        organization=None,
        project=None,
        default_headers=None,
    ) -> AsyncOpenAI:
        """Create an AsyncOpenAI client instance.

        Args:
            api_key: OpenAI API key.
            base_url: Custom base URL for the API.
            organization: OpenAI organization ID.
            project: OpenAI project ID.
            default_headers: Additional HTTP headers.

        Returns:
            Configured AsyncOpenAI client instance.
        """
        return AsyncOpenAI(
            api_key=api_key,
            base_url=base_url,
            organization=organization,
            project=project,
            http_client=DefaultAsyncHttpxClient(
                limits=httpx.Limits(
                    max_keepalive_connections=100,
                    max_connections=1000,
                    keepalive_expiry=None,
                )
            ),
            default_headers=default_headers,
        )

    def create_context_aggregator(
        self,
        context: LLMContext,
        *,
        user_params: LLMUserAggregatorParams = LLMUserAggregatorParams(),
        assistant_params: LLMAssistantAggregatorParams = LLMAssistantAggregatorParams(),
    ) -> LLMContextAggregatorPair:
        """Create a universal context aggregator pair for this service.

        Args:
            context: The LLM context to create an aggregator for.
            user_params: Parameters for user message aggregation.
            assistant_params: Parameters for assistant message aggregation.

        Returns:
            A LLMContextAggregatorPair instance.
        """
        return LLMContextAggregatorPair(
            context,
            user_params=user_params,
            assistant_params=assistant_params,
        )

    def can_generate_metrics(self) -> bool:
        """Check if this service can generate processing metrics."""
        return True

    async def process_frame(self, frame: Frame, direction: FrameDirection):
        """Process frames for LLM completion requests.

        Args:
            frame: The frame to process.
            direction: The direction of frame processing.
        """
        await super().process_frame(frame, direction)

        context = None
        if isinstance(frame, WarmupLLMFrame):
            asyncio.create_task(self._handle_warmup_frame(frame))
            return
        elif isinstance(frame, LLMContextFrame):
            context = frame.context
        elif isinstance(frame, LLMUpdateSettingsFrame):
            await self._update_settings(frame.settings)
        else:
            await self.push_frame(frame, direction)

        if context:
            try:
                await self.push_frame(LLMFullResponseStartFrame())
                await self.start_processing_metrics()
                await self._process_context(context)
            except httpx.TimeoutException as e:
                await self._call_event_handler("on_completion_timeout")
                await self.push_error(error_msg="LLM completion timeout", exception=e)
            except Exception as e:
                await self.push_error(error_msg=f"Error during completion: {e}", exception=e)
            finally:
                await self.stop_processing_metrics()
                await self.push_frame(LLMFullResponseEndFrame())

    @traced_llm
    async def _process_context(self, context: LLMContext):
        """Process an LLM context through the Responses API.

        Streams the response, handling text deltas and function calls.

        Args:
            context: The universal LLM context to process.
        """
        adapter: OpenAIResponsesLLMAdapter = self.get_llm_adapter()
        logger.debug(
            f"{self}: Generating response from universal context "
            f"{adapter.get_messages_for_logging(context)}"
        )

        invocation_params = adapter.get_llm_invocation_params(
            context, system_instruction=self._system_instruction
        )

        params = self._build_response_params(invocation_params)

        await self.start_ttfb_metrics()
        await self.start_connection_metrics()

        try:
            if self._retry_on_timeout:
                try:
                    stream = await asyncio.wait_for(
                        self._client.responses.create(**params),
                        timeout=self._retry_timeout_secs,
                    )
                    await self.stop_connection_metrics(success=True, connection_type="http")
                except (APITimeoutError, asyncio.TimeoutError):
                    logger.debug(f"{self}: Retrying responses.create due to timeout")
                    stream = await self._client.responses.create(**params)
                    await self.stop_connection_metrics(success=True, connection_type="http")
            else:
                stream = await self._client.responses.create(**params)
                await self.stop_connection_metrics(success=True, connection_type="http")
        except Exception as e:
            await self.stop_connection_metrics(success=False, error=str(e), connection_type="http")
            raise

        # Track function calls across stream events
        function_calls: Dict[str, Dict[str, str]] = {}  # item_id -> {name, call_id, arguments}
        current_arguments: Dict[str, str] = {}  # item_id -> accumulated arguments

        # Ensure stream and its async iterator are closed on cancellation/exception
        # to prevent socket leaks and uvloop crashes. Closing the iterator first
        # cascades cleanup through nested async generators (httpx/httpcore internals),
        # preventing uvloop's broken asyncgen finalizer from firing on Python 3.12+
        # (MagicStack/uvloop#699).
        @asynccontextmanager
        async def _closing(stream):
            chunk_iter = stream.__aiter__()
            try:
                yield chunk_iter
            finally:
                # Close the iterator first to cascade cleanup through
                # nested async generators (httpx/httpcore internals).
                if hasattr(chunk_iter, "aclose"):
                    await chunk_iter.aclose()
                # Then close the stream to release HTTP resources.
                if hasattr(stream, "close"):
                    await stream.close()
                elif hasattr(stream, "aclose"):
                    await stream.aclose()

        async with _closing(stream) as event_iter:
            async for event in event_iter:
                if isinstance(event, ResponseTextDeltaEvent):
                    await self.stop_ttfb_metrics()
                    await self.push_frame(LLMTextFrame(event.delta))

                elif isinstance(event, ResponseOutputItemAddedEvent):
                    await self.stop_ttfb_metrics()
                    item = event.item
                    if isinstance(item, ResponseFunctionToolCall):
                        item_id = item.id or ""
                        function_calls[item_id] = {
                            "name": item.name,
                            "call_id": item.call_id,
                            "arguments": "",
                        }
                        current_arguments[item_id] = ""

                elif isinstance(event, ResponseFunctionCallArgumentsDeltaEvent):
                    item_id = event.item_id
                    if item_id in current_arguments:
                        current_arguments[item_id] += event.delta

                elif isinstance(event, ResponseFunctionCallArgumentsDoneEvent):
                    item_id = event.item_id
                    if item_id in function_calls:
                        function_calls[item_id]["arguments"] = event.arguments

                elif isinstance(event, ResponseOutputItemDoneEvent):
                    item = event.item
                    if isinstance(item, ResponseFunctionToolCall):
                        item_id = item.id or ""
                        if item_id in function_calls:
                            function_calls[item_id]["name"] = item.name
                            function_calls[item_id]["call_id"] = item.call_id
                            function_calls[item_id]["arguments"] = item.arguments

                elif isinstance(event, ResponseCompletedEvent):
                    response = event.response
                    if response.usage:
                        cached = response.usage.input_tokens_details.cached_tokens if response.usage.input_tokens_details else 0
                        reasoning = response.usage.output_tokens_details.reasoning_tokens if response.usage.output_tokens_details else 0
                        tokens = LLMTokenUsage(
                            prompt_tokens=response.usage.input_tokens,
                            completion_tokens=response.usage.output_tokens,
                            total_tokens=response.usage.total_tokens,
                            cache_read_input_tokens=cached,
                            reasoning_tokens=reasoning,
                        )
                        await self.start_llm_usage_metrics(tokens)

                    # Used by @traced_llm for more detailed model name in tracing spans
                    self._full_model_name = response.model

        # Process any function calls
        if function_calls:
            fc_list: List[FunctionCallFromLLM] = []
            for item_id, fc in function_calls.items():
                try:
                    arguments = json.loads(fc["arguments"]) if fc["arguments"] else {}
                except json.JSONDecodeError:
                    logger.warning(
                        f"{self}: Failed to parse function call arguments: {fc['arguments']}"
                    )
                    arguments = {}
                fc_list.append(
                    FunctionCallFromLLM(
                        context=context,
                        tool_call_id=fc["call_id"],
                        function_name=fc["name"],
                        arguments=arguments,
                    )
                )
            await self.run_function_calls(fc_list)

    async def _update_settings(self, settings: Mapping[str, Any]):
        """Update model settings from an LLMUpdateSettingsFrame.

        Args:
            settings: Mapping of setting names to new values.
        """
        for key, value in settings.items():
            if key in self._settings:
                logger.info(f"{self}: Updating setting {key} to: [{value}]")
                self._settings[key] = value
            else:
                logger.warning(f"{self}: Unknown setting {key}, ignoring")

    def _build_response_params(
        self, invocation_params: OpenAIResponsesLLMInvocationParams
    ) -> dict:
        """Build parameters for the responses.create() call.

        Args:
            invocation_params: Parameters derived from the LLM context.

        Returns:
            Dictionary of parameters for the Responses API call.
        """
        params: Dict[str, Any] = {
            "model": self.model_name,
            "stream": True,
            "store": False,
            "input": invocation_params["input"],
        }

        # instructions (set by the adapter when input is non-empty)
        if "instructions" in invocation_params:
            params["instructions"] = invocation_params["instructions"]

        # Optional parameters — only include if explicitly given
        temperature = self._settings.get("temperature")
        if isinstance(temperature, (int, float)):
            params["temperature"] = temperature

        top_p = self._settings.get("top_p")
        if isinstance(top_p, (int, float)):
            params["top_p"] = top_p

        max_completion_tokens = self._settings.get("max_completion_tokens")
        if isinstance(max_completion_tokens, int):
            params["max_output_tokens"] = max_completion_tokens

        frequency_penalty = self._settings.get("frequency_penalty")
        if isinstance(frequency_penalty, (int, float)):
            params["frequency_penalty"] = frequency_penalty

        presence_penalty = self._settings.get("presence_penalty")
        if isinstance(presence_penalty, (int, float)):
            params["presence_penalty"] = presence_penalty

        seed = self._settings.get("seed")
        if isinstance(seed, int):
            params["seed"] = seed

        if self._service_tier is not None:
            params["service_tier"] = self._service_tier

        # Tools
        tools = invocation_params.get("tools")
        if tools is not None and not isinstance(tools, type(NOT_GIVEN)):
            params["tools"] = tools

        # Tool choice (e.g., "auto", "none", "required", or {"type": "function", "name": "..."})
        tool_choice = invocation_params.get("tool_choice")
        if tool_choice is not None and not isinstance(tool_choice, type(NOT_GIVEN)):
            params["tool_choice"] = tool_choice

        # Extra settings (e.g., reasoning_effort, parallel_tool_calls)
        extra = self._settings.get("extra", {})
        if extra:
            params.update(extra)

        return params

    async def run_inference(
        self,
        context: LLMContext,
        max_tokens: Optional[int] = None,
        system_instruction: Optional[str] = None,
    ) -> Optional[str]:
        """Run a one-shot, out-of-band inference with the given LLM context.

        Args:
            context: The LLM context containing conversation history.
            max_tokens: Optional maximum number of tokens to generate.
            system_instruction: Optional system instruction for this inference.

        Returns:
            The LLM's response as a string, or None if no response is generated.
        """
        adapter: OpenAIResponsesLLMAdapter = self.get_llm_adapter()
        effective_instruction = system_instruction or self._system_instruction
        invocation_params = adapter.get_llm_invocation_params(
            context, system_instruction=effective_instruction
        )

        params = self._build_response_params(invocation_params)

        # Override for non-streaming
        params["stream"] = False

        if max_tokens is not None:
            params["max_output_tokens"] = max_tokens

        response = await self._client.responses.create(**params)

        return response.output_text

    def _is_gpt5_model(self) -> bool:
        """Check if the current model is a GPT-5 series model."""
        model = (self.model_name or "").lower()
        return model.startswith("gpt-5")

    async def _handle_warmup_frame(self, frame: WarmupLLMFrame):
        """Handle WarmupLLMFrame to prime the LLM cache via the Responses API.

        Sends a minimal non-streaming request to warm up provider-side caches.
        The response is discarded and no frames are emitted.

        Args:
            frame: WarmupLLMFrame containing the messages to cache.
        """
        if not self._enable_warmup:
            self.logger.debug("LLM warmup is disabled, skipping")
            return

        try:
            # Convert warmup messages to Responses API input format
            adapter: OpenAIResponsesLLMAdapter = self.get_llm_adapter()
            # Create a minimal LLMContext from warmup messages
            warmup_context = LLMContext(messages=frame.messages)
            invocation_params = adapter.get_llm_invocation_params(
                warmup_context, system_instruction=self._system_instruction
            )

            warmup_params: Dict[str, Any] = {
                "model": self.model_name,
                "input": invocation_params["input"],
                "stream": False,
                "store": False,
                "max_output_tokens": self._warmup_max_tokens,
            }

            if "instructions" in invocation_params:
                warmup_params["instructions"] = invocation_params["instructions"]

            await self._client.responses.create(**warmup_params)

            self.logger.info("LLM cache warmed successfully via Responses API")

        except Exception as e:
            self.logger.error(f"Failed to warm LLM cache: {e}")
            # Don't propagate error — warmup failure shouldn't break the bot


__all__ = ["OpenAIResponsesLLMService"]
