""" Implementation of the OpenAI Responses APImodels.

Feature table:
    - Async chat:       YES
    - Return JSON:      NO
    - Structured types: YES, via Pydantic  TODO: Add support for native Python types
    - Token count:      YES
    - Image support:    YES
    - Tool use:         YES
    - File input:       NOT YET IMPLEMENTED
    - Web search:       NOT YET IMPLEMENTED
"""

import asyncio
import json
import os
import time
from typing import Any, AsyncGenerator
import pydantic

import instructor
import tiktoken
from dotenv import dotenv_values
from openai import OpenAI, NOT_GIVEN, APIConnectionError, \
    RateLimitError, APITimeoutError, AuthenticationError, PermissionDeniedError, BadRequestError
from openai.types.beta.realtime import ResponseTextDeltaEvent

from justai.model.message import Message, ToolUseMessage
from justai.models.basemodel import ImageInput
from justai.tools.display import color_print, ERROR_COLOR, DEBUG_COLOR1, DEBUG_COLOR2
from justai.models.basemodel import BaseModel, ConnectionException, AuthorizationException, \
    ModelOverloadException, RatelimitException, BadRequestException, GeneralException




class OpenAIResponsesModel(BaseModel):
    def __init__(self, model_name: str, params: dict = None):
        system_message = f"You are {model_name}, a large language model trained by OpenAI."
        super().__init__(model_name, params, system_message)

        # Authentication
        api_key = params.get("OPENAI_API_KEY") or os.getenv("OPENAI_API_KEY") or dotenv_values()["OPENAI_API_KEY"]
        if not api_key:
            color_print("No OpenAI API key found. Create one at https://platform.openai.com/account/api-keys and " +
                        "set it in the .env file like OPENAI_API_KEY=here_comes_your_key.", color=ERROR_COLOR)

        # instructor.patch makes the OpenAI client compatible with structured output via response_model=”
        # Works only for OpenAI models

        # Not sure if this works, or is needed, for the Responses API
        # self.client = instructor.patch(OpenAI(api_key=api_key))
        self.client = OpenAI()

        # Diversions from the features that are supported or not supported by default
        self.supports_function_calling = True

        self.last_response_id = None

    def prompt(self, prompt: str, images: list[ImageInput], tools, return_json: bool, response_format, _chat=False) \
            -> tuple[Any, int|None, int|None]:

        # _chat can only be called from chat()

        if return_json:
            raise NotImplementedError("OpenAI Responses API does not support return_json. Pass a response_format instead.")
        if response_format and not issubclass(response_format, pydantic.BaseModel):
            raise NotImplementedError("OpenAI Responses API requires response_format to be a Pydantic model.")

        content = self.create_content(prompt, images)
        input_list = [{'role': 'system', 'content': self.system_message},
                      {'role': 'user', 'content': content}]
        tool_spec = self.create_tool_spec(tools)

        last_response_id = self.last_response_id if _chat else None

        for run in range(3):  # Max 3 function calls to prevent infinite loop
            try:
                if response_format:
                    response = self.client.responses.parse(model=self.model_name, input=input_list, tools=tool_spec,
                                                           text_format=response_format,
                                                           previous_response_id=last_response_id)
                else:
                    response = self.client.responses.create(model=self.model_name, input=input_list, tools=tool_spec,
                                                            previous_response_id=last_response_id)
            except APITimeoutError as e:
                raise ModelOverloadException(e)
            except APIConnectionError as e:
                raise ConnectionException(e)
            except (AuthenticationError, PermissionDeniedError) as e:
                raise AuthorizationException(e)
            except RateLimitError as e:
                raise RatelimitException(e)
            except BadRequestError as e:
                raise BadRequestException(e)
            except Exception as e:
                raise GeneralException(e)

            # Save the response id for subsequent requests
            self.last_response_id = response.id if _chat else None

            # Save function call outputs for subsequent requests
            function_call = None
            function_call_arguments = None
            input_list += response.output

            for item in response.output:
                if item.type == "function_call":
                    function_call = item
                    function_call_arguments = json.loads(item.arguments)

            if not function_call or run == 2:
                output = response.output_parsed if response_format else response.output_text
                return output, response.usage.input_tokens, response.usage.output_tokens

            # The rest of the function is to process the function call output
            # Numbering is from OpenAI's docs: https://platform.openai.com/docs/guides/function-calling#function-tool-example

            for tool in tools:
                if tool['function'].__name__ == function_call.name:
                    function = tool['function']
                    break
            else:
                raise ValueError(f"Function {function_call.name} not found")

            # 3. Execute the function logic for get_horoscope
            result = {function_call.name: function(*function_call_arguments.values())}

            # 4. Provide function call results to the model
            input_list.append(
                {
                    "type": "function_call_output",
                    "call_id": function_call.call_id,
                    "output": json.dumps(result),
                }
            )
        # Response looks like this:
        # [{
        #   "id": "resp_67ccd2bed1ec8190b14f964abc0542670bb6a6b452d3795b",
        #   "object": "response",
        #   "created_at": 1741476542,
        #   "status": "completed",
        #   "error": null,
        #   "incomplete_details": null,
        #   "instructions": null,
        #   "max_output_tokens": null,
        #   "model": "gpt-4.1-2025-04-14",
        #   "output": [
        #     {
        #       "type": "message",
        #       "id": "msg_67ccd2bf17f0819081ff3bb2cf6508e60bb6a6b452d3795b",
        #       "status": "completed",
        #       "role": "assistant",
        #       "content": [
        #         {
        #           "type": "output_text",
        #           "text": "In a peaceful grove beneath a silver moon, a unicorn named Lumina discovered a hidden pool that reflected the stars. As she dipped her horn into the water, the pool began to shimmer, revealing a pathway to a magical realm of endless night skies. Filled with wonder, Lumina whispered a wish for all who dream to find their own hidden magic, and as she glanced back, her hoofprints sparkled like stardust.",
        #           "annotations": []
        #         }
        #       ]
        #     }
        #   ],
        #   "parallel_tool_calls": true,
        #   "previous_response_id": null,
        #   "reasoning": {
        #     "effort": null,
        #     "summary": null
        #   },
        #   "store": true,
        #   "temperature": 1.0,
        #   "text": {
        #     "format": {
        #       "type": "text"
        #     }
        #   },
        #   "tool_choice": "auto",
        #   "tools": [],
        #   "top_p": 1.0,
        #   "truncation": "disabled",
        #   "usage": {
        #     "input_tokens": 36,
        #     "input_tokens_details": {
        #       "cached_tokens": 0
        #     },
        #     "output_tokens": 87,
        #     "output_tokens_details": {
        #       "reasoning_tokens": 0
        #     },
        #     "total_tokens": 123
        #   },
        #   "user": null,
        #   "metadata": {}
        # }]


    async def prompt_async(self, prompt: str, images: list[ImageInput]=None, _chat=False) -> AsyncGenerator[str]:

        content = self.create_content(prompt, images)
        input_ = [{"role": "user", "content": content}]

        last_response_id = self.last_response_id if _chat else None

        response = self.client.responses.create(model=self.model_name, input=input_, stream=True,
                                                previous_response_id=last_response_id)

        # Save the response id for subsequent requests
        self.last_response_id = response.id if _chat else None

        for event in response:
            if hasattr(event, 'delta'):
                yield event.delta
        # Events can have different types:
        # event: response.created
        # data: {
        #     "type":"response.created",
        #     "response":{
        #         "id":"resp_67c9fdcecf488190bdd9a0409de3a1ec07b8b0ad4e5eb654",
        #         "object":"response",
        #         "created_at":1741290958,
        #         "status":"in_progress",
        #         "error":null,
        #         "incomplete_details":null,
        #         "instructions":"You are a helpful assistant.",
        #         "max_output_tokens":null,
        #         "model":"gpt-4.1-2025-04-14","
        #         output":[],"p
        #         arallel_tool_calls":true,
        #         "previous_response_id":null,
        #         "reasoning":{
        #             "effort":null,
        #             "summary":null},
        #         "store":true,
        #         "temperature":1.0,
        #         "text":{
        #             "format":{
        #                 "type":"text"}},
        #         "tool_choice":"auto",
        #         "tools":[],
        #         "top_p":1.0,
        #         "truncation":
        #         "disabled",
        #         "usage":null,
        #         "user":null,
        #         "metadata":{}}}
        #
        # event: response.in_progress
        # data: {
        #     "type":"response.in_progress",
        #     "response":{"
        #         id":"resp_67c9fdcecf488190bdd9a0409de3a1ec07b8b0ad4e5eb654",
        #         "object":"response",
        #         "created_at":1741290958,
        #         "status":"in_progress",
        #         "error":null,
        #         "incomplete_details":null,
        #         "instructions":"You are a helpful assistant.",
        #         "max_output_tokens":null,
        #         "model":"gpt-4.1-2025-04-14",
        #         "output":[],
        #         "parallel_tool_calls":true,
        #         "previous_response_id":null,
        #         "reasoning":{
        #             "effort":null,
        #             "summary":null},
        #         "store":true,
        #         "temperature":1.0,
        #         "text":{
        #             "format":{
        #                 "type":"text"}},
        #         "tool_choice":"auto",
        #         "tools":[],
        #         "top_p":1.0,
        #         "truncation":"disabled",
        #         "usage":null,
        #         "user":null,
        #         "metadata":{}}}
        #
        # event: response.output_item.added
        # data: {
        #     "type":"response.output_item.added",
        #     "output_index":0,
        #     "item":{
        #         "id":"msg_67c9fdcf37fc8190ba82116e33fb28c507b8b0ad4e5eb654",
        #         "type":"message","status":"
        #         in_progress","
        #         role":"assistant",
        #         "content":[]}}
        #
        # event: response.content_part.added
        # data: {
        #     "type":"response.content_part.added",
        #     "item_id":"msg_67c9fdcf37fc8190ba82116e33fb28c507b8b0ad4e5eb654",
        #     "output_index":0,
        #     "content_index":0,
        #     "part":{
        #         "type":"output_text",
        #         "text":"",
        #         "annotations":[]}}
        #
        # event: response.output_text.delta
        # data: {
        #     "type":"response.output_text.delta",
        #     "item_id":"msg_67c9fdcf37fc8190ba82116e33fb28c507b8b0ad4e5eb654",
        #     "output_index":0,
        #     "content_index":0,
        #     "delta":"Hi"}
        #
        # ...
        #
        # event: response.output_text.done
        # data: {"type":"response.output_text.done","item_id":"msg_67c9fdcf37fc8190ba82116e33fb28c507b8b0ad4e5eb654","output_index":0,"content_index":0,"text":"Hi there! How can I assist you today?"}
        #
        # event: response.content_part.done
        # data: {"type":"response.content_part.done","item_id":"msg_67c9fdcf37fc8190ba82116e33fb28c507b8b0ad4e5eb654","output_index":0,"content_index":0,"part":{"type":"output_text","text":"Hi there! How can I assist you today?","annotations":[]}}
        #
        # event: response.output_item.done
        # data: {"type":"response.output_item.done","output_index":0,"item":{"id":"msg_67c9fdcf37fc8190ba82116e33fb28c507b8b0ad4e5eb654","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"Hi there! How can I assist you today?","annotations":[]}]}}
        #
        # event: response.completed
        # data: {"type":"response.completed","response":{"id":"resp_67c9fdcecf488190bdd9a0409de3a1ec07b8b0ad4e5eb654","object":"response","created_at":1741290958,"status":"completed","error":null,"incomplete_details":null,"instructions":"You are a helpful assistant.","max_output_tokens":null,"model":"gpt-4.1-2025-04-14","output":[{"id":"msg_67c9fdcf37fc8190ba82116e33fb28c507b8b0ad4e5eb654","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"Hi there! How can I assist you today?","annotations":[]}]}],"parallel_tool_calls":true,"previous_response_id":null,"reasoning":{"effort":null,"summary":null},"store":true,"temperature":1.0,"text":{"format":{"type":"text"}},"tool_choice":"auto","tools":[],"top_p":1.0,"truncation":"disabled","usage":{"input_tokens":37,"output_tokens":11,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":48},"user":null,"metadata":{}}}

    def chat(self, prompt: str, images: list[ImageInput], tools, return_json: bool, response_format) \
             -> tuple[Any, int|None, int|None]:
        return self.prompt(prompt, images, tools, return_json, response_format, _chat=True)


    async def chat_async(self, prompt: str, images: list[ImageInput]) -> AsyncGenerator[tuple[str, str], None]:
        async for chunk in self.prompt_async(prompt, images, _chat=True):
            yield chunk
               
    @staticmethod
    def create_content(prompt, images:list[ImageInput]) -> list[dict]:
        content = [{"type": "input_text", "text": prompt}]

        if images:
            for image in images:
                image_type = Message.get_image_type(image)
                image_url = image if image_type == "image_url" else Message.to_base64_image(image)
                content += [{"type": "input_image", "image_url": image_url}]
        return content

    @staticmethod
    def create_tool_spec(tools: list[dict]) -> list[dict]:
        if not tools:
            return []

        type_mapping = {
            int: "integer",
            str: "string",
            float: "number",
            bool: "boolean",
            list: "array",
            dict: "object"
        }
        tool_spec = []
        for tool in tools:
            tool_spec += [{
                "type": "function",
                "name": tool["function"].__name__,
                "description": tool['description'],
                "parameters": {
                    "type": "object",
                    "properties": {
                        param_name: {
                            "type": type_mapping.get(_type, "string"),  # Default to string if type not in map
                            "description": param_name,  # Or a more descriptive text if available
                        }
                        for param_name, _type in tool['parameters'].items()
                    },
                    "required": tool['required_parameters'] or [],
                }
            }]
        return tool_spec

    @staticmethod
    def transform_messages(messages: list[Message]) -> list[dict]:
        transformed_messages = []

        for message in messages:
            msg = {"role": message.role}

            # Handle tool messages (function calls and their results)
            # Todo: tool use is geen onderdeel meer van message maar staat in de Model class.
            if isinstance(message, ToolUseMessage):
                if message.role == 'assistant' and 'function_to_call' in message.tool_use:
                    # This is a function call from the assistant
                    msg["content"] = None
                    msg["tool_calls"] = [{
                        "id": message.tool_use.get('call_id', 'call_' + str(hash(str(message.tool_use)))),
                        "type": "function",
                        "function": {
                            "name": message.tool_use['function_to_call'],
                            "arguments": json.dumps(message.tool_use['function_parameters'])
                        }
                    }]
                elif message.role == 'tool':
                    # This is a function result
                    function_result = message.tool_use.get('function_result', '')
                    if not isinstance(function_result, str):
                        function_result = json.dumps(function_result)
                    msg["content"] = function_result
                    msg["tool_call_id"] = message.tool_use.get('call_id', '')
                    msg["name"] = message.tool_use.get('function_to_call', '')
            # Handle regular messages
            else:
                if message.images:
                    content = [{"type": "text", "text": message.content or ""}]
                    for image in message.images:
                        content.append({
                            "type": "image_url",
                            "image_url": {'url': f"data:image/jpeg;base64,{Message.to_base64_image(image)}"}
                        })
                    msg["content"] = content
                else:
                    msg["content"] = message.content or ""

            transformed_messages.append(msg)

        return transformed_messages

    def token_count(self, text: str) -> int:
        """ Returns the number of tokens in a string. """
        try:
            encoding = tiktoken.encoding_for_model(self.model_name)
        except KeyError:
            # Fall back to cl100k_base encoding for newer models not yet in tiktoken
            encoding = tiktoken.get_encoding("cl100k_base")
        return len(encoding.encode(text))
