import logging
from collections.abc import AsyncIterator, Iterator, Sequence
from typing import Any

from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk, ChoiceDelta

from dhenara.ai.providers.openai.base import OpenAIClientBase
from dhenara.ai.providers.openai.chat_completions_api.message_converter import (
    OpenAIMessageConverterChatCompletions as OpenAIMessageConverter,
)
from dhenara.ai.types.genai import (
    AIModelCallResponseMetaData,
    ChatResponse,
    ChatResponseChoice,
    ChatResponseChoiceDelta,
    ChatResponseContentItem,
    ChatResponseContentItemDelta,
    ChatResponseReasoningContentItemDelta,
    ChatResponseTextContentItemDelta,
    ChatResponseUsage,
    StreamingChatResponse,
)
from dhenara.ai.types.genai.ai_model import AIModelAPIProviderEnum, AIModelProviderEnum
from dhenara.ai.types.genai.dhenara.request import MessageItem, Prompt, SystemInstruction
from dhenara.ai.types.shared.api import SSEErrorResponse

logger = logging.getLogger(__name__)


class OpenAIChatCompletions(OpenAIClientBase):
    """OpenAI-compatible Chat Completions client.

    Prefer `OpenAIResponses` for OpenAI's Responses API.
    """

    def get_api_call_params(
        self,
        prompt: str | dict | Prompt | None,
        context: Sequence[str | dict | Prompt] | None = None,
        instructions: dict[str, Any] | list[str | dict | SystemInstruction] | None = None,
        messages: Sequence[MessageItem] | None = None,
    ) -> dict[str, Any]:
        if not self._client:
            raise RuntimeError("Client not initialized. Use with 'async with' context manager")

        formatter = self.formatter
        if formatter is None:
            raise RuntimeError("Formatter not initialized")

        if self._input_validation_pending:
            raise ValueError("inputs must be validated with `self.validate_inputs()` before api calls")

        messages_list: list[dict[str, Any]] = []
        user = self.config.get_user()

        # Process instructions
        if instructions:
            if not isinstance(instructions, dict):
                raise ValueError(
                    "Invalid Instructions format. Instructions should be processed and passed in prompt format (dict)."
                )
            messages_list.append(instructions)

        # Add previous messages and current prompt
        if messages is not None:
            formatted_messages = formatter.format_messages(
                messages=messages,
                model_endpoint=self.model_endpoint,
            )
            messages_list.extend(formatted_messages)
        else:
            if context:
                context_dicts: list[dict[str, Any]] = []
                for item in context:
                    if not isinstance(item, dict):
                        raise ValueError(
                            "Invalid context format. Context should be processed and passed in prompt format (dict)."
                        )
                    context_dicts.append(item)
                messages_list.extend(context_dicts)
            if prompt is not None:
                if not isinstance(prompt, dict):
                    raise ValueError(
                        "Invalid prompt format. Prompt should be processed and passed in prompt format (dict)."
                    )
                messages_list.append(prompt)

        chat_args: dict[str, Any] = {
            "model": self.model_name_in_api_calls,
            "messages": messages_list,
            "stream": self.config.streaming,
        }

        if user:
            if self.model_endpoint.api.provider != AIModelAPIProviderEnum.MICROSOFT_AZURE_AI:
                chat_args["safety_identifier"] = user

        max_output_tokens, _max_reasoning_tokens = self.config.get_max_output_tokens(self.model_endpoint.ai_model)

        if self.model_endpoint.api.provider != AIModelAPIProviderEnum.MICROSOFT_AZURE_AI:
            # NOTE: With reasoning models, `max_output_tokens` is deprecated in favour of `max_completion_tokens`.
            chat_args["max_completion_tokens"] = max_output_tokens
        else:
            chat_args["max_tokens"] = max_output_tokens

        model_settings = self.model_endpoint.ai_model.get_settings()
        if self.config.reasoning and model_settings.supports_reasoning and self.config.reasoning_effort is not None:
            effort = self.config.reasoning_effort
            if effort == "minimal":
                effort = "low"
            elif effort == "max":
                effort = "high"
            chat_args["reasoning_effort"] = effort

        if self.config.streaming:
            if self.model_endpoint.api.provider != AIModelAPIProviderEnum.MICROSOFT_AZURE_AI:
                chat_args["stream_options"] = {"include_usage": True}

        if self.config.options:
            chat_args.update(self.config.options)

        # --- Tools ---
        if self.config.tools:
            chat_args["tools"] = formatter.format_tools(
                tools=self.config.tools,
                model_endpoint=self.model_endpoint,
            )

        if self.config.tool_choice:
            chat_args["tool_choice"] = formatter.format_tool_choice(
                tool_choice=self.config.tool_choice,
                model_endpoint=self.model_endpoint,
            )

        # --- Structured Output ---
        structured_output_cfg = self._get_structured_output_config()
        if structured_output_cfg is not None:
            chat_args["response_format"] = formatter.format_structured_output(
                structured_output=structured_output_cfg,
                model_endpoint=self.model_endpoint,
            )

        return {"chat_args": chat_args}

    def do_api_call_sync(
        self,
        api_call_params: dict,
    ) -> ChatCompletion:
        chat_args = api_call_params["chat_args"]
        client = self._client
        if client is None:
            raise RuntimeError("Client not initialized. Use with 'async with' context manager")
        if self.model_endpoint.api.provider != AIModelAPIProviderEnum.MICROSOFT_AZURE_AI:
            response = client.chat.completions.create(**chat_args)
        else:
            response = client.complete(**chat_args)
        return response

    async def do_api_call_async(
        self,
        api_call_params: dict,
    ) -> ChatCompletion:
        chat_args = api_call_params["chat_args"]
        client = self._client
        if client is None:
            raise RuntimeError("Client not initialized. Use with 'async with' context manager")
        if self.model_endpoint.api.provider != AIModelAPIProviderEnum.MICROSOFT_AZURE_AI:
            response = await client.chat.completions.create(**chat_args)
        else:
            response = await client.complete(**chat_args)
        return response

    def do_streaming_api_call_sync(
        self,
        api_call_params: dict,
    ) -> Iterator[ChatCompletionChunk]:
        chat_args = api_call_params["chat_args"]
        client = self._client
        if client is None:
            raise RuntimeError("Client not initialized. Use with 'async with' context manager")
        if self.model_endpoint.api.provider != AIModelAPIProviderEnum.MICROSOFT_AZURE_AI:
            stream = client.chat.completions.create(**chat_args)
        else:
            stream = client.complete(**chat_args)

        return stream

    async def do_streaming_api_call_async(
        self,
        api_call_params: dict,
    ) -> AsyncIterator[ChatCompletionChunk]:
        chat_args = api_call_params["chat_args"]
        client = self._client
        if client is None:
            raise RuntimeError("Client not initialized. Use with 'async with' context manager")
        if self.model_endpoint.api.provider != AIModelAPIProviderEnum.MICROSOFT_AZURE_AI:
            stream = await client.chat.completions.create(**chat_args)
        else:
            stream = await client.complete(**chat_args)

        return stream

    def parse_stream_chunk(
        self,
        chunk: object,
    ) -> StreamingChatResponse | SSEErrorResponse | list[StreamingChatResponse | SSEErrorResponse] | None:
        """Handle streaming response with progress tracking and final response."""
        if not isinstance(chunk, ChatCompletionChunk):
            raise TypeError(f"Unexpected OpenAI stream chunk type: {type(chunk)}")
        processed_chunks: list[StreamingChatResponse | SSEErrorResponse] = []

        sm = self.streaming_manager
        if sm is None:
            raise RuntimeError("Streaming manager not initialized")

        sm.provider_metadata = {}
        sm.persistant_choice_metadata_list = []

        if not sm.provider_metadata:
            sm.provider_metadata = {
                "id": chunk.id,
                "created": str(chunk.created),
                "object": chunk.object if hasattr(chunk, "object") else None,
                "system_fingerprint": chunk.system_fingerprint if hasattr(chunk, "system_fingerprint") else None,
            }

        # Process usage (some SDKs lag behind and may omit usage on early chunks)
        if hasattr(chunk, "usage") and chunk.usage:
            usage = ChatResponseUsage(
                total_tokens=chunk.usage.total_tokens,
                prompt_tokens=chunk.usage.prompt_tokens,
                completion_tokens=chunk.usage.completion_tokens,
            )
            sm.update_usage(usage)

        if chunk.choices:
            choice_deltas = []
            for choice in chunk.choices:
                if len(sm.persistant_choice_metadata_list) < choice.index + 1:
                    sm.persistant_choice_metadata_list.append(
                        {
                            "role": choice.delta.role,
                            "refusal": choice.delta.refusal if hasattr(choice.delta, "refusal") else None,
                        }
                    )

                role = choice.delta.role or sm.persistant_choice_metadata_list[choice.index]["role"]

                choice_deltas.append(
                    ChatResponseChoiceDelta(
                        index=choice.index,
                        finish_reason=choice.finish_reason if hasattr(choice, "finish_reason") else None,
                        stop_sequence=None,
                        content_deltas=[
                            self.process_content_item_delta(
                                index=0,
                                role=role,
                                delta=choice.delta,
                            ),
                        ],
                        metadata={},
                    )
                )

            response_chunk = sm.update(choice_deltas=choice_deltas)
            stream_response = StreamingChatResponse(
                id=chunk.id,
                data=response_chunk,
            )

            processed_chunks.append(stream_response)

        return processed_chunks

    def _get_usage_from_provider_response(
        self,
        response: ChatCompletion,
    ) -> ChatResponseUsage:
        usage = getattr(response, "usage", None)
        if usage is None:
            return ChatResponseUsage(total_tokens=0, prompt_tokens=0, completion_tokens=0)
        return ChatResponseUsage(
            total_tokens=getattr(usage, "total_tokens", 0) or 0,
            prompt_tokens=getattr(usage, "prompt_tokens", 0) or 0,
            completion_tokens=getattr(usage, "completion_tokens", 0) or 0,
        )

    def parse_response(
        self,
        response: object,
    ) -> ChatResponse:
        """Parse the provider response into our standard format."""
        if not isinstance(response, ChatCompletion):
            raise TypeError(f"Unexpected OpenAI chat completion response type: {type(response)}")
        usage, usage_charge = self.get_usage_and_charge(response)
        usage_chat = usage if isinstance(usage, ChatResponseUsage) else None

        choices = []
        for choice in response.choices:
            content_items = self.process_content_item(
                index=0,
                role=choice.message.role,
                content_item=choice.message,
            )

            contents = content_items if isinstance(content_items, list) else [content_items]

            choice_obj = ChatResponseChoice(
                index=choice.index,
                finish_reason=choice.finish_reason if hasattr(choice, "finish_reason") else None,
                stop_sequence=None,
                contents=contents,
                metadata={},
            )
            choices.append(choice_obj)

        return ChatResponse(
            model=response.model,
            provider=self.model_endpoint.ai_model.provider,
            api_provider=self.model_endpoint.api.provider,
            usage=usage_chat,
            usage_charge=usage_charge,
            choices=choices,
            metadata=AIModelCallResponseMetaData(
                streaming=False,
                duration_seconds=0,
                provider_metadata={
                    "id": response.id,
                    "created": str(response.created),
                    "object": response.object if hasattr(response, "object") else None,
                    "system_fingerprint": response.system_fingerprint
                    if hasattr(response, "system_fingerprint")
                    else None,
                },
            ),
        )

    def process_content_item(
        self,
        index: int,
        role: str,
        content_item: ChatCompletionMessage,
    ) -> ChatResponseContentItem | list[ChatResponseContentItem]:
        converted_items = OpenAIMessageConverter.provider_message_to_content_items(
            message=content_item,
            role=role,
            index_start=index,
            ai_model_provider=self.model_endpoint.ai_model.provider,
            structured_output_config=self._get_structured_output_config(),
        )

        if not converted_items:
            return self.get_unknown_content_type_item(
                index=index,
                role=role,
                unknown_item=content_item,
                streaming=False,
            )

        if len(converted_items) == 1:
            return converted_items[0]

        return converted_items

    def process_content_item_delta(
        self,
        index: int,
        role: str,
        delta: ChoiceDelta,
    ) -> ChatResponseContentItemDelta:
        if hasattr(delta, "content"):
            sm = self.streaming_manager
            if sm is None:
                raise RuntimeError("Streaming manager not initialized")
            if self.model_endpoint.ai_model.provider == AIModelProviderEnum.DEEPSEEK:
                content = delta.content or ""

                think_start = "<think>" in content
                think_end = "</think>" in content

                if think_start:
                    sm.progress.in_thinking_block = True
                    _, reasoning_part = content.split("<think>", 1)
                    return ChatResponseReasoningContentItemDelta(
                        index=index,
                        role=role,
                        text_delta=reasoning_part,
                    )

                elif think_end:
                    sm.progress.in_thinking_block = False
                    reasoning_part, answer_part = content.split("</think>", 1)
                    return (
                        ChatResponseReasoningContentItemDelta(
                            index=index,
                            role=role,
                            text_delta=reasoning_part,
                        )
                        if reasoning_part
                        else ChatResponseTextContentItemDelta(
                            index=index + 1,
                            role=role,
                            text_delta=answer_part,
                        )
                    )

                elif sm.progress.in_thinking_block:
                    return ChatResponseReasoningContentItemDelta(
                        index=index,
                        role=role,
                        text_delta=content,
                    )

                else:
                    return ChatResponseTextContentItemDelta(
                        index=index + 1,
                        role=role,
                        text_delta=content,
                    )
            else:
                return ChatResponseTextContentItemDelta(
                    index=index,
                    role=role,
                    text_delta=delta.content,
                )

        # NOTE: There is no way to identify content type for OpenAI streaming response.
        # Also, providers may send a few extra chunks with no content.
        # Eg: The very first chunk will only have the `role` set (and role in other chunks will be None)
        # Therefore, don't send a generic delta item here as it will mess up streaming manager updates.

        return self.get_unknown_content_type_item(
            index=index,
            role=role,
            unknown_item=delta,
            streaming=True,
        )
