import os
import json
import base64
from functools import wraps
from typing import AsyncIterator, Iterator, Callable, ParamSpec, Awaitable

from pydantic import BaseModel
from openai import OpenAI, AsyncOpenAI, Stream, AsyncStream, APIError as OpenAIAPIError
from openai.types.responses import ResponseStreamEvent

from promptbuilder.llm_client.base_client import BaseLLMClient, BaseLLMClientAsync, ResultType
from promptbuilder.llm_client.types import Response, Content, Candidate, UsageMetadata, Part, ThinkingConfig, Tool, ToolConfig, FunctionCall, MessageDict, Model
from promptbuilder.llm_client.config import DecoratorConfigs
from promptbuilder.llm_client.utils import inherited_decorator
from promptbuilder.llm_client.exceptions import APIError


P = ParamSpec("P")


@inherited_decorator
def _error_handler(func: Callable[P, Response]) -> Callable[P, Response]:
    """
    Decorator to catch error from openai and transform it into unified one
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except OpenAIAPIError as e:
            code = getattr(e, "status_code", None) or e.code
            response = getattr(e, "response", None)
            status = getattr(response, "reason_phrase", None)
            response_json = {
                "status": status,
                "message": e.message,
            }
            raise APIError(code, response_json, response)
    return wrapper

class OpenaiStreamIterator:
    def __init__(self, openai_iterator: Stream[ResponseStreamEvent]):
        self._openai_iterator = openai_iterator

    def __iter__(self) -> Iterator[Response]:
        output_tokens: int | None = None
        input_tokens: int | None = None
        total_tokens: int | None = None
        for next_event in self._openai_iterator:
            if next_event.type == "response.output_text.delta":
                parts = [Part(text=next_event.delta)]
                yield Response(candidates=[Candidate(content=Content(parts=parts, role="model"))])
            elif next_event.type == "response.completed":
                output_tokens = next_event.response.usage.output_tokens
                input_tokens = next_event.response.usage.input_tokens
                total_tokens = next_event.response.usage.total_tokens
        
        usage_metadata = UsageMetadata(
            candidates_token_count=output_tokens,
            prompt_token_count=input_tokens,
            total_token_count=total_tokens,
        )
        yield Response(candidates=[Candidate(content=Content(parts=[Part(text="")], role="model"))], usage_metadata=usage_metadata)


class OpenaiLLMClient(BaseLLMClient):
    PROVIDER = "openai"

    def __init__(
        self,
        model: str,
        api_key: str | None = None,
        decorator_configs: DecoratorConfigs | None = None,
        default_thinking_config: ThinkingConfig | None = None,
        default_max_tokens: int | None = None,
        **kwargs,
    ):
        if api_key is None:
            api_key = os.getenv("OPENAI_API_KEY")
        if api_key is None or not isinstance(api_key, str):
            raise ValueError("To create an openai llm client you need to either set the environment variable OPENAI_API_KEY or pass the api_key in string format")
        super().__init__(OpenaiLLMClient.PROVIDER, model, decorator_configs=decorator_configs, default_thinking_config=default_thinking_config, default_max_tokens=default_max_tokens)
        self._api_key = api_key
        self.client = OpenAI(api_key=api_key)
    
    @property
    def api_key(self) -> str:
        return self._api_key
    
    @staticmethod
    def _content_to_openai_messages(messages: list[Content], system_message: str | None = None) -> list[MessageDict]:
        openai_messages: list[MessageDict] = []
        if system_message is not None:
            openai_messages.append({"role": "developer", "content": system_message})
        for message in messages:
            role = "user" if message.role == "user" else "assistant"
            if message.parts is None:
                openai_messages.append({"role": role, "content": message.as_str()})
            else:
                content = []
                for part in message.parts:
                    if part.inline_data is not None and part.inline_data.data is not None:
                        base64_file_content = base64.b64encode(part.inline_data.data).decode('utf-8')
                        file_data = {
                            "file_data": f"data:{part.inline_data.mime_type};base64,{base64_file_content}"
                        }
                        if part.inline_data.display_name is not None:
                            file_data["filename"] = part.inline_data.display_name
                        match part.inline_data.mime_type:
                            case "application/pdf":
                                file_data["type"] = "input_file"
                            case "image/png" | "image/jpeg" | "image/webp":
                                file_data["type"] = "input_image"
                            case _:
                                raise ValueError(f"Unsupported inline data mime type: {part.inline_data.mime_type}. Supported types are: application/pdf, image/png, image/jpeg, image/webp.")
                        content.append(file_data)
                    else:
                        text_type = "input_text" if message.role == "user" else "output_text"
                        content.append({"type": text_type, "text": part.as_str()})
                openai_messages.append({"role": role, "content": content})
        return openai_messages

    @staticmethod
    def _process_thinking_config(thinking_config: ThinkingConfig | None) -> dict[str, str]:
        if thinking_config is None:
            return {}
        
        openai_thinking_config = {}
        if thinking_config.include_thoughts:
            # openai_thinking_config["summary"] = "auto"
            match thinking_config.thinking_budget:
                case 0 | None:
                    openai_thinking_config["reasoning"] = {"effort": "high"} # default
                case 1:
                    openai_thinking_config["reasoning"] = {"effort": "low"}
                case 2:
                    openai_thinking_config["reasoning"] = {"effort": "medium"}
                case 3:
                    openai_thinking_config["reasoning"] = {"effort": "high"}
                case _:
                    openai_thinking_config["reasoning"] = {"effort": "high"}
        return openai_thinking_config

    @_error_handler
    def _create(
        self,
        messages: list[Content],
        result_type: ResultType = None,
        *,
        thinking_config: ThinkingConfig | None = None,
        system_message: str | None = None,
        max_tokens: int | None = None,
        tools: list[Tool] | None = None,
        tool_config: ToolConfig = ToolConfig(),
    ) -> Response:
        openai_messages: list[MessageDict] = OpenaiLLMClient._content_to_openai_messages(messages, system_message)
        
        if max_tokens is None:
            max_tokens = self.default_max_tokens
        
        openai_kwargs = {
            "model": self.model,
            "max_output_tokens": max_tokens,
            "input": openai_messages,
        }
        
        if thinking_config is None:
            thinking_config = self.default_thinking_config
        openai_kwargs.update(OpenaiLLMClient._process_thinking_config(thinking_config))

        if tools is not None:
            openai_kwargs["parallel_tool_calls"] = True
            
            openai_tools = []
            allowed_function_names = None
            if tool_config.function_calling_config is not None:
                allowed_function_names = tool_config.function_calling_config.allowed_function_names
            for tool in tools:
                for func_decl in tool.function_declarations:
                    if allowed_function_names is None or func_decl.name in allowed_function_names:
                        parameters = func_decl.parameters
                        if parameters is not None:
                            parameters = parameters.model_dump(exclude_none=True)
                            parameters["additionalProperties"] = False
                        else:
                            parameters = {"type": "object", "properties": {}, "required": [], "additionalProperties": False}
                        openai_tools.append({
                            "type": "function",
                            "name": func_decl.name,
                            "description": func_decl.description,
                            "strict": True,
                            "parameters": parameters,
                        })
            openai_kwargs["tools"] = openai_tools
            
            tool_choice_mode = "AUTO"
            if tool_config.function_calling_config is not None:
                if tool_config.function_calling_config.mode is not None:
                    tool_choice_mode = tool_config.function_calling_config.mode
            if tool_choice_mode == "NONE":
                openai_kwargs["tool_choice"] = "none"
            elif tool_choice_mode == "AUTO":
                openai_kwargs["tool_choice"] = "auto"
            elif tool_choice_mode == "ANY":
                openai_kwargs["tool_choice"] = "required"
        
        if result_type is None or result_type == "json":
            response = self.client.responses.create(**openai_kwargs)
            
            parts: list[Part] = []
            for output_item in response.output:
                if output_item.type == "message":
                    for content in output_item.content:
                        parts.append(Part(text=content.text))
                elif output_item.type == "reasoning":
                    for summary in output_item.summary:
                        parts.append(Part(text=summary.text, thought=True))
                elif output_item.type == "function_call":
                    parts.append(Part(function_call=FunctionCall(args=json.loads(output_item.arguments), name=output_item.name)))
            
            return Response(
                candidates=[Candidate(content=Content(parts=parts, role="model"))],
                usage_metadata=UsageMetadata(
                    candidates_token_count=response.usage.output_tokens,
                    prompt_token_count=response.usage.input_tokens,
                    total_token_count=response.usage.total_tokens,
                ),
            )
        elif isinstance(result_type, type(BaseModel)):
            response = self.client.responses.parse(**openai_kwargs, text_format=result_type)
            
            parts: list[Part] = []
            for output_item in response.output:
                if output_item.type == "message":
                    for content in output_item.content:
                        parts.append(Part(text=content.text))
                elif output_item.type == "reasoning":
                    for summary in output_item.summary:
                        parts.append(Part(text=summary.text, thought=True))
                elif output_item.type == "function_call":
                    parts.append(Part(function_call=FunctionCall(args=json.loads(output_item.arguments), name=output_item.name)))
            parsed = response.output_parsed
            
            return Response(
                candidates=[Candidate(content=Content(parts=parts, role="model"))],
                usage_metadata=UsageMetadata(
                    candidates_token_count=response.usage.output_tokens,
                    prompt_token_count=response.usage.input_tokens,
                    total_token_count=response.usage.total_tokens,
                ),
                parsed=parsed,
            )
        else:
            raise ValueError(f"Unsupported result type: {result_type}. Supported types are None, 'json', or a Pydantic model class.")
    
    @_error_handler
    def create_stream(
        self,
        messages: list[Content],
        *,
        thinking_config: ThinkingConfig | None = None,
        system_message: str | None = None,
        max_tokens: int | None = None,
    ) -> Iterator[Response]:
        openai_messages = OpenaiLLMClient._content_to_openai_messages(messages, system_message)
        
        if max_tokens is None:
            max_tokens = self.default_max_tokens
        
        openai_kwargs = {
            "model": self.model,
            "max_output_tokens": max_tokens,
            "input": openai_messages,
        }
        
        if thinking_config is None:
            thinking_config = self.default_thinking_config
        openai_kwargs.update(OpenaiLLMClient._process_thinking_config(thinking_config))
        
        response = self.client.responses.create(**openai_kwargs, stream=True)
        return OpenaiStreamIterator(response)

    @staticmethod
    def models_list() -> list[Model]:
        models: list[Model] = []
        client = OpenAI()
        for openai_model in client.models.list():
            model_name = openai_model.id
            if "tts" in model_name or "whisper" in model_name:
                continue
            if "emb" in model_name:
                continue
            if "davinci" in model_name or "babbage" in model_name:
                continue
            if "image" in model_name or "dall" in model_name:
                continue
            if "moderation" in model_name:
                continue
            if "2024" in model_name or "2025" in model_name:
                continue
            
            models.append(Model(
                full_model_name=OpenaiLLMClient.PROVIDER + ":" + model_name,
                model=model_name,
                provider=OpenaiLLMClient.PROVIDER,
            ))
        return models


@inherited_decorator
def _error_handler_async(func: Callable[P, Awaitable[Response]]) -> Callable[P, Awaitable[Response]]:
    """
    Decorator to catch error from openai and transform it into unified one
    """
    @wraps(func)
    async def wrapper(*args, **kwargs):
        try:
            return await func(*args, **kwargs)
        except OpenAIAPIError as e:
            code = getattr(e, "status_code", None) or e.code
            response = getattr(e, "response", None)
            status = getattr(response, "reason_phrase", None)
            response_json = {
                "status": status,
                "message": e.message,
            }
            raise APIError(code, response_json, response)
    return wrapper

class OpenaiStreamIteratorAsync:
    def __init__(self, openai_iterator: AsyncStream[ResponseStreamEvent]):
        self._openai_iterator = openai_iterator

    async def __aiter__(self) -> AsyncIterator[Response]:
        output_tokens: int | None = None
        input_tokens: int | None = None
        total_tokens: int | None = None
        async for next_event in self._openai_iterator:
            if next_event.type == "response.output_text.delta":
                parts = [Part(text=next_event.delta)]
                yield Response(candidates=[Candidate(content=Content(parts=parts, role="model"))])
            elif next_event.type == "response.completed":
                output_tokens = next_event.response.usage.output_tokens
                input_tokens = next_event.response.usage.input_tokens
                total_tokens = next_event.response.usage.total_tokens
                
        usage_metadata = UsageMetadata(
            candidates_token_count=output_tokens,
            prompt_token_count=input_tokens,
            total_token_count=total_tokens,
        )
        yield Response(candidates=[Candidate(content=Content(parts=[Part(text="")], role="model"))], usage_metadata=usage_metadata)


class OpenaiLLMClientAsync(BaseLLMClientAsync):
    PROVIDER = "openai"
    
    def __init__(
        self,
        model: str,
        api_key: str | None = None,
        decorator_configs: DecoratorConfigs | None = None,
        default_thinking_config: ThinkingConfig | None = None,
        default_max_tokens: int | None = None,
        **kwargs,
    ):
        if api_key is None:
            api_key = os.getenv("OPENAI_API_KEY")
        if api_key is None or not isinstance(api_key, str):
            raise ValueError("To create an openai llm client you need to either set the environment variable OPENAI_API_KEY or pass the api_key in string format")
        super().__init__(OpenaiLLMClientAsync.PROVIDER, model, decorator_configs=decorator_configs, default_thinking_config=default_thinking_config, default_max_tokens=default_max_tokens)
        self._api_key = api_key
        self.client = AsyncOpenAI(api_key=api_key)
    
    @property
    def api_key(self) -> str:
        return self._api_key
    
    @_error_handler_async
    async def _create(
        self,
        messages: list[Content],
        result_type: ResultType = None,
        *,
        thinking_config: ThinkingConfig | None = None,
        system_message: str | None = None,
        max_tokens: int | None = None,
        tools: list[Tool] | None = None,
        tool_config: ToolConfig = ToolConfig(),
    ) -> Response:
        openai_messages = OpenaiLLMClient._content_to_openai_messages(messages, system_message)
        if system_message is not None:
            openai_messages.append({"role": "developer", "content": system_message})
        for message in messages:
            if message.role == "user":
                openai_messages.append({"role": "user", "content": message.as_str()})
            elif message.role == "model":
                openai_messages.append({"role": "assistant", "content": message.as_str()})
        
        if max_tokens is None:
            max_tokens = self.default_max_tokens
        
        openai_kwargs = {
            "model": self.model,
            "max_output_tokens": max_tokens,
            "input": openai_messages,
        }
        
        if thinking_config is None:
            thinking_config = self.default_thinking_config
        openai_kwargs.update(OpenaiLLMClient._process_thinking_config(thinking_config))
        
        if tools is not None:
            openai_kwargs["parallel_tool_calls"] = True
            
            openai_tools = []
            allowed_function_names = None
            if tool_config.function_calling_config is not None:
                allowed_function_names = tool_config.function_calling_config.allowed_function_names
            for tool in tools:
                for func_decl in tool.function_declarations:
                    if allowed_function_names is None or func_decl.name in allowed_function_names:
                        parameters = func_decl.parameters
                        if parameters is not None:
                            parameters = parameters.model_dump(exclude_none=True)
                            parameters["additionalProperties"] = False
                        else:
                            parameters = {"type": "object", "properties": {}, "required": [], "additionalProperties": False}
                        openai_tools.append({
                            "type": "function",
                            "name": func_decl.name,
                            "description": func_decl.description,
                            "strict": True,
                            "parameters": parameters,
                        })
            openai_kwargs["tools"] = openai_tools
            
            tool_choice_mode = "AUTO"
            if tool_config.function_calling_config is not None:
                if tool_config.function_calling_config.mode is not None:
                    tool_choice_mode = tool_config.function_calling_config.mode
            if tool_choice_mode == "NONE":
                openai_kwargs["tool_choice"] = "none"
            elif tool_choice_mode == "AUTO":
                openai_kwargs["tool_choice"] = "auto"
            elif tool_choice_mode == "ANY":
                openai_kwargs["tool_choice"] = "required"
        
        if result_type is None or result_type == "json":
            response = await self.client.responses.create(**openai_kwargs)
            
            parts: list[Part] = []
            for output_item in response.output:
                if output_item.type == "message":
                    for content in output_item.content:
                        parts.append(Part(text=content.text))
                elif output_item.type == "reasoning":
                    for summary in output_item.summary:
                        parts.append(Part(text=summary.text, thought=True))
                elif output_item.type == "function_call":
                    parts.append(Part(function_call=FunctionCall(args=json.loads(output_item.arguments), name=output_item.name)))
            
            return Response(
                candidates=[Candidate(content=Content(parts=parts, role="model"))],
                usage_metadata=UsageMetadata(
                    candidates_token_count=response.usage.output_tokens,
                    prompt_token_count=response.usage.input_tokens,
                    total_token_count=response.usage.total_tokens,
                ),
            )
        elif isinstance(result_type, type(BaseModel)):
            response = await self.client.responses.parse(**openai_kwargs, text_format=result_type)
            
            parts: list[Part] = []
            for output_item in response.output:
                if output_item.type == "message":
                    for content in output_item.content:
                        parts.append(Part(text=content.text))
                elif output_item.type == "reasoning":
                    for summary in output_item.summary:
                        parts.append(Part(text=summary.text, thought=True))
                elif output_item.type == "function_call":
                    parts.append(Part(function_call=FunctionCall(args=json.loads(output_item.arguments), name=output_item.name)))
            parsed = response.output_parsed
            
            return Response(
                candidates=[Candidate(content=Content(parts=parts, role="model"))],
                usage_metadata=UsageMetadata(
                    candidates_token_count=response.usage.output_tokens,
                    prompt_token_count=response.usage.input_tokens,
                    total_token_count=response.usage.total_tokens,
                ),
                parsed=parsed,
            )
        else:
            raise ValueError(f"Unsupported result_type: {result_type}. Supported types are: None, 'json', or a Pydantic model.")
    
    @_error_handler_async
    async def create_stream(
        self,
        messages: list[Content],
        *,
        thinking_config: ThinkingConfig | None = None,
        system_message: str | None = None,
        max_tokens: int | None = None,
    ) -> AsyncIterator[Response]:
        openai_messages = OpenaiLLMClient._content_to_openai_messages(messages, system_message)
        if system_message is not None:
            openai_messages.append({"role": "developer", "content": system_message})
        for message in messages:
            if message.role == "user":
                openai_messages.append({"role": "user", "content": message.as_str()})
            elif message.role == "model":
                openai_messages.append({"role": "assistant", "content": message.as_str()})
        
        if max_tokens is None:
            max_tokens = self.default_max_tokens
        
        openai_kwargs = {
            "model": self.model,
            "max_output_tokens": max_tokens,
            "input": openai_messages,
        }
        
        if thinking_config is None:
            thinking_config = self.default_thinking_config
        openai_kwargs.update(OpenaiLLMClient._process_thinking_config(thinking_config))
        
        response = await self.client.responses.create(**openai_kwargs, stream=True)
        return OpenaiStreamIteratorAsync(response)

    @staticmethod
    def models_list() -> list[Model]:
        return OpenaiLLMClient.models_list()
