import asyncio
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime
import json
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

from pydantic import BaseModel, ConfigDict, Field, ValidationError


class CurrentTimeArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")

    timezone: str = Field(default="UTC", description="IANA timezone, e.g. Asia/Shanghai")


@dataclass(frozen=True, slots=True)
class ToolSpec:
    name: str
    description: str
    arguments_model: type[BaseModel]
    handler: Callable[[BaseModel], dict[str, object]]

    def api_definition(self) -> dict[str, object]:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.arguments_model.model_json_schema(),
            },
        }


@dataclass(frozen=True, slots=True)
class ToolExecutionResult:
    content: str
    is_error: bool


class ToolRegistry:
    def __init__(self, timeout_seconds: float = 5.0) -> None:
        if timeout_seconds <= 0:
            raise ValueError("tool timeout must be greater than zero")
        self.timeout_seconds = timeout_seconds
        specs = [
            ToolSpec(
                name="get_current_time",
                description="Get the current date and time in an IANA timezone.",
                arguments_model=CurrentTimeArguments,
                handler=_get_current_time,
            )
        ]
        self._specs = {spec.name: spec for spec in specs}

    def definitions(self) -> list[dict[str, object]]:
        return [spec.api_definition() for spec in self._specs.values()]

    async def execute(self, name: str, arguments: str) -> ToolExecutionResult:
        spec = self._specs.get(name)
        if spec is None:
            return _tool_error("unknown_tool", f"Unknown tool: {name}")

        try:
            raw_arguments = json.loads(arguments)
            if not isinstance(raw_arguments, dict):
                raise ValueError("arguments must be a JSON object")
            parsed_arguments = spec.arguments_model.model_validate(raw_arguments)
        except (json.JSONDecodeError, ValidationError, ValueError) as exc:
            return _tool_error("invalid_arguments", str(exc))

        try:
            result = await asyncio.wait_for(
                asyncio.to_thread(spec.handler, parsed_arguments),
                timeout=self.timeout_seconds,
            )
            return ToolExecutionResult(
                content=json.dumps(result, ensure_ascii=False),
                is_error=False,
            )
        except TimeoutError:
            return _tool_error("timeout", f"Tool {name} timed out")
        except Exception:
            return _tool_error("execution_failed", f"Tool {name} failed")


def _tool_error(code: str, message: str) -> ToolExecutionResult:
    return ToolExecutionResult(
        content=json.dumps(
            {"error": {"code": code, "message": message}},
            ensure_ascii=False,
        ),
        is_error=True,
    )


def _get_current_time(arguments: BaseModel) -> dict[str, object]:
    assert isinstance(arguments, CurrentTimeArguments)
    try:
        timezone = ZoneInfo(arguments.timezone)
    except ZoneInfoNotFoundError as exc:
        raise ValueError("unknown timezone") from exc
    now = datetime.now(timezone)
    return {"timezone": arguments.timezone, "datetime": now.isoformat()}