from __future__ import annotations

import json
import logging
import os
from collections.abc import Callable
from typing import Any

from autogen_ext.models.openai import BaseOpenAIChatCompletionClient
from autogen_core import EVENT_LOGGER_NAME, FunctionCall
from autogen_core.models import ModelFamily, CreateResult
from openai import AsyncOpenAI

logger = logging.getLogger(EVENT_LOGGER_NAME)

class PapyrusChatCompletionClient(BaseOpenAIChatCompletionClient):
    """Papyrus tool-calling adapter built on top of BaseOpenAIChatCompletionClient.

    Papyrus returns tool calls as plain text, typically as multiple JSON lines, for example:

        {"name": "functions.submit_aipa_eval", "arguments": {...}}

    This adapter inspects the final text content, parses such JSON lines into a
    ``List[FunctionCall]`` and writes it back to ``CreateResult.content`` so that
    Autogen's ``FunctionTool`` mechanism can be triggered correctly.
    """

    def __init__(self, *args: Any, token_provider: Callable[[], str] | None = None, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._token_provider = token_provider

    def _refresh_api_key(self) -> None:
        if not self._token_provider:
            return
        token = self._token_provider()
        if token:
            self._client.api_key = token

    async def create(self, *args, **kwargs):  # type: ignore[override]
        """Intercept non-streaming calls and convert Papyrus text output to a FunctionCall list."""
        self._refresh_api_key()
        result = await super().create(*args, **kwargs)
        if isinstance(result, CreateResult) and isinstance(result.content, str):
            result = self._convert_papyrus_text_to_tool_calls(result)
        return result

    async def create_stream(self, *args, **kwargs):  # type: ignore[override]
        """Intercept streaming calls and convert the final CreateResult once."""
        self._refresh_api_key()
        async for item in super().create_stream(*args, **kwargs):
            if isinstance(item, CreateResult) and isinstance(item.content, str):
                item = self._convert_papyrus_text_to_tool_calls(item)
            yield item

    @staticmethod
    def _convert_papyrus_text_to_tool_calls(result: CreateResult) -> CreateResult:
        text = result.content
        if not isinstance(text, str):
            return result

        tool_calls: list[FunctionCall] = []
        for idx, line in enumerate(text.splitlines()):
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except Exception:
                continue
            if not isinstance(obj, dict):
                continue
            name = obj.get("name")
            args = obj.get("arguments")
            if not isinstance(name, str) or args is None:
                continue

            # Normalize "functions.submit_aipa_eval" -> "submit_aipa_eval"
            if name.startswith("functions."):
                name = name.split(".", 1)[1]

            if isinstance(args, (dict, list)):
                arg_str = json.dumps(args, ensure_ascii=False)
            else:
                arg_str = str(args)

            tool_calls.append(
                FunctionCall(
                    id=f"tool_{idx}",
                    name=name,
                    arguments=arg_str,
                )
            )

        if not tool_calls:
            return result

        # Use Pydantic's model_copy to update only the content field and keep others unchanged
        return result.model_copy(update={"content": tool_calls})

