"""OneShot account tools (notifications, balance)."""

from __future__ import annotations

import json
from typing import Any, Type

from langchain_core.tools import BaseTool
from pydantic import BaseModel

from langchain_oneshot._types import (
    GetBalanceInput,
    MarkNotificationReadInput,
    NotificationsInput,
)
from oneshot import OneShotClient


class NotificationsTool(BaseTool):
    """List agent notifications. Free."""

    name: str = "oneshot_notifications"
    description: str = (
        "List notifications for the agent (job completions, failures, warnings). "
        "Optionally filter to unread only. Free."
    )
    args_schema: Type[BaseModel] = NotificationsInput
    handle_tool_error: bool = True

    client: OneShotClient

    def _run(self, **kwargs: Any) -> str:
        params: dict[str, str] = {}
        if kwargs.get("unread") is not None:
            params["unread"] = str(kwargs["unread"]).lower()
        if kwargs.get("limit") is not None:
            params["limit"] = str(kwargs["limit"])
        result = self.client.call_free_get("/v1/tools/notifications", params or None)
        return json.dumps(result)

    async def _arun(self, **kwargs: Any) -> str:
        params: dict[str, str] = {}
        if kwargs.get("unread") is not None:
            params["unread"] = str(kwargs["unread"]).lower()
        if kwargs.get("limit") is not None:
            params["limit"] = str(kwargs["limit"])
        result = await self.client.acall_free_get("/v1/tools/notifications", params or None)
        return json.dumps(result)


class MarkNotificationReadTool(BaseTool):
    """Mark a notification as read. Free."""

    name: str = "oneshot_mark_notification_read"
    description: str = "Mark a specific notification as read by its ID. Free."
    args_schema: Type[BaseModel] = MarkNotificationReadInput
    handle_tool_error: bool = True

    client: OneShotClient

    def _run(self, notification_id: str, **kwargs: Any) -> str:
        result = self.client.call_free_patch(
            f"/v1/tools/notifications/{notification_id}/read"
        )
        return json.dumps(result)

    async def _arun(self, notification_id: str, **kwargs: Any) -> str:
        result = await self.client.acall_free_patch(
            f"/v1/tools/notifications/{notification_id}/read"
        )
        return json.dumps(result)


class GetBalanceTool(BaseTool):
    """Get the agent's USDC balance (on-chain and credits). Free."""

    name: str = "oneshot_get_balance"
    description: str = (
        "Get the agent's USDC balance including on-chain balance and credits. "
        "Returns on_chain_balance, credits_balance, currency, address, and chain_id. Free."
    )
    args_schema: Type[BaseModel] = GetBalanceInput
    handle_tool_error: bool = True

    client: OneShotClient

    def _run(self, **kwargs: Any) -> str:
        result = self.client.get_balance()
        return json.dumps(result)

    async def _arun(self, **kwargs: Any) -> str:
        result = await self.client.aget_balance()
        return json.dumps(result)
