# retoor <retoor@molodetz.nl>

from __future__ import annotations

import json
from typing import Any, Callable, Optional

from ..errors import ToolInputError
from .tracker import get_tracker


class CostController:
    def __init__(
        self, quota_provider: Optional[Callable[[], dict[str, Any]]] = None
    ) -> None:
        self._quota_provider = quota_provider

    async def dispatch(self, name: str, arguments: dict[str, Any]) -> str:
        if name == "usage_quota":
            if self._quota_provider is None:
                return json.dumps(
                    {
                        "status": "unavailable",
                        "message": "No quota information is available for this session.",
                    }
                )
            return json.dumps(
                {"status": "success", **self._quota_provider()}, ensure_ascii=False
            )
        if name == "cost_stats":
            tracker = get_tracker()
            if tracker is None:
                return json.dumps(
                    {
                        "status": "unavailable",
                        "message": "No cost tracker is active for this session.",
                    }
                )
            return json.dumps(
                {"status": "success", **tracker.stats()}, ensure_ascii=False
            )
        raise ToolInputError(f"Unknown cost tool: {name}")
