"""OneShot browser automation tools."""

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 (
    BrowserCreateProfileInput,
    BrowserDeleteProfileInput,
    BrowserInput,
)
from oneshot import OneShotClient


class BrowserTool(BaseTool):
    """Run a browser automation task via OneShot. Paid tool."""

    name: str = "oneshot_browser"
    description: str = (
        "Run a browser automation task. The agent navigates web pages, fills forms, "
        "clicks buttons, and extracts data. Supports persistent profiles and "
        "domain-scoped credentials. Paid tool."
    )
    args_schema: Type[BaseModel] = BrowserInput
    handle_tool_error: bool = True

    client: OneShotClient

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

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


class BrowserCreateProfileTool(BaseTool):
    """Create a persistent browser profile. Free tool."""

    name: str = "oneshot_browser_create_profile"
    description: str = (
        "Create a persistent browser profile that retains cookies and session state "
        "across tasks. Returns a profile_id to reuse. Free tool."
    )
    args_schema: Type[BaseModel] = BrowserCreateProfileInput
    handle_tool_error: bool = True

    client: OneShotClient

    def _run(self, name: str) -> str:
        result = self.client.create_browser_profile(name)
        return json.dumps(result)

    async def _arun(self, name: str) -> str:
        result = await self.client.acreate_browser_profile(name)
        return json.dumps(result)


class BrowserListProfilesTool(BaseTool):
    """List all browser profiles. Free tool."""

    name: str = "oneshot_browser_list_profiles"
    description: str = "List all persistent browser profiles for this agent. Free tool."
    handle_tool_error: bool = True

    client: OneShotClient

    def _run(self) -> str:
        result = self.client.list_browser_profiles()
        return json.dumps(result)

    async def _arun(self) -> str:
        result = await self.client.alist_browser_profiles()
        return json.dumps(result)


class BrowserDeleteProfileTool(BaseTool):
    """Delete a browser profile. Free tool."""

    name: str = "oneshot_browser_delete_profile"
    description: str = "Delete a persistent browser profile by ID. Free tool."
    args_schema: Type[BaseModel] = BrowserDeleteProfileInput
    handle_tool_error: bool = True

    client: OneShotClient

    def _run(self, profile_id: str) -> str:
        result = self.client.delete_browser_profile(profile_id)
        return json.dumps(result)

    async def _arun(self, profile_id: str) -> str:
        result = await self.client.adelete_browser_profile(profile_id)
        return json.dumps(result)


def _strip_none(d: dict[str, Any]) -> dict[str, Any]:
    return {k: v for k, v in d.items() if v is not None}
