"""
OpenAI chat completion client implementation.
"""

import json
import logging
import time
from collections.abc import AsyncGenerator
from typing import Any, Dict, List, Optional

from openai import AsyncOpenAI

from agentbyte.messages import AssistantMessage, Message, ToolCallRequest

from agentbyte.component import Component

from .._retry_observability import build_retry_summary
from ..base import (
    BaseChatCompletionClient,
    BaseChatCompletionClientConfig,
    RateLimitError,
    AuthenticationError,
    InvalidRequestError,
)
from ..pricing import PricingRegistry
from ..retry_policy import RetryMixin, RetryPolicy
from ..types import (
    ChatCompletionChunk,
    ChatCompletionResult,
    ModelPricing,
    ModelClientError,
    Usage,
)

logger = logging.getLogger(__name__)


class OpenAIChatCompletionClientConfig(BaseChatCompletionClientConfig):
    """Configuration for OpenAIChatCompletionClient serialization."""

    retry_policy: RetryPolicy = RetryPolicy()
    key_env_var: str = "OPENAI_API_KEY"


class OpenAIChatCompletionClient(
    BaseChatCompletionClient,
    Component[OpenAIChatCompletionClientConfig],
    RetryMixin,
):
    """OpenAI chat completion client."""

    component_type = "model_client"
    component_config_schema = OpenAIChatCompletionClientConfig
    component_provider_override = "agentbyte.llm.OpenAIChatCompletionClient"

    def __init__(
        self,
        model: str,
        client: AsyncOpenAI,
        config: Optional[Dict[str, Any]] = None,
        model_pricing: Optional[ModelPricing] = None,
        max_retries: int = 3,
        initial_retry_delay: float = 1.0,
        max_retry_delay: float = 60.0,
        retry_policy: Optional[RetryPolicy] = None,
    ):
        super().__init__(model, client, config)
        self.model_pricing = model_pricing
        if retry_policy is not None:
            self.retry_policy = retry_policy
        else:
            self.retry_policy = RetryPolicy(
                max_retries=max_retries,
                initial_retry_delay=initial_retry_delay,
                max_retry_delay=max_retry_delay,
            )

    @property
    def max_retries(self) -> int:
        return self.retry_policy.max_retries

    @property
    def initial_retry_delay(self) -> float:
        return self.retry_policy.initial_retry_delay

    @property
    def max_retry_delay(self) -> float:
        return self.retry_policy.max_retry_delay

    @classmethod
    def from_api_key(
        cls,
        model: str,
        api_key: Optional[str] = None,
        org: Optional[str] = None,
        config: Optional[Dict[str, Any]] = None,
        model_pricing: Optional[ModelPricing] = None,
        max_retries: int = 3,
        initial_retry_delay: float = 1.0,
        max_retry_delay: float = 60.0,
        retry_policy: Optional[RetryPolicy] = None,
    ) -> "OpenAIChatCompletionClient":
        """Create a client from an explicit or env-backed API key."""
        from pydantic import ValidationError

        from .settings import OpenAISettings

        if api_key is None:
            try:
                settings = OpenAISettings()
                resolved_key = settings.api_key.get_secret_value()
                resolved_org = org or settings.org
            except ValidationError:
                raise ValueError(
                    "No OpenAI API key provided. "
                    "Pass api_key= explicitly or set the OPENAI_API_KEY environment variable."
                )
        else:
            resolved_key = api_key
            resolved_org = org

        openai_client = AsyncOpenAI(api_key=resolved_key, organization=resolved_org)
        return cls(
            model=model,
            client=openai_client,
            config=config,
            model_pricing=model_pricing,
            max_retries=max_retries,
            initial_retry_delay=initial_retry_delay,
            max_retry_delay=max_retry_delay,
            retry_policy=retry_policy,
        )

    async def create(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        """Make a single OpenAI API call."""
        retry_history: List = []
        try:
            result = await self._retry_with_backoff(
                self._create_internal,
                messages,
                tools,
                output_format,
                _retry_history=retry_history,
                **kwargs,
            )
            if retry_history:
                result.metadata["retry_observability"] = {
                    "history": retry_history,
                    "latest": build_retry_summary(retry_history),
                }
            return result
        except Exception as e:
            self._handle_error(e)

    async def _create_internal(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        """Internal create implementation (retried by wrapper)."""
        api_kwargs = {**self.config, **kwargs}
        api_messages = self._convert_messages_to_api_format(messages)

        structured_output = None
        if output_format:
            try:
                schema = output_format.model_json_schema()
                schema = self._make_schema_compatible(schema)
                api_kwargs["response_format"] = {
                    "type": "json_schema",
                    "json_schema": {
                        "name": schema.get("title", output_format.__name__),
                        "description": schema.get(
                            "description",
                            f"Structured output for {output_format.__name__}",
                        ),
                        "strict": True,
                        "schema": schema,
                    },
                }
            except Exception as e:
                logger.warning(
                    f"Failed to convert {output_format.__name__} to JSON schema: {e}. "
                    f"Continuing without structured output."
                )

        start_time = time.perf_counter()
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=api_messages,
            tools=tools,
            **api_kwargs,
        )
        duration_ms = int((time.perf_counter() - start_time) * 1000)

        choice = response.choices[0]
        message_content = choice.message.content or ""

        tool_calls = None
        if choice.message.tool_calls:
            tool_calls = [
                ToolCallRequest(
                    call_id=tc.id,
                    tool_name=tc.function.name,
                    parameters=json.loads(tc.function.arguments),
                )
                for tc in choice.message.tool_calls
            ]

        if output_format and message_content:
            try:
                structured_output = output_format.model_validate_json(message_content)
            except Exception as e:
                logger.warning(f"Failed to parse structured output: {e}. Using raw content.")
                structured_output = None

        return ChatCompletionResult(
            message=AssistantMessage(
                content=message_content,
                source=self.model,
                tool_calls=tool_calls,
            ),
            usage=Usage(
                duration_ms=duration_ms,
                llm_calls=1,
                tokens_input=response.usage.prompt_tokens,
                tokens_output=response.usage.completion_tokens,
                tool_calls=len(tool_calls) if tool_calls else 0,
                cost_estimate=self._estimate_cost(
                    response.usage.prompt_tokens,
                    response.usage.completion_tokens,
                ),
            ),
            model=response.model,
            metadata={"finish_reason": choice.finish_reason},
            finish_reason=choice.finish_reason,
            structured_output=structured_output,
        )

    async def create_stream(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        stream_options: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        """Make a streaming OpenAI API call."""
        try:
            async for chunk in self._retry_stream_with_backoff(
                self._create_stream_internal,
                messages,
                tools,
                output_format,
                stream_options,
                **kwargs,
            ):
                yield chunk
        except Exception as e:
            self._handle_error(e)

    async def _create_stream_internal(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        stream_options: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        """Internal streaming implementation (retried by wrapper)."""
        api_kwargs = {**self.config, **kwargs}
        if "stream_options" in api_kwargs and stream_options is None:
            stream_options = api_kwargs.pop("stream_options")
        else:
            api_kwargs.pop("stream_options", None)
        if stream_options is None:
            stream_options = {"include_usage": True}
        if stream_options:
            api_kwargs["stream_options"] = stream_options

        api_messages = self._convert_messages_to_api_format(messages)

        start_time = time.perf_counter()
        stream = await self.client.chat.completions.create(
            model=self.model,
            messages=api_messages,
            tools=tools,
            stream=True,
            **api_kwargs,
        )

        tool_call_chunks: Dict[Any, Dict[str, Any]] = {}

        async for chunk in stream:
            if hasattr(chunk, "usage") and chunk.usage and (not chunk.choices or len(chunk.choices) == 0):
                duration_ms = int((time.perf_counter() - start_time) * 1000)
                usage = Usage(
                    duration_ms=duration_ms,
                    llm_calls=1,
                    tokens_input=chunk.usage.prompt_tokens,
                    tokens_output=chunk.usage.completion_tokens,
                    cost_estimate=self._estimate_cost(
                        chunk.usage.prompt_tokens,
                        chunk.usage.completion_tokens,
                    ),
                )
                yield ChatCompletionChunk(
                    content="",
                    model=chunk.model,
                    is_complete=True,
                    tool_call_chunk=None,
                    usage=usage,
                    metadata={},
                )
                break

            if not chunk.choices:
                continue

            chunk_choice = chunk.choices[0]
            delta = chunk_choice.delta

            if delta.content:
                yield ChatCompletionChunk(
                    content=delta.content,
                    model=chunk.model,
                    is_complete=False,
                    tool_call_chunk=None,
                    usage=None,
                    metadata={"finish_reason": chunk_choice.finish_reason},
                )

            if delta.tool_calls:
                for tc_delta in delta.tool_calls:
                    index = getattr(tc_delta, "index", None)
                    call_id = tc_delta.id
                    tracking_key = index if index is not None else call_id

                    if tracking_key not in tool_call_chunks:
                        tool_call_chunks[tracking_key] = {
                            "id": call_id,
                            "function": {"name": "", "arguments": ""},
                        }

                    if call_id:
                        tool_call_chunks[tracking_key]["id"] = call_id

                    if tc_delta.function:
                        if tc_delta.function.name:
                            tool_call_chunks[tracking_key]["function"]["name"] = tc_delta.function.name
                        if tc_delta.function.arguments:
                            tool_call_chunks[tracking_key]["function"]["arguments"] += tc_delta.function.arguments

                    yield ChatCompletionChunk(
                        content="",
                        model=chunk.model,
                        is_complete=False,
                        tool_call_chunk=tool_call_chunks[tracking_key],
                        usage=None,
                        metadata={},
                    )

            if chunk_choice.finish_reason and not stream_options.get("include_usage"):
                yield ChatCompletionChunk(
                    content="",
                    model=chunk.model,
                    is_complete=True,
                    tool_call_chunk=None,
                    usage=None,
                    metadata={"finish_reason": chunk_choice.finish_reason},
                )

    def _make_schema_compatible(self, schema: Dict[str, Any]) -> Dict[str, Any]:
        """Make JSON schema compatible with OpenAI Structured Outputs."""
        compatible_schema: Dict[str, Any] = {}
        for key, value in schema.items():
            if isinstance(value, dict):
                compatible_schema[key] = self._make_schema_compatible(value)
            elif isinstance(value, list):
                compatible_schema[key] = [
                    self._make_schema_compatible(item) if isinstance(item, dict) else item
                    for item in value
                ]
            else:
                compatible_schema[key] = value

        if compatible_schema.get("type") == "object":
            compatible_schema["additionalProperties"] = False

        return compatible_schema

    def _estimate_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
        """Estimate cost from token usage."""
        if self.model_pricing is not None:
            pricing = self.model_pricing
        else:
            pricing = PricingRegistry.resolve_openai_chat_pricing(self.model)
        return pricing.calculate_cost(prompt_tokens, completion_tokens)

    def _to_config(self) -> OpenAIChatCompletionClientConfig:
        """Serialize to config without persisting secrets."""
        return OpenAIChatCompletionClientConfig(
            model=self.model,
            config=self.config.copy(),
            retry_policy=self.retry_policy,
            key_env_var=getattr(self, "_key_env_var", "OPENAI_API_KEY"),
        )

    @classmethod
    def _from_config(cls, config: OpenAIChatCompletionClientConfig) -> "OpenAIChatCompletionClient":
        """Reconstruct client from saved config."""
        import os

        api_key = os.environ.get(config.key_env_var)
        if not api_key:
            raise ValueError(
                f"Cannot reconstruct OpenAIChatCompletionClient: environment variable "
                f"'{config.key_env_var}' is not set."
            )
        instance = cls.from_api_key(
            model=config.model,
            api_key=api_key,
            config=config.config or None,
            retry_policy=config.retry_policy,
        )
        instance._key_env_var = config.key_env_var
        return instance

    def _handle_error(self, error: Exception) -> None:
        """Convert OpenAI exceptions to unified error hierarchy."""
        try:
            from openai import (
                RateLimitError as OpenAIRateLimitError,
                AuthenticationError as OpenAIAuthenticationError,
                BadRequestError,
                APIError,
            )
        except ImportError:
            raise ModelClientError(f"OpenAI error: {str(error)}") from error

        error_msg = str(error)
        error_type = type(error).__name__

        if isinstance(error, OpenAIRateLimitError):
            raise RateLimitError(f"OpenAI rate limit exceeded: {error_msg}") from error
        elif isinstance(error, OpenAIAuthenticationError):
            raise AuthenticationError(f"OpenAI authentication failed: {error_msg}") from error
        elif isinstance(error, BadRequestError):
            raise InvalidRequestError(f"OpenAI request invalid: {error_msg}") from error
        elif isinstance(error, APIError):
            raise ModelClientError(f"OpenAI API error ({error_type}): {error_msg}") from error
        else:
            raise ModelClientError(
                f"Unexpected error from OpenAI ({error_type}): {error_msg}"
            ) from error


__all__ = ["OpenAIChatCompletionClient", "OpenAIChatCompletionClientConfig"]
