from __future__ import annotations

import inspect
import json
import traceback
from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable, Sequence
from enum import StrEnum
from functools import lru_cache
from typing import Self, cast

from judy_core.types import ContentUnion, Promptable

# from jsonschema.exceptions import ValidationError
from openai.types.chat import ChatCompletionFunctionToolParam
from openai.types.shared_params import FunctionDefinition
from pydantic import BaseModel, TypeAdapter
from pydantic import ValidationError as PydanticValidationError
from tongcell.util import getLogger

logger = getLogger(__name__)


@lru_cache(maxsize=5000)
def get_cached_typeadapter[T](cls: T) -> TypeAdapter[T]:
    """
    TypeAdapters are heavy objects, and in an application context we'd typically
    create them once in a global scope and reuse them as often as possible.
    However, this isn't feasible for user-generated functions. Instead, we use a
    cache to minimize the cost of creating them as much as possible.
    """
    return TypeAdapter(cls)


type JsonSchema = dict[str, object]


class ToolBasicDefinition(BaseModel):
    name: str
    description: str


class ToolDefinition(ToolBasicDefinition):
    parameters: JsonSchema

    def to_openai_format(self) -> ChatCompletionFunctionToolParam:
        return ChatCompletionFunctionToolParam(
            type="function",
            function=FunctionDefinition(
                name=self.name,
                description=self.description,
                parameters=self.parameters,
            ),
        )


def to_openai_format(
    tools_definition: Sequence[ToolDefinition],
) -> list[ChatCompletionFunctionToolParam]:
    return [tool_def.to_openai_format() for tool_def in tools_definition]


TOOL_RESULT_SUCCESS_CONTENT = "success"


class ToolErrorKind(StrEnum):
    NOT_FOUND = "not_found"
    INVALID_ARGS = "invalid_args"
    EXEC_FAIL = "tool_exec_fail"


class ToolErrorBase(ABC, Exception):
    kind: ToolErrorKind

    @abstractmethod
    def detail_message(self) -> str: ...


class ToolNotFound(ToolErrorBase):
    kind = ToolErrorKind.NOT_FOUND

    def __init__(self, name: str):
        super().__init__(f"tool not found:{name}")
        self.name = name

    def detail_message(self) -> str:
        return str(self)


class InvalidToolArgs(ToolErrorBase):
    kind = ToolErrorKind.INVALID_ARGS

    def __init__(self, ex: Exception):
        super().__init__(f"invalid tool args: {str(ex)}")
        self.ex = ex

    def detail_message(self) -> str:
        return str(self) + "\n" + "\n".join(traceback.format_exception(self.ex))


class ToolRuntimeError(ToolErrorBase):
    kind = ToolErrorKind.EXEC_FAIL

    def __init__(self, ex: Exception):
        super().__init__(f"tool runtime error: {str(ex)}")
        self.ex = ex

    def detail_message(self) -> str:
        return str(self) + "\n" + "\n".join(traceback.format_exception(self.ex))


class ToolExecResult:
    def __init__(self, content: ContentUnion, error: ToolErrorBase | None = None):
        self.content = content
        self.error = error

    @classmethod
    def success(cls, content: ContentUnion) -> Self:
        return cls(content=content, error=None)

    @classmethod
    def fail(cls, tool_error: ToolErrorBase) -> Self:
        text = "type: {}, msg:{}".format(type(tool_error).__name__, tool_error)
        return cls(
            content=text,
            error=tool_error,
        )

    def is_ok(self) -> bool:
        return self.error is None


def try_fix_model_args(args: str):
    decoded_args = json.loads(args)
    if isinstance(decoded_args, dict):
        for k, v in decoded_args.items():
            if isinstance(v, str):
                try:
                    decoded_args[k] = json.loads(v)
                except Exception:
                    pass
    return decoded_args


async def call_single_tool[T](tool: Callable[..., Awaitable[T]], args: str) -> T:
    type_adapter = get_cached_typeadapter(tool)
    try:
        # jsonschema.validate(dict_args, type_adapter.json_schema())
        coro = type_adapter.validate_json(args)
    except PydanticValidationError as e:
        logger.warning("tool args validate failed, trying autofix")
        try:
            fixed_args = try_fix_model_args(args)
            coro = type_adapter.validate_python(fixed_args)
        except Exception:
            raise InvalidToolArgs(e) from None

    coro = cast(Awaitable[T], coro)
    try:
        result_content = await coro
    except Exception as e:
        raise ToolRuntimeError(e)
    return result_content


def camel_to_snake(name: str) -> str:
    result = []

    n = len(name)
    for i, c in enumerate(name):
        if c.isupper():
            if i > 0:
                # 如果前一个字符或者后一个字符是小写，添加下划线
                if name[i - 1].islower() or (i + 1 < n and name[i + 1].islower()):
                    result.append("_")
            result.append(c.lower())
        else:
            result.append(c)

    return "".join(result)


type AsyncToolCallable[R] = Callable[..., Awaitable[R]]


class Tool[R]:
    def __init__(
        self, function: AsyncToolCallable[R], basic_definition: ToolBasicDefinition
    ):
        self._function = function
        self._definition = ToolDefinition(
            name=basic_definition.name,
            description=basic_definition.description,
            parameters=Tool.get_function_parameters(function),
        )

    def get_function(self) -> AsyncToolCallable[R]:
        return self._function

    def get_definition(self) -> ToolDefinition:
        return self._definition

    @staticmethod
    def get_function_parameters(function: AsyncToolCallable) -> JsonSchema:
        return get_cached_typeadapter(function).json_schema()

    @staticmethod
    def from_callable[R2](
        callable: AsyncToolCallable[R2],
        name: str | None = None,
        description: str | None = None,
    ) -> Tool[R2]:
        if inspect.iscoroutinefunction(callable):
            tool_function = callable
            default_name = callable.__name__
            default_description = callable.__doc__ or ""
        else:
            tool_function = callable.__call__
            callable_cls = type(callable)
            default_name = camel_to_snake(callable_cls.__name__)
            default_description = tool_function.__doc__ or ""

        tool_basic_def = ToolBasicDefinition(
            name=name or default_name,
            description=description or default_description,
        )

        return Tool(tool_function, tool_basic_def)


ToolReturnType = ContentUnion | Promptable


class ToolRegistry:
    def __init__(
        self,
        tools: Sequence[Tool[ToolReturnType] | AsyncToolCallable[ToolReturnType]],
    ):
        self._tool_dict = dict[str, Tool[ToolReturnType]]()
        for tool in tools:
            if not isinstance(tool, Tool):
                tool = Tool.from_callable(tool)
            self._tool_dict[tool.get_definition().name] = tool

    def get_tool_names(self) -> set[str]:
        return set(self._tool_dict.keys())

    def get_tools_definition(self) -> list[ToolDefinition]:
        definitions = []
        for name, tool in self._tool_dict.items():
            definitions.append(tool.get_definition())
        return definitions

    async def call_tool(self, name: str, args: str) -> ToolExecResult:
        tool = self._tool_dict.get(name)
        if tool is None:
            return ToolExecResult.fail(ToolNotFound(name))
        try:
            tool_result_content = await call_single_tool(tool.get_function(), args)
            if isinstance(tool_result_content, Promptable):
                tool_result_content = tool_result_content.to_prompt()
            return ToolExecResult.success(tool_result_content)
        except ToolErrorBase as tool_error:
            logger.warning(
                "⛔⛔⛔ name: %s, args: %s, error: %s",
                name,
                args,
                tool_error.detail_message(),
            )
            return ToolExecResult.fail(tool_error)
