import json
import os
from collections.abc import AsyncIterator
from typing import Any, Dict, Iterator, List, Optional, Union

from openai import AsyncOpenAI, OpenAI
from openai.types.chat import ChatCompletion, ChatCompletionChunk
from openai.types.chat.chat_completion_chunk import ChoiceDelta

from gwenflow.llms.base import ChatBase
from gwenflow.logger import logger
from gwenflow.telemetry import tracer
from gwenflow.types import ItemHelpers, Message, ModelResponse, ToolCall, Usage
from gwenflow.utils import extract_json_str


class ChatOpenAI(ChatBase):

    model: str = "gpt-4o-mini"

    # model parameters
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    n: Optional[int] = None
    stop: Optional[Union[str, List[str]]] = None
    max_completion_tokens: Optional[int] = None
    max_tokens: Optional[int] = None
    presence_penalty: Optional[float] = None
    frequency_penalty: Optional[float] = None
    logit_bias: Optional[Dict[int, float]] = None
    response_format: Optional[Dict[str, Any]] = None
    seed: Optional[int] = None
    logprobs: Optional[bool] = None
    top_logprobs: Optional[int] = None

    # clients
    client: Optional[OpenAI] = None
    async_client: Optional[AsyncOpenAI] = None

    # client parameters
    api_key: Optional[str] = None
    organization: Optional[str] = None
    base_url: Optional[str] = None
    timeout: Optional[Union[float, int]] = None
    max_retries: Optional[int] = None

    def _get_client_params(self) -> Dict[str, Any]:

        api_key = self.api_key
        if api_key is None:
            api_key = os.environ.get("OPENAI_API_KEY")
        if api_key is None:
            raise ValueError(
                "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable"
            )

        organization = self.organization
        if organization is None:
            organization = os.environ.get('OPENAI_ORG_ID')

        client_params = {
            "api_key": api_key,
            "organization": organization,
            "base_url": self.base_url,
            "timeout": self.timeout,
            "max_retries": self.max_retries,
        }

        client_params = {k: v for k, v in client_params.items() if v is not None}

        return client_params

    @property
    def _model_params(self) -> Dict[str, Any]:

        model_params = {
            "temperature": self.temperature,
            "top_p": self.top_p,
            "n": self.n,
            "stop": self.stop,
            "max_tokens": self.max_tokens or self.max_completion_tokens,
            "presence_penalty": self.presence_penalty,
            "frequency_penalty": self.frequency_penalty,
            "logit_bias": self.logit_bias,
            "response_format": self.response_format,
            "seed": self.seed,
            "logprobs": self.logprobs,
            "top_logprobs": self.top_logprobs,
        }

        if self.tools and self.tool_type == "fncall":
            model_params["tools"] = [tool.to_openai() for tool in self.tools]
            model_params["tool_choice"] = self.tool_choice or "auto"

        model_params = {k: v for k, v in model_params.items() if v is not None}

        return model_params

    def get_client(self) -> OpenAI:
        if self.client:
            return self.client
        client_params = self._get_client_params()
        self.client = OpenAI(**client_params)
        return self.client

    def get_async_client(self) -> AsyncOpenAI:
        if self.async_client:
            return self.async_client
        client_params = self._get_client_params()
        self.async_client = AsyncOpenAI(**client_params)
        return self.async_client

    def _format_response(self, response: str, response_format: dict = None) -> Any:
        """Process the response."""

        if response_format.get("type") == "json_object":
            try:
                json_str = extract_json_str(response)
                text_response = json.loads(json_str)
                return text_response
            except:
                pass

        return response

    def _format_message(self, message: Message) -> Dict[str, Any]:
        """Format a message into the format expected by OpenAI."""
        return message.to_openai()

    def _get_openai_usage(self, completion: Union[ChatCompletion, ChatCompletionChunk]) -> Usage:
        if not completion.usage:
            return None
        return Usage(
            requests=1,
            input_tokens=completion.usage.prompt_tokens,
            output_tokens=completion.usage.completion_tokens,
            total_tokens=completion.usage.total_tokens,
        )

    def _parse_response(self, completion: ChatCompletion) -> ModelResponse:
        model_response = ModelResponse(
            role=completion.choices[0].message.role,
            content=completion.choices[0].message.content,
            finish_reason="stop",
            usage=self._get_openai_usage(completion),
        )

        if hasattr(completion.choices[0].message, 'reasoning_content'):
            model_response.reasoning_content = completion.choices[0].message.reasoning_content

        tool_calls = completion.choices[0].message.tool_calls
        if tool_calls is not None and len(tool_calls) > 0:
            try:
                model_response.tool_calls = [ToolCall(**t.model_dump()) for t in tool_calls]
            except Exception as e:
                logger.warning(f"Error processing tool calls: {e}")

        if self.response_format:
            model_response.parsed = self._format_response(completion.choices[0].message.content, response_format=self.response_format)

        return model_response

    @tracer.llm(name="LLM Invoke")
    def invoke(self, input: Union[str, List[Message], List[Dict[str, str]]]) -> ModelResponse:
        try:
            messages_for_model = ItemHelpers.input_to_message_list(input)
            response: ChatCompletion = self.get_client().chat.completions.create(
                model=self.model,
                messages=[self._format_message(m) for m in messages_for_model],
                **self._model_params,
            )
        except Exception as e:
            raise RuntimeError(f"Error in calling openai API: {e}")

        return self._parse_response(response)

    @tracer.llm(name="LLM Async Invoke")
    async def ainvoke(self, input: Union[str, List[Message], List[Dict[str, str]]]) -> ModelResponse:
        try:
            messages_for_model = ItemHelpers.input_to_message_list(input)
            response: ChatCompletion = await self.get_async_client().chat.completions.create(
                model=self.model,
                messages=[self._format_message(m) for m in messages_for_model],
                **self._model_params,
            )
        except Exception as e:
            raise RuntimeError(f"Error in calling openai API: {e}")

        return self._parse_response(response)

    @tracer.llm(name="LLM Stream")
    def stream(self, input: Union[str, List[Message], List[Dict[str, str]]]) -> Iterator[ModelResponse]:
        try:
            messages_for_model = ItemHelpers.input_to_message_list(input)
            completion = self.get_client().chat.completions.create(
                model=self.model,
                messages=[self._format_message(m) for m in messages_for_model],
                stream=True,
                stream_options={"include_usage": True},
                **self._model_params,
            )

            _full_tool_calls = []

            for chunk in completion:

                response = ModelResponse(role="assistant")

                if chunk.choices:

                    delta: ChoiceDelta = chunk.choices[0].delta

                    if hasattr(delta, 'content') and delta.content:
                        response.content = delta.content

                    if hasattr(delta, 'tool_calls') and delta.tool_calls:
                        for tool_call in delta.tool_calls or []:
                            if _full_tool_calls and (not tool_call.id or tool_call.id == _full_tool_calls[-1].id):
                                if tool_call.function.name:
                                    _full_tool_calls[-1].function.name += tool_call.function.name
                                if tool_call.function.arguments:
                                    _full_tool_calls[-1].function.arguments += tool_call.function.arguments
                            else:
                                _full_tool_calls.append(
                                    ToolCall(**tool_call.model_dump())
                                )

                    if _full_tool_calls:
                        response.tool_calls = _full_tool_calls

                    if response.content or response.reasoning_content:
                        yield response
                    elif len(response.tool_calls)>0 and chunk.choices[0].finish_reason == "tool_calls":
                        yield response

                else:
                    if hasattr(response, 'usage'):
                        response.usage = self._get_openai_usage(chunk)
                        yield response

        except Exception as e:
            raise RuntimeError(f"Error in calling openai API: {e}")

    @tracer.llm(name="LLM Async Astream")
    async def astream(self, input: Union[str, List[Message], List[Dict[str, str]]]) -> AsyncIterator[ModelResponse]:
        try:
            messages_for_model = ItemHelpers.input_to_message_list(input)
            completion = await self.get_async_client().chat.completions.create(
                model=self.model,
                messages=[self._format_message(m) for m in messages_for_model],
                stream=True,
                stream_options={"include_usage": True},
                **self._model_params,
            )

            _full_tool_calls = []

            async for chunk in completion:

                response = ModelResponse(role="assistant")

                if chunk.choices:

                    delta: ChoiceDelta = chunk.choices[0].delta

                    if hasattr(delta, 'content') and delta.content:
                        response.content = delta.content

                    if hasattr(delta, 'tool_calls') and delta.tool_calls:
                        for tool_call in delta.tool_calls or []:
                            if _full_tool_calls and (not tool_call.id or tool_call.id == _full_tool_calls[-1].id):
                                if tool_call.function.name:
                                    _full_tool_calls[-1].function.name += tool_call.function.name
                                if tool_call.function.arguments:
                                    _full_tool_calls[-1].function.arguments += tool_call.function.arguments
                            else:
                                _full_tool_calls.append(
                                    ToolCall(**tool_call.model_dump())
                                )

                    if _full_tool_calls:
                        response.tool_calls = _full_tool_calls

                    if response.content or response.reasoning_content:
                        yield response
                    elif len(response.tool_calls)>0 and chunk.choices[0].finish_reason == "tool_calls":
                        yield response

                else:
                    if hasattr(response, 'usage'):
                        response.usage = self._get_openai_usage(chunk)
                        yield response

        except Exception as e:
            raise RuntimeError(f"Error in calling openai API: {e}")