"""Tools for LangSmith billing/usage (trace counts). Uses REST API only."""

import copy
import json
import urllib.parse
import urllib.request
from typing import Any

_DEFAULT_ENDPOINT = "https://api.smith.langchain.com"


def _request(
    api_key: str,
    endpoint: str,
    path: str,
    params: dict[str, str] | None = None,
) -> dict[str, Any] | list[Any]:
    """GET request to LangSmith API. Returns JSON (dict or list)."""
    base = (endpoint or _DEFAULT_ENDPOINT).rstrip("/")
    url = f"{base}{path}"
    if params:
        url += "?" + urllib.parse.urlencode(params)
    req = urllib.request.Request(url, method="GET")
    req.add_header("Accept", "application/json")
    req.add_header("X-API-Key", api_key)
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        body = e.read().decode() if e.fp else ""
        return {"error": f"HTTP {e.code}: {body}"}
    except OSError as e:
        return {"error": str(e)}


def _list_workspaces(api_key: str, endpoint: str) -> dict[str, Any] | list[Any]:
    """GET /api/v1/workspaces."""
    return _request(api_key, endpoint, "/api/v1/workspaces")


def _get_workspace_by_id(api_key: str, endpoint: str, workspace_id: str) -> dict[str, Any]:
    """GET /api/v1/workspaces/{id}."""
    out = _request(api_key, endpoint, f"/api/v1/workspaces/{workspace_id}")
    return out if isinstance(out, dict) else {"error": "Unexpected response"}


def _build_workspace_id_to_name(
    api_key: str,
    endpoint: str,
    single_workspace: str | None,
) -> dict[str, str]:
    """Build workspace_id -> name. If single_workspace set, fetch only that one."""
    id_to_name: dict[str, str] = {}
    if single_workspace:
        single_workspace = single_workspace.strip()
        if len(single_workspace) == 36 and single_workspace.count("-") == 4:
            resp = _get_workspace_by_id(api_key, endpoint, single_workspace)
            if isinstance(resp, dict) and "error" not in resp and resp.get("id"):
                name = resp.get("display_name") or resp.get("name") or single_workspace
                id_to_name[str(resp["id"])] = name
                return id_to_name
        ws_resp = _list_workspaces(api_key, endpoint)
        workspaces: list[dict] = []
        if isinstance(ws_resp, list):
            workspaces = [w for w in ws_resp if isinstance(w, dict)]
        elif isinstance(ws_resp, dict) and "error" not in ws_resp:
            workspaces = list(ws_resp.get("workspaces") or ws_resp.get("items") or [])
        single_lower = single_workspace.lower()
        for w in workspaces:
            if not w.get("id"):
                continue
            wid = str(w["id"])
            name = w.get("display_name") or w.get("name") or wid
            if wid == single_workspace or (name or "").lower() == single_lower:
                id_to_name[wid] = name
                return id_to_name
        return id_to_name
    ws_resp = _list_workspaces(api_key, endpoint)
    if isinstance(ws_resp, list):
        for w in ws_resp:
            if isinstance(w, dict) and w.get("id"):
                id_to_name[str(w["id"])] = w.get("display_name") or w.get("name") or str(w["id"])
    elif isinstance(ws_resp, dict) and "error" not in ws_resp:
        for w in ws_resp.get("workspaces") or ws_resp.get("items") or []:
            if isinstance(w, dict) and w.get("id"):
                id_to_name[str(w["id"])] = w.get("display_name") or w.get("name") or str(w["id"])
    return id_to_name


def _augment_usage_groups_with_names(
    usage: list[dict],
    workspace_id_to_name: dict[str, str],
    only_workspace_id: str | None = None,
) -> list[dict]:
    """Put workspace_name next to each group value; optionally filter to one workspace.

    All workspace UUIDs from the billing API are preserved. If a UUID is not in
    workspace_id_to_name (e.g. workspace not in /api/v1/workspaces list), it is
    still included with workspace_name set to the raw UUID.
    """
    result = copy.deepcopy(usage)
    for item in result:
        groups = (item or {}).get("groups")
        if not isinstance(groups, dict):
            continue
        new_groups: dict[str, Any] = {}
        for uid, val in groups.items():
            if only_workspace_id and uid != only_workspace_id:
                continue
            # Keep every UUID; use id-to-name map or fall back to UUID as name
            name = workspace_id_to_name.get(uid) or uid
            new_groups[uid] = {"workspace_name": name, "value": val}
        item["groups"] = new_groups
    return result


def get_billing_usage_tool(
    api_key: str,
    endpoint: str,
    starting_on: str,
    ending_before: str,
    on_current_plan: bool = True,
    workspace: str | None = None,
) -> dict[str, Any] | list[dict]:
    """
    Fetch org billing usage (trace counts) with workspace names inline.

    Uses GET /api/v1/orgs/current/billing/usage and optionally
    GET /api/v1/workspaces (or single workspace) to resolve names.
    When workspace is provided (UUID or name), only that workspace's
    entries are included in each metric's groups.

    Args:
        api_key: LangSmith API key.
        endpoint: API base URL (e.g. https://api.smith.langchain.com).
        starting_on: Start of range (ISO 8601).
        ending_before: End of range (ISO 8601).
        on_current_plan: If true, only usage on current plan.
        workspace: Optional single workspace UUID or name to filter to.

    Returns:
        List of billing metrics with groups as
        { uuid: { "workspace_name": "<name>", "value": <number> } },
        or error dict with "error" key.
    """
    params = {
        "starting_on": starting_on,
        "ending_before": ending_before,
        "on_current_plan": "true" if on_current_plan else "false",
    }
    raw = _request(api_key, endpoint, "/api/v1/orgs/current/billing/usage", params)
    if isinstance(raw, dict) and "error" in raw:
        return raw
    if not isinstance(raw, list) or not raw:
        return {"error": "Unexpected billing usage response"}
    workspace_id_to_name = _build_workspace_id_to_name(api_key, endpoint, workspace)
    only_workspace_id: str | None = None
    if workspace and workspace_id_to_name:
        only_workspace_id = next(iter(workspace_id_to_name.keys()), None)
    return _augment_usage_groups_with_names(raw, workspace_id_to_name, only_workspace_id)
