import json
from typing import Final

from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import MCPToolsetRepository
from litellm.types.mcp_server.mcp_toolset import (
    MCPToolset,
    NewMCPToolsetRequest,
    UpdateMCPToolsetRequest,
)


def _toolset_from_row(row) -> MCPToolset:
    data: Final = row.model_dump()
    tools = data.get("tools") or []
    if isinstance(tools, str):
        tools = json.loads(tools)
    data["tools"] = tools
    return MCPToolset(**data)


async def create_mcp_toolset(
    prisma_client: PrismaClient,
    data: NewMCPToolsetRequest,
    touched_by: str,
) -> MCPToolset:
    data_dict: Final = data.model_dump(exclude_none=True)
    data_dict["toolset_id"] = str(uuid.uuid4())
    data_dict["tools"] = json.dumps(data_dict.get("tools", []))
    data_dict["created_by"] = touched_by
    data_dict["updated_by"] = touched_by
    row: Final = await MCPToolsetRepository(prisma_client).table.create(data=data_dict)
    return _toolset_from_row(row)


async def get_mcp_toolset(
    prisma_client: PrismaClient,
    toolset_id: str,
) -> MCPToolset | None:
    row: Final = await MCPToolsetRepository(prisma_client).table.find_unique(where={"toolset_id": toolset_id})
    if row is None:
        return None
    return _toolset_from_row(row)


async def list_mcp_toolsets(
    prisma_client: PrismaClient,
    toolset_ids: list[str] | None = None,
) -> list[MCPToolset]:
    try:
        where = {}
        if toolset_ids is not None:
            where = {"toolset_id": {"in": toolset_ids}}
        rows: Final = await MCPToolsetRepository(prisma_client).table.find_many(where=where)
        return [_toolset_from_row(r) for r in rows]
    except Exception as e:
        verbose_proxy_logger.warning("litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - %s", e)
        return []


async def get_mcp_toolset_by_name(
    prisma_client: PrismaClient,
    toolset_name: str,
) -> MCPToolset | None:
    row: Final = await MCPToolsetRepository(prisma_client).table.find_first(where={"toolset_name": toolset_name})
    if row is None:
        return None
    return _toolset_from_row(row)


async def update_mcp_toolset(
    prisma_client: PrismaClient,
    data: UpdateMCPToolsetRequest,
    touched_by: str,
) -> MCPToolset | None:
    data_dict: Final = data.model_dump(exclude_none=True, exclude={"toolset_id"})
    if "tools" in data_dict:
        data_dict["tools"] = json.dumps(data_dict["tools"])
    data_dict["updated_by"] = touched_by
    try:
        row: Final = await MCPToolsetRepository(prisma_client).table.update(
            where={"toolset_id": data.toolset_id},
            data=data_dict,
        )
    except Exception as e:
        from prisma.errors import RecordNotFoundError

        if isinstance(e, RecordNotFoundError):
            return None
        raise
    return _toolset_from_row(row)


async def delete_mcp_toolset(
    prisma_client: PrismaClient,
    toolset_id: str,
) -> MCPToolset | None:
    try:
        row: Final = await MCPToolsetRepository(prisma_client).table.delete(where={"toolset_id": toolset_id})
    except Exception as e:
        from prisma.errors import RecordNotFoundError

        if isinstance(e, RecordNotFoundError):
            return None
        raise
    return _toolset_from_row(row)