from typing import override, MutableSequence, Any
from textwrap import dedent
from copy import deepcopy

from pydantic import BaseModel
from agent_framework.openai import OpenAIChatClient
from agent_framework import ChatMessage, ChatOptions, TextContent

class OpenAILikeChatClient(OpenAIChatClient):
    @override
    def _prepare_options(self, messages: MutableSequence[ChatMessage], options: dict[str, Any]) -> dict[str, Any]:
        chat_options_copy = deepcopy(options)
        response_format = chat_options_copy.get("response_format")

        if (
            response_format
            and isinstance(response_format, type)
            and issubclass(response_format, BaseModel)
        ):
            structured_output_prompt = self._build_structured_prompt(response_format)
            if old_instructions := chat_options_copy.get("instructions"):
                chat_options_copy["instructions"] = f"{old_instructions}\n\n{structured_output_prompt}"
            else:
                messages = [ChatMessage(role=Role.SYSTEM, text=structured_output_prompt), *messages]
            
            chat_options_copy["response_format"] = {"type": "json_object"}

        return super()._prepare_options(messages, chat_options_copy)

    @staticmethod
    def _build_structured_prompt(response_format: type[BaseModel]) -> str:
        json_schema = response_format.model_json_schema()
        structured_output_prompt = dedent(f"""
        <output-format>\n
        Your output must adhere to the following JSON schema format,
        without any Markdown syntax, and without any preface or explanation:\n
        {json_schema}\n
        </output-format>
        """)

        return structured_output_prompt