iri_client.py

"""
IRI Facility API Client
=======================

Handles HTTP communication with IRI Facility API, Bearer token
resolution (env var or per-call parameter), and async task polling
for filesystem operations.

Token priority (highest to lowest):
  1. ``iri_token`` argument passed directly to the tool call.
  2. ``IRI_BEARER_TOKEN`` environment variable.

If neither is set, a ToolError is raised with instructions.

Multi-site routing
------------------
Pass ``site=`` to ``make_iri_request`` (or ``get_base_url``) to target a
specific IRI deployment.  The default site is controlled by the ``IRI_SITE``
env var (default: ``esnet-east``).  ``IRI_BASE_URL`` overrides everything when
set (legacy compatibility for single-site deployments).

Known sites
~~~~~~~~~~~
  nersc       — https://api.iri.nersc.gov/api/v1
  esnet-east  — https://iri-dev.ppg.es.net/api/v1
  esnet-west  — https://esnet-west.sdn-sense.net/api/v1
  alcf        — https://api.alcf.anl.gov/api/v1

Capability sets per site
~~~~~~~~~~~~~~~~~~~~~~~~~
  nersc, esnet-east, esnet-west: facility, status, account, compute, filesystem
  alcf:                          facility, status  (no account/compute/filesystem)
"""
import os
import asyncio
import logging
from typing import Any, Dict, Optional, Set

import httpx
from fastmcp.exceptions import ToolError

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Site registry
# ---------------------------------------------------------------------------

SITES: Dict[str, str] = {
    "nersc": "https://api.iri.nersc.gov/api/v1",
    "esnet-east": "https://iri-dev.ppg.es.net/api/v1",
    "esnet-west": "https://esnet-west.sdn-sense.net/api/v1",
    "alcf": "https://api.alcf.anl.gov/api/v1",
}

# Capabilities each site is known to support
SITE_CAPABILITIES: Dict[str, Set[str]] = {
    "nersc":      {"facility", "status", "account", "compute", "filesystem"},
    "esnet-east": {"facility", "status", "account", "compute", "filesystem"},
    "esnet-west": {"facility", "status", "account", "compute", "filesystem"},
    "alcf":       {"facility", "status"},
}

_DEFAULT_SITE = os.environ.get("IRI_SITE", "esnet-east")

# Legacy single-site override — takes priority when set
_IRI_BASE_URL_OVERRIDE: Optional[str] = os.environ.get("IRI_BASE_URL", "").strip() or None


# ---------------------------------------------------------------------------
# URL / capability helpers
# ---------------------------------------------------------------------------


def get_base_url(site: Optional[str] = None) -> str:
    """
    Resolve the IRI base URL for a given site name.

    Priority:
      1. ``site`` argument (explicit per-call routing).
      2. ``IRI_SITE`` environment variable (session default).
      3. ``IRI_BASE_URL`` environment variable (legacy override; skips site lookup).

    Raises:
        ToolError: If ``site`` is not in the known site registry.
    """
    if site:
        key = site.lower()
        if key not in SITES:
            raise ToolError(
                f"Unknown site '{site}'. Valid sites: {', '.join(sorted(SITES))}."
            )
        return SITES[key]

    # Legacy: explicit IRI_BASE_URL takes priority over IRI_SITE
    if _IRI_BASE_URL_OVERRIDE:
        return _IRI_BASE_URL_OVERRIDE.rstrip("/")

    key = _DEFAULT_SITE.lower()
    if key not in SITES:
        raise ToolError(
            f"IRI_SITE='{_DEFAULT_SITE}' is not a known site. "
            f"Valid sites: {', '.join(sorted(SITES))}."
        )
    return SITES[key]


def check_capability(site: Optional[str], capability: str) -> None:
    """
    Raise a ToolError if the resolved site is known NOT to support ``capability``.

    Does nothing if the site is unknown (gives the API a chance to respond).
    """
    if _IRI_BASE_URL_OVERRIDE and not site:
        return  # legacy URL — no capability info
    key = (site or _DEFAULT_SITE).lower()
    caps = SITE_CAPABILITIES.get(key)
    if caps is not None and capability not in caps:
        raise ToolError(
            f"Site '{key}' does not support {capability} operations. "
            f"Sites with {capability}: "
            + ", ".join(s for s, c in SITE_CAPABILITIES.items() if capability in c)
            + "."
        )


# ---------------------------------------------------------------------------
# Token resolution
# ---------------------------------------------------------------------------


def get_token(iri_token: Optional[str] = None) -> str:
    """
    Resolve the IRI Bearer token.

    Checks (in order):
      1. The ``iri_token`` argument (passed per tool call).
      2. The ``IRI_BEARER_TOKEN`` environment variable.

    Raises:
        ToolError: If no token is available.
    """
    token = iri_token or os.environ.get("IRI_BEARER_TOKEN", "").strip()
    if not token:
        raise ToolError(
            "No IRI Bearer token found. Provide it via:\n"
            "  • Set the IRI_BEARER_TOKEN environment variable, OR\n"
            "  • Pass iri_token='<your-token>' in the tool call.\n\n"
            "You can obtain a Globus token via:\n"
            "  globus-automate-flow run ... --label test\n"
            "or from your facility's token endpoint."
        )
    return token


# ---------------------------------------------------------------------------
# HTTP client
# ---------------------------------------------------------------------------


async def make_iri_request(
    method: str,
    path: str,
    token: str,
    *,
    site: Optional[str] = None,
    params: Optional[Dict[str, Any]] = None,
    json: Optional[Any] = None,
    data: Optional[Any] = None,
    files: Optional[Any] = None,
    timeout: float = 30.0,
) -> Any:
    """
    Make an authenticated request to an IRI Facility API deployment.

    Args:
        method:  HTTP method (GET, POST, PUT, DELETE).
        path:    API path, e.g. "/status/resources".
        token:   IRI Bearer token.
        site:    Site name (nersc, esnet-east, esnet-west, alcf).
                 Defaults to IRI_SITE env var, or "esnet-east".
        params:  Query parameters dict.
        json:    JSON body (mutually exclusive with data/files).
        data:    Form data.
        files:   Multipart files.
        timeout: Request timeout in seconds.

    Returns:
        Parsed JSON response (dict or list).

    Raises:
        ToolError: On HTTP errors or connection failures.
    """
    base_url = get_base_url(site)
    url = f"{base_url}{path}"
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept": "application/json",
    }

    logger.debug("%s %s  params=%s", method.upper(), url, params)

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.request(
                method,
                url,
                headers=headers,
                params=params,
                json=json,
                data=data,
                files=files,
            )

        if response.status_code == 401:
            raise ToolError(
                "IRI API returned 401 Unauthorized. "
                "Your Bearer token may have expired — update IRI_BEARER_TOKEN and try again."
            )
        if response.status_code == 403:
            raise ToolError(
                "IRI API returned 403 Forbidden. "
                "You may not have permission for this resource."
            )
        if response.status_code == 404:
            raise ToolError(f"IRI API returned 404 Not Found for: {path}")
        if response.status_code == 501:
            resolved = site or _DEFAULT_SITE
            raise ToolError(
                f"IRI API at site '{resolved}' returned 501 Not Implemented for: {path}. "
                "This endpoint may not be supported by that deployment."
            )

        if response.is_error:
            try:
                body = response.json()
            except Exception:
                body = response.text
            raise ToolError(
                f"IRI API error {response.status_code} for {path}: {body}"
            )

        # Return raw text for non-JSON responses (e.g. empty 204)
        if not response.content:
            return {}
        return response.json()

    except httpx.TimeoutException:
        raise ToolError(f"Request to IRI API timed out after {timeout}s: {url}")
    except httpx.ConnectError as exc:
        raise ToolError(f"Cannot connect to IRI API at {base_url}: {exc}")
    except ToolError:
        raise
    except Exception as exc:
        raise ToolError(f"IRI API request failed: {exc}")


# ---------------------------------------------------------------------------
# Async task polling (used by filesystem operations)
# ---------------------------------------------------------------------------


async def _get_task(task_uri: str, token: str, timeout: float = 30.0) -> Dict[str, Any]:
    """
    Fetch a task status using its absolute URI.

    IRI returns task_uri as a full URL (e.g. https://host/api/v1/task/task-123).
    We query it directly with httpx rather than going through make_iri_request,
    which always prepends a base URL.
    """
    headers = {"Authorization": f"Bearer {token}", "Accept": "application/json"}
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.get(task_uri, headers=headers)
        if response.status_code == 401:
            raise ToolError(
                "IRI API returned 401 Unauthorized polling task. "
                "Update IRI_BEARER_TOKEN and retry."
            )
        response.raise_for_status()
        return response.json()
    except httpx.TimeoutException:
        raise ToolError(f"Timed out polling IRI task: {task_uri}")
    except httpx.RequestError as exc:
        raise ToolError(f"Connection error polling IRI task: {exc}")
    except ToolError:
        raise
    except Exception as exc:
        raise ToolError(f"Failed to poll IRI task: {exc}")


async def poll_task(
    task_uri: str,
    token: str,
    timeout: int = 300,
    poll_interval: float = 2.0,
) -> Any:
    """
    Poll an IRI async task until it completes or times out.

    IRI filesystem operations return a TaskSubmitResponse with `task_id` and
    `task_uri` (absolute URL). This function polls `task_uri` every
    `poll_interval` seconds until status is terminal, then returns the
    task's `result` field (the actual operation output).

    Task status values (from IRI TaskStatus enum):
        pending, active, completed, failed, canceled

    Args:
        task_uri:      Absolute URL to the task, e.g.
                       "https://host/api/v1/task/task-123"
        token:         IRI Bearer token.
        timeout:       Maximum seconds to wait (default 300).
        poll_interval: Seconds between polls (default 2.0, matching IRI notebooks).

    Returns:
        The task's `result` field (dict/any) — the actual operation output.

    Raises:
        ToolError: On task failure, cancellation, or timeout.
    """
    elapsed = 0.0
    while elapsed < timeout:
        task = await _get_task(task_uri, token)
        status = task.get("status", "").lower()

        logger.debug("Task %s status: %s", task.get("id", "?"), status)

        if status == "completed":
            return task.get("result")

        if status in ("failed", "canceled"):
            task_id = task.get("id", "?")
            raise ToolError(
                f"IRI task {task_id} ended with status '{status}'. "
                f"Details: {task}"
            )

        await asyncio.sleep(poll_interval)
        elapsed += poll_interval

    raise ToolError(
        f"IRI task timed out after {timeout}s ({task_uri}). "
        "The operation may still be running — check the facility dashboard."
    )