"""
API client for HLA-Compass platform

This module provides the APIClient class that handles all communication
with the HLA-Compass REST API. It's used internally by the data access
classes (PeptideData, ProteinData, SampleData) to fetch scientific data.
"""

from __future__ import annotations

import json
import logging
import os
import platform
import re
import time
import uuid
from typing import Any, Callable, Dict, Iterator, List, Optional

import requests

from .auth import Auth
from .config import Config
from .utils import RateLimiter, parse_api_error

logger = logging.getLogger(__name__)

_COMPLETED_RUN_STATUSES = {"completed"}
_FAILED_RUN_STATUSES = {"failed", "cancelled", "cancel_requested", "timed_out", "timeout"}
_HLA_BINDING_METHODS = {"netmhcpan", "mhcflurry", "ensemble"}
_HLA_PEPTIDE_RE = re.compile(r"^[ACDEFGHIKLMNPQRSTVWY]+$")


class APIError(Exception):
    """API request error"""

    def __init__(
        self,
        message: str,
        status_code: Optional[int] = None,
        details: Optional[Dict] = None,
    ):
        super().__init__(message)
        self.status_code = status_code
        self.details = details or {}


class APIClient:
    """Client for interacting with the HLA-Compass API.

    Highlights
    ---------
    • **Iterators** – Use :meth:`iter_peptides`, :meth:`iter_proteins`, etc. to stream
      large result sets without manual paging::

          client = APIClient()
          for peptide in client.iter_peptides({"sequence": "SIINFEKL"}):
              print(peptide["sequence"], peptide["mass"])

    • **Rate limits** – Override defaults via environment variables
      ``HLA_RATE_LIMIT_MAX_REQUESTS`` and ``HLA_RATE_LIMIT_TIME_WINDOW``; inspect the
      active settings with :meth:`get_rate_limit_state`.

    • **Telemetry** – Attach a callback to capture structured metrics for every
      request (status, duration, retries)::

          def log_request(data: Dict[str, Any]) -> None:
              print("API", data["status_code"], data["duration_seconds"], data["url"])

          client = APIClient()
          client.set_telemetry_callback(log_request)
          client.get_peptide("pep_123")

    """

    def __init__(self, provider: Optional[str] = None, catalog: Optional[str] = None):
        """Initialize API client with authentication and a persistent session

        Args:
            provider: Data provider name (default: from config or 'alithea-bio')
            catalog: Data catalog name (default: from config or 'immunopeptidomics')
        """
        self.auth = Auth()
        self.config = Config()
        self.base_url = self.config.get_api_endpoint()

        # Set provider and catalog from args, config, or defaults
        self.provider = provider or self.config.get("data_provider", "alithea-bio")
        self.catalog = catalog or self.config.get("data_catalog", "immunopeptidomics")
        # Facades for data access (SQL + storage) via the same API client
        from .data import DataClient

        self.data = DataClient(provider=self.provider, catalog=self.catalog, api_client=self)
        self.catalog_storage = self.data.catalog_storage
        # Convenience facades for module + run operations
        self.modules = ModulesFacade(self)
        self.runs = RunsFacade(self)
        self.dataset_versions = DatasetVersionsFacade(self)
        self._scientific_query = None
        self._scientific_query_checked = False

        rl_settings = Config.get_rate_limit_settings()
        max_requests = rl_settings.get("max_requests", 100)
        time_window = rl_settings.get("time_window", 60)

        self.rate_limiter = RateLimiter(
            max_requests=max_requests,
            time_window=time_window,
        )
        circuit_cfg = Config.get_circuit_breaker_settings()
        self._cb_threshold = circuit_cfg.get("threshold", 5)
        self._cb_reset = circuit_cfg.get("reset_seconds", 60.0)
        self._cb_failures = 0
        self._cb_opened_at: Optional[float] = None
        self._rate_limit_state: Dict[str, Any] = {
            "max_requests": max_requests,
            "time_window": time_window,
            "last_wait_seconds": 0.0,
        }
        self._telemetry_callback: Optional[Callable[[Dict[str, Any]], None]] = None

        # Create a persistent HTTP session for connection reuse and default headers
        self.session = requests.Session()

        # Build a descriptive User-Agent
        try:
            from . import (
                __version__ as SDK_VERSION,  # Avoid hard dependency if import path changes
            )
        except Exception:
            SDK_VERSION = "unknown"

        ua = (
            f"hla-compass-sdk/{SDK_VERSION} "
            f"python/{platform.python_version()} "
            f"os/{platform.system()}-{platform.release()}"
        )

        self.session.headers.update(
            {
                "Accept": "application/json",
                "User-Agent": ua,
            }
        )

        # Configure retries for idempotent requests (GET/HEAD/OPTIONS)
        try:  # Optional dependency present in requests.adapters
            from requests.adapters import HTTPAdapter  # type: ignore
            from urllib3.util import Retry  # type: ignore

            retry = Retry(
                total=3,
                connect=3,
                read=3,
                backoff_factor=0.5,
                status_forcelist=[429, 500, 502, 503, 504],
                allowed_methods={"GET", "HEAD", "OPTIONS"},
                respect_retry_after_header=True,
            )
            adapter = HTTPAdapter(max_retries=retry)
            self.session.mount("https://", adapter)
            self.session.mount("http://", adapter)
        except Exception:
            # If retry configuration fails, continue without it
            pass

    def _headers(self) -> Dict[str, str]:
        """Compose request headers by combining session defaults with auth headers"""
        headers = dict(self.session.headers)
        headers.update(self.auth.get_headers())
        # Optional correlation id for cross-service tracing
        try:
            corr = self.config.get_correlation_id()
            if corr:
                headers["X-Correlation-Id"] = corr
        except Exception:
            pass
        return headers

    def _build_data_endpoint(self, entity: str, entity_id: Optional[str] = None) -> str:
        """Build a data API endpoint path.

        Args:
            entity: Entity name (e.g., 'peptides', 'proteins')
            entity_id: Optional entity ID for specific resource

        Returns:
            Complete endpoint path
        """
        base = f"/v1/data/{self.provider}/{self.catalog}/{entity}"
        if entity_id:
            return f"{base}/{entity_id}"
        return base

    def _get_scientific_query(self):
        if self._scientific_query_checked:
            return self._scientific_query

        self._scientific_query_checked = True
        cluster_arn = os.environ.get("DB_CLUSTER_ARN")
        secret_arn = os.environ.get("DB_APP_SECRET_ARN")
        organization_id = self.config.get("active_org_id")

        if not (cluster_arn and secret_arn and organization_id):
            return None

        try:
            from .database import ScientificQuery

            self._scientific_query = ScientificQuery(
                cluster_arn=cluster_arn,
                secret_arn=secret_arn,
                organization_id=str(organization_id),
            )
        except Exception as exc:  # pragma: no cover - best effort fallback
            logger.debug("ScientificQuery fallback unavailable: %s", exc)
            self._scientific_query = None

        return self._scientific_query

    def _uses_api_key_auth(self) -> bool:
        try:
            headers = self.auth.get_headers()
        except Exception:
            return False
        return any(str(key).lower() == "x-api-key" for key in (headers or {}).keys())

    def _execute_scientific_rows(self, sql: str, params: Optional[List[Any]] = None) -> List[Dict[str, Any]]:
        scientific_query = self._get_scientific_query()
        if scientific_query is not None:
            result = scientific_query.execute_readonly(sql, params or [])
            if isinstance(result, list):
                return [row for row in result if isinstance(row, dict)]
            raise APIError("Unexpected scientific query result format")

        if self._uses_api_key_auth():
            result = self.data.sql.query(sql, params=params)
            if isinstance(result, dict):
                rows = result.get("data", [])
                if isinstance(rows, list):
                    return [row for row in rows if isinstance(row, dict)]
            raise APIError("Unexpected scientific query result format")

        raise APIError(
            "This scientific helper requires either direct database access "
            "(DB_CLUSTER_ARN + DB_APP_SECRET_ARN + active organization) "
            "or API-key authentication. User-session JWT access does not yet expose "
            "a dedicated backend route for this operation."
        )

    @staticmethod
    def _normalize_hla_locus(locus: Optional[str]) -> Optional[str]:
        if locus is None:
            return None
        normalized = str(locus).strip().upper()
        if not normalized:
            return None
        if not normalized.startswith("HLA-"):
            normalized = f"HLA-{normalized}"
        return normalized

    @staticmethod
    def _normalize_hla_resolution(resolution: str) -> str:
        normalized = str(resolution or "2-digit").strip().lower()
        if normalized not in {"2-digit", "4-digit", "full"}:
            raise APIError("Unsupported HLA resolution. Use one of: '2-digit', '4-digit', 'full'.")
        return normalized

    @staticmethod
    def _format_hla_allele_resolution(allele_name: str, resolution: str) -> str:
        if resolution == "full" or "*" not in allele_name:
            return allele_name

        prefix, suffix = allele_name.split("*", 1)
        segments = [segment for segment in suffix.split(":") if segment]
        if not segments:
            return allele_name
        if resolution == "2-digit":
            return f"{prefix}*{segments[0]}"
        return f"{prefix}*{':'.join(segments[:2])}" if len(segments) >= 2 else f"{prefix}*{segments[0]}"

    @staticmethod
    def _extract_hla_predictions(result_payload: Dict[str, Any]) -> List[Dict[str, Any]]:
        if isinstance(result_payload, dict):
            data = result_payload.get("data")
            if isinstance(data, dict) and isinstance(data.get("predictions"), list):
                return [item for item in data["predictions"] if isinstance(item, dict)]
            if isinstance(data, list):
                return [item for item in data if isinstance(item, dict)]
            if isinstance(result_payload.get("predictions"), list):
                return [item for item in result_payload["predictions"] if isinstance(item, dict)]
        raise APIError("Binding prediction run completed without a usable 'predictions' payload")

    @staticmethod
    def _normalize_module_run_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
        normalized = dict(payload)

        run_id = normalized.get("run_id") or normalized.get("runId") or normalized.get("id")
        if run_id:
            normalized.setdefault("run_id", str(run_id))

        module_id = normalized.get("module_id") or normalized.get("moduleId")
        module_block = normalized.get("module")
        if not module_id and isinstance(module_block, dict):
            module_id = module_block.get("id")
        if module_id:
            normalized.setdefault("module_id", str(module_id))

        return normalized

    def _resolve_hla_binding_module_id(self, module_id: Optional[str] = None) -> Optional[str]:
        if module_id:
            return str(module_id)

        for candidate in (
            os.environ.get("HLA_BINDING_MODULE_ID"),
            self.config.get("hla_binding_module_id"),
            self.config.get("hlaBindingModuleId"),
        ):
            if candidate:
                return str(candidate)
        return None

    def _make_request(
        self,
        method: str,
        endpoint: str,
        params: Dict = None,
        json_data: Dict = None,
        max_retries: int = 3,
        idempotency_key: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Make an authenticated API request with retries and timeouts

        Args:
            method: HTTP method (GET, POST, etc.)
            endpoint: API endpoint path (e.g., /v1/data/alithea-bio/immunopeptidomics/peptides)
            params: Query parameters
            json_data: JSON body data
            max_retries: Maximum number of retry attempts for transient errors

        Returns:
            Parsed JSON response

        Raises:
            APIError: If request fails
        """
        # Ensure we have authentication
        if not self.auth.is_authenticated():
            raise APIError("Not authenticated. Please run 'hla-compass auth login' first")

        # Circuit breaker guard
        now = time.time()
        if self._cb_opened_at and (now - self._cb_opened_at) < self._cb_reset:
            raise APIError("Circuit open due to repeated failures")
        if self._cb_opened_at and (now - self._cb_opened_at) >= self._cb_reset:
            self._cb_opened_at = None
            self._cb_failures = 0

        # Build full URL
        url = f"{self.base_url}{endpoint}"

        # Apply rate limiting
        self.rate_limiter.acquire()
        self._rate_limit_state["last_wait_seconds"] = self.rate_limiter.last_wait

        # Get combined headers (session defaults + auth)
        headers = self._headers()

        # Retry logic for transient errors
        # For POST requests, only allow retry when an idempotency key is provided.
        # Generate a stable idempotency key per call if not provided so that
        # internal retries are safe and deduplicated server-side.
        method_upper = (method or "").upper()
        if method_upper == "POST" and not idempotency_key:
            idempotency_key = str(uuid.uuid4())
        can_retry_post = method_upper != "POST" or (idempotency_key is not None)

        for attempt in range(max_retries):
            try:
                start_time = time.perf_counter()
                # Create a fresh set of headers per attempt and add a unique request id
                attempt_headers = dict(headers)
                request_id = str(uuid.uuid4())
                attempt_headers["X-Request-Id"] = request_id
                # Add idempotency key for POSTs to make retries safe
                if method_upper == "POST" and idempotency_key:
                    attempt_headers.setdefault("Idempotency-Key", idempotency_key)
                # Make request with configurable timeout (defaults: 5s connect, 30s read)
                timeouts = self.config.get_request_timeouts()
                response = self.session.request(
                    method=method,
                    url=url,
                    headers=attempt_headers,
                    params=params,
                    json=json_data,
                    timeout=(timeouts.get("connect", 5.0), timeouts.get("read", 30.0)),
                )
                duration = time.perf_counter() - start_time

                # Handle 401 - try to refresh token once
                if response.status_code == 401 and attempt == 0:
                    logger.info("Token expired, attempting refresh")
                    new_token = self.auth.refresh_token()
                    if new_token:
                        # Rebuild base headers to include new Authorization and retry
                        headers = self._headers()
                        continue  # Retry with new token
                    else:
                        raise APIError("Authentication expired. Please run 'hla-compass auth login' again")

                # Handle rate limiting with exponential backoff
                if response.status_code == 429:
                    if attempt < max_retries - 1:
                        retry_after = int(response.headers.get("Retry-After", 2**attempt))
                        logger.warning(f"Rate limited, retrying after {retry_after} seconds")
                        time.sleep(retry_after)
                        continue
                    else:
                        self._record_failure()
                        raise APIError("Rate limit exceeded. Please try again later.", 429)

                # Handle server errors with retry
                if response.status_code >= 500 and attempt < max_retries - 1:
                    # Only retry POSTs when idempotency is guaranteed
                    if method_upper == "POST" and not can_retry_post:
                        raise APIError(f"Server error {response.status_code}", response.status_code)
                    wait_time = 2**attempt  # Exponential backoff
                    logger.warning(f"Server error {response.status_code}, retrying in {wait_time} seconds")
                    time.sleep(wait_time)
                    continue

                # Check for other errors
                if response.status_code >= 400:
                    # Sanitize error message to avoid info disclosure
                    error_msg = parse_api_error(response, "API request failed")
                    # Include correlation ID in error message for debugging
                    error_msg_with_id = f"{error_msg} (request_id={request_id})"
                    error_details = None
                    try:
                        error_details = response.json()
                    except Exception:  # pragma: no cover - defensive
                        error_details = None
                    self._emit_telemetry(
                        method=method_upper,
                        url=url,
                        status_code=response.status_code,
                        duration=duration,
                        attempt=attempt,
                        retries=max_retries,
                        params=params,
                        idempotency_key=idempotency_key,
                        request_id=request_id,
                        error=error_msg,
                    )
                    self._record_failure()
                    raise APIError(error_msg_with_id, response.status_code, details=error_details)

                # Parse successful response (204 = empty body)
                if response.status_code == 204:
                    data = {}
                    self._record_success()
                else:
                    try:
                        data = response.json()
                    except ValueError as exc:
                        # Non-JSON response body
                        if not response.text:
                            data = {}
                        else:
                            error_msg = parse_api_error(response, "Invalid JSON response from API")
                            raise APIError(error_msg, response.status_code) from exc

                # Handle success wrapper format
                payload = data.get("data", data) if isinstance(data, dict) and data.get("success") else data

                self._emit_telemetry(
                    method=method_upper,
                    url=url,
                    status_code=response.status_code,
                    duration=duration,
                    attempt=attempt,
                    retries=max_retries,
                    params=params,
                    idempotency_key=idempotency_key,
                    request_id=request_id,
                )

                self._record_success()
                return payload

            except requests.Timeout as e:
                duration = time.perf_counter() - start_time
                self._emit_telemetry(
                    method=method_upper,
                    url=url,
                    status_code=None,
                    duration=duration,
                    attempt=attempt,
                    retries=max_retries,
                    params=params,
                    idempotency_key=idempotency_key,
                    request_id=request_id,
                    error="timeout",
                )

                if attempt < max_retries - 1:
                    wait_time = 2**attempt
                    logger.warning(f"Request timeout, retrying in {wait_time} seconds")
                    time.sleep(wait_time)
                    continue
                else:
                    self._record_failure()
                    raise APIError("Request timed out. Please check your connection and try again.") from e

            except requests.RequestException as e:
                duration = time.perf_counter() - start_time
                self._emit_telemetry(
                    method=method_upper,
                    url=url,
                    status_code=None,
                    duration=duration,
                    attempt=attempt,
                    retries=max_retries,
                    params=params,
                    idempotency_key=idempotency_key,
                    request_id=request_id,
                    error=str(e),
                )

                if attempt < max_retries - 1 and "connection" in str(e).lower():
                    wait_time = 2**attempt
                    logger.warning(f"Connection error, retrying in {wait_time} seconds")
                    time.sleep(wait_time)
                    continue
                else:
                    self._record_failure()
                    raise APIError(f"Network error: {str(e)}") from e

    def _iterate_paginated(
        self,
        fetcher: Callable[[int, int], List[Dict[str, Any]]],
        page_size: int,
        max_results: Optional[int],
    ) -> Iterator[Dict[str, Any]]:
        offset = 0
        yielded = 0

        while True:
            if max_results is not None:
                remaining = max_results - yielded
                if remaining <= 0:
                    break
                limit = min(page_size, remaining)
            else:
                limit = page_size

            page = fetcher(limit, offset)
            if not page:
                break

            for item in page:
                yield item
                yielded += 1
                if max_results is not None and yielded >= max_results:
                    return

            if len(page) < limit:
                break

            offset += limit

    # Peptide endpoints

    def get_peptides(self, filters: Dict = None, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """
        Search peptides with filters

        Args:
            filters: Search filters (sequence, min_length, max_length, etc.)
            limit: Maximum results to return
            offset: Pagination offset

        Returns:
            List of peptide records
        """
        params = {"limit": limit, "offset": offset}

        if filters:
            params.update(filters)

        result = self._make_request("GET", self._build_data_endpoint("peptides"), params=params)

        # Extract peptides from response
        if isinstance(result, dict):
            return result.get("peptides", result.get("items", []))
        return result if isinstance(result, list) else []

    def get_peptide(self, peptide_id: str) -> Dict[str, Any]:
        """Get a single peptide record.

        The data API exposes list-style GET endpoints (with filters) for core
        entities. Historically the SDK also offered ``/peptides/{id}`` lookups,
        but those paths are not consistently available. To keep this call useful
        across deployments, we translate the request into a filtered list query.

        Args:
            peptide_id: A numeric peptide id (peptide_id) or a sequence string.
        """
        value = str(peptide_id)
        if value.isdigit():
            items = self.get_peptides(filters={"peptide_ids": value}, limit=1, offset=0)
        else:
            items = self.get_peptides(filters={"sequence": value}, limit=1, offset=0)

        if items:
            return items[0]
        raise APIError("Not Found", status_code=404)

    def get_peptide_samples(self, peptide_id: str, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """List samples associated with a peptide id."""
        value = str(peptide_id)
        return self.get_samples(filters={"peptide_ids": value}, limit=limit, offset=offset)

    def get_peptide_proteins(self, peptide_id: str, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """List proteins associated with a peptide id."""
        value = str(peptide_id)
        return self.get_proteins(filters={"peptide_ids": value}, limit=limit, offset=offset)

    def iter_peptides(
        self,
        filters: Dict = None,
        page_size: int = 100,
        max_results: Optional[int] = None,
    ) -> Iterator[Dict[str, Any]]:
        """Yield peptides across pages without manual pagination."""

        def fetch(limit: int, offset: int) -> List[Dict[str, Any]]:
            return self.get_peptides(filters=filters, limit=limit, offset=offset)

        yield from self._iterate_paginated(fetch, page_size, max_results)

    # Protein endpoints

    def get_proteins(self, filters: Dict = None, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """
        Search proteins with filters

        Args:
            filters: Search filters (accession, gene_name, organism, etc.)
            limit: Maximum results to return
            offset: Pagination offset

        Returns:
            List of protein records
        """
        params = {"limit": limit, "offset": offset}

        if filters:
            params.update(filters)

        result = self._make_request("GET", self._build_data_endpoint("proteins"), params=params)

        # Extract proteins from response
        if isinstance(result, dict):
            return result.get("proteins", result.get("items", []))
        return result if isinstance(result, list) else []

    def get_protein(self, protein_id: str) -> Dict[str, Any]:
        """Get a single protein record.

        Args:
            protein_id: A numeric protein id (protein_id) or an accession.
        """
        value = str(protein_id)
        if value.isdigit():
            items = self.get_proteins(filters={"protein_ids": value}, limit=1, offset=0)
        else:
            items = self.get_proteins(filters={"accession": value}, limit=1, offset=0)

        if items:
            return items[0]
        raise APIError("Not Found", status_code=404)

    def get_protein_peptides(self, protein_id: str, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """List peptides associated with a protein id."""
        value = str(protein_id)
        return self.get_peptides(filters={"protein_ids": value}, limit=limit, offset=offset)

    def get_protein_coverage(self, protein_id: str) -> Dict[str, Any]:
        """Compute protein coverage from peptide mapping data."""
        value = str(protein_id).strip()
        if not value:
            raise APIError("protein_id is required")

        if value.isdigit():
            protein_rows = self._execute_scientific_rows(
                """
                SELECT protein_id, accession, protein_name, gene_name,
                       CHAR_LENGTH(COALESCE(full_sequence, '')) AS total_residues
                  FROM proteins
                 WHERE protein_id = %s
                 LIMIT 1
                """,
                [int(value)],
            )
        else:
            protein_rows = self._execute_scientific_rows(
                """
                SELECT protein_id, accession, protein_name, gene_name,
                       CHAR_LENGTH(COALESCE(full_sequence, '')) AS total_residues
                  FROM proteins
                 WHERE accession = %s
                 LIMIT 1
                """,
                [value],
            )

        if not protein_rows:
            raise APIError("Not Found", status_code=404)

        protein = protein_rows[0]
        protein_pk = protein.get("protein_id")
        total_residues = int(protein.get("total_residues") or 0)
        mapping_rows = self._execute_scientific_rows(
            """
            SELECT p.sequence, ppm.start_position, ppm.end_position
              FROM protein_peptide_mappings ppm
              JOIN peptides p ON p.peptide_id = ppm.peptide_id
             WHERE ppm.protein_id = %s
               AND ppm.start_position IS NOT NULL
               AND ppm.end_position IS NOT NULL
             ORDER BY ppm.start_position ASC, ppm.end_position ASC
            """,
            [protein_pk],
        )

        covered_positions: set[int] = set()
        coverage_counts: Dict[int, int] = {}
        peptide_lengths: List[int] = []
        for row in mapping_rows:
            start = row.get("start_position")
            end = row.get("end_position")
            sequence = str(row.get("sequence") or "")
            if start is None or end is None:
                continue
            try:
                start_i = int(start)
                end_i = int(end)
            except (TypeError, ValueError):
                continue
            if start_i <= 0 or end_i < start_i:
                continue
            peptide_lengths.append(len(sequence))
            for position in range(start_i, end_i + 1):
                covered_positions.add(position)
                coverage_counts[position] = coverage_counts.get(position, 0) + 1

        covered_residues = len(covered_positions)
        percent_covered = round((covered_residues / total_residues) * 100, 2) if total_residues else 0.0

        gaps: List[Dict[str, int]] = []
        gap_start: Optional[int] = None
        if total_residues > 0:
            for position in range(1, total_residues + 1):
                if position not in covered_positions:
                    if gap_start is None:
                        gap_start = position
                    continue
                if gap_start is not None:
                    gaps.append(
                        {
                            "start": gap_start,
                            "end": position - 1,
                            "length": position - gap_start,
                        }
                    )
                    gap_start = None
            if gap_start is not None:
                gaps.append(
                    {
                        "start": gap_start,
                        "end": total_residues,
                        "length": total_residues - gap_start + 1,
                    }
                )

        coverage_map = [
            {
                "position": position,
                "covered": position in covered_positions,
                "peptideCount": coverage_counts.get(position, 0),
            }
            for position in range(1, min(total_residues, 100) + 1)
        ]

        average_peptide_length = round(sum(peptide_lengths) / len(peptide_lengths), 2) if peptide_lengths else 0.0

        return {
            "proteinId": str(protein_pk),
            "proteinAccession": protein.get("accession"),
            "proteinName": protein.get("protein_name"),
            "geneName": protein.get("gene_name"),
            "coverageStatistics": {
                "percentCovered": percent_covered,
                "coveredResidues": covered_residues,
                "totalResidues": total_residues,
                "peptideCount": len(mapping_rows),
                "averagePeptideLength": average_peptide_length,
            },
            "gaps": gaps,
            "coverageMap": coverage_map,
        }

    def iter_proteins(
        self,
        filters: Dict = None,
        page_size: int = 100,
        max_results: Optional[int] = None,
    ) -> Iterator[Dict[str, Any]]:
        """Yield proteins across pages without manual pagination."""

        def fetch(limit: int, offset: int) -> List[Dict[str, Any]]:
            return self.get_proteins(filters=filters, limit=limit, offset=offset)

        yield from self._iterate_paginated(fetch, page_size, max_results)

    # Sample endpoints

    def get_samples(self, filters: Dict = None, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """
        Search samples with filters

        Args:
            filters: Search filters (tissue, disease, cell_line, etc.)
            limit: Maximum results to return
            offset: Pagination offset

        Returns:
            List of sample records
        """
        params = {"limit": limit, "offset": offset}

        if filters:
            params.update(filters)

        result = self._make_request("GET", self._build_data_endpoint("samples"), params=params)

        # Extract samples from response
        if isinstance(result, dict):
            return result.get("samples", result.get("items", []))
        return result if isinstance(result, list) else []

    def get_sample(self, sample_id: str) -> Dict[str, Any]:
        """Get a single sample record.

        Args:
            sample_id: A numeric sample id (sample_id).
        """
        value = str(sample_id)
        if not value.isdigit():
            raise APIError(
                "Sample lookup by name is not supported by the public API. "
                "Use get_samples(filters={'search': ...}) or pass a numeric sample id.",
                status_code=400,
            )

        items = self.get_samples(filters={"sample_ids": value}, limit=1, offset=0)

        if items:
            return items[0]
        raise APIError("Not Found", status_code=404)

    def get_sample_peptides(self, sample_id: str, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """List peptides observed in a sample."""
        value = str(sample_id)
        return self.get_peptides(filters={"sample_ids": value}, limit=limit, offset=offset)

    def iter_samples(
        self,
        filters: Dict = None,
        page_size: int = 100,
        max_results: Optional[int] = None,
    ) -> Iterator[Dict[str, Any]]:
        """Yield samples across pages without manual pagination."""

        def fetch(limit: int, offset: int) -> List[Dict[str, Any]]:
            return self.get_samples(filters=filters, limit=limit, offset=offset)

        yield from self._iterate_paginated(fetch, page_size, max_results)

    # HLA endpoints

    def get_hla_alleles(self, locus: Optional[str] = None, resolution: str = "2-digit") -> List[str]:
        """Get HLA alleles through scientific-query capable auth paths.

        This helper works when the client can execute scientific queries:
        module-runtime DB access or API-key backed scientific-query access.
        Ordinary JWT-only user sessions do not currently expose a dedicated
        REST endpoint for this data.
        """
        normalized_locus = self._normalize_hla_locus(locus)
        normalized_resolution = self._normalize_hla_resolution(resolution)

        sql = "SELECT allele_name FROM hla_alleles"
        params: List[Any] = []
        if normalized_locus:
            sql += " WHERE gene = %s"
            params.append(normalized_locus)
        sql += " ORDER BY allele_name ASC"

        rows = self._execute_scientific_rows(sql, params or None)

        alleles: List[str] = []
        seen: set[str] = set()
        for row in rows:
            allele_name = str(row.get("allele_name") or "").strip()
            if not allele_name:
                continue
            resolved = self._format_hla_allele_resolution(allele_name, normalized_resolution)
            if resolved not in seen:
                seen.add(resolved)
                alleles.append(resolved)
        return alleles

    def get_hla_frequencies(self, population: Optional[str] = None) -> Dict[str, float]:
        """Get HLA allele frequencies from visible sample typing data.

        This helper requires the same scientific-query capability as
        ``get_hla_alleles()``.
        """
        normalized_population = str(population).strip().lower() if population else None
        params: List[Any] = []
        filter_clause = ""
        if normalized_population:
            filter_clause = " WHERE LOWER(COALESCE(h.ethnicity, '')) = %s"
            params.append(normalized_population)

        rows = self._execute_scientific_rows(
            f"""
            WITH filtered_samples AS (
                SELECT DISTINCT s.sample_id
                  FROM samples s
                  LEFT JOIN hosts h ON h.host_id = s.host_id
                  {filter_clause}
            ),
            sample_totals AS (
                SELECT COUNT(*)::float AS total_samples
                  FROM filtered_samples
            )
            SELECT sht.allele_name,
                   COUNT(DISTINCT sht.sample_id)::float / NULLIF((SELECT total_samples FROM sample_totals), 0)
                       AS frequency
              FROM sample_hla_typing sht
              JOIN filtered_samples fs ON fs.sample_id = sht.sample_id
             GROUP BY sht.allele_name
             ORDER BY frequency DESC NULLS LAST, sht.allele_name ASC
            """,
            params or None,
        )

        frequencies: Dict[str, float] = {}
        for row in rows:
            allele_name = str(row.get("allele_name") or "").strip()
            if not allele_name:
                continue
            frequency = row.get("frequency")
            frequencies[allele_name] = float(frequency or 0.0)
        return frequencies

    def wait_for_module_run(
        self,
        run_id: str,
        *,
        timeout: float = 600.0,
        poll_interval: float = 5.0,
    ) -> Dict[str, Any]:
        """Poll a module run until it reaches a terminal state."""
        deadline = time.monotonic() + max(float(timeout), 0.0)
        interval = max(float(poll_interval), 0.0)

        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise APIError(f"Timed out waiting for module run {run_id} after {timeout}s")

            status_response = self.get_module_run(run_id)
            status = str(status_response.get("status") or status_response.get("state") or "").lower()

            if status in _COMPLETED_RUN_STATUSES:
                return self.get_module_run_result(run_id)
            if status in _FAILED_RUN_STATUSES:
                error_info = status_response.get("error") or {}
                if isinstance(error_info, dict):
                    message = str(error_info.get("message") or status or "module run failed")
                else:
                    message = str(error_info or status or "module run failed")
                raise APIError(message)

            time.sleep(min(interval, remaining))

    def start_hla_binding_prediction(
        self,
        peptides: List[str],
        alleles: List[str],
        *,
        method: str = "netmhcpan",
        module_id: Optional[str] = None,
        version: Optional[str] = None,
        compute_profile: Optional[str] = None,
        wait: bool = False,
        timeout: float = 600.0,
        poll_interval: float = 5.0,
    ) -> Dict[str, Any]:
        """Start a binding-prediction module run using a configured module id.

        This is a convenience wrapper over the normal async module-run surface.
        It requires either an explicit ``module_id`` or a configured
        ``HLA_BINDING_MODULE_ID`` / ``hla_binding_module_id`` value.
        """
        resolved_module_id = self._resolve_hla_binding_module_id(module_id)
        if not resolved_module_id:
            raise APIError(
                "No HLA binding module is configured. Set HLA_BINDING_MODULE_ID "
                "or configure hla_binding_module_id in the SDK config, or call "
                "start_module_run() directly with your binding module."
            )

        normalized_method = str(method or "netmhcpan").strip().lower()
        if normalized_method not in _HLA_BINDING_METHODS:
            raise APIError(
                f"Unsupported HLA binding method {method!r}. "
                f"Supported methods: {', '.join(sorted(_HLA_BINDING_METHODS))}."
            )

        cleaned_peptides: List[str] = []
        for peptide in peptides or []:
            value = str(peptide or "").strip().upper()
            if not value:
                continue
            if not _HLA_PEPTIDE_RE.match(value):
                raise APIError(f"Invalid peptide sequence {value!r}. Peptides must contain only canonical amino acids.")
            if not 8 <= len(value) <= 15:
                raise APIError(f"Invalid peptide length for {value!r}. Peptides must be 8-15 amino acids long.")
            cleaned_peptides.append(value)
        if not cleaned_peptides:
            raise APIError("At least one peptide is required")

        cleaned_alleles = [str(allele).strip() for allele in (alleles or []) if str(allele).strip()]
        if not cleaned_alleles:
            raise APIError("At least one HLA allele is required")

        run_response = self.start_module_run(
            resolved_module_id,
            parameters={
                "peptides": cleaned_peptides,
                "alleles": cleaned_alleles,
                "method": normalized_method,
            },
            mode="workflow",
            compute_profile=compute_profile,
            version=version,
        )

        if not wait:
            return run_response

        run_id = run_response.get("runId") or run_response.get("id") or run_response.get("run_id")
        if not run_id:
            raise APIError("Binding prediction run started without returning a run identifier")
        return self.wait_for_module_run(str(run_id), timeout=timeout, poll_interval=poll_interval)

    def predict_hla_binding(
        self,
        peptides: List[str],
        alleles: List[str],
        method: str = "netmhcpan",
        batch_size: int = 100,
        module_id: Optional[str] = None,
        version: Optional[str] = None,
        compute_profile: Optional[str] = None,
        timeout: float = 600.0,
        poll_interval: float = 5.0,
    ) -> List[Dict[str, Any]]:
        """Predict HLA-peptide binding by dispatching a configured module run.

        Args:
            peptides: List of peptide sequences to evaluate.
            alleles: List of HLA alleles.
            method: Prediction method (default ``"netmhcpan"``).
            batch_size: Number of peptides per request; values <=0 disable batching.
            module_id: Optional binding module override. Falls back to
                ``HLA_BINDING_MODULE_ID`` or ``hla_binding_module_id`` config.
            version: Optional module version override.
            compute_profile: Optional compute profile override.
            timeout: Maximum seconds to wait for each batch result.
            poll_interval: Seconds between module-run status polls.

        Notes:
            This helper is not backed by a dedicated synchronous REST endpoint.
            It submits one or more workflow-mode module runs and waits for them.
        """
        if not peptides:
            return []

        if batch_size and batch_size > 0:
            batches = [peptides[i : i + batch_size] for i in range(0, len(peptides), batch_size)]
        else:
            batches = [list(peptides)]

        predictions: List[Dict[str, Any]] = []
        for batch in batches:
            result_payload = self.start_hla_binding_prediction(
                batch,
                alleles,
                method=method,
                module_id=module_id,
                version=version,
                compute_profile=compute_profile,
                wait=True,
                timeout=timeout,
                poll_interval=poll_interval,
            )
            predictions.extend(self._extract_hla_predictions(result_payload))

        return predictions

    # ------------------------------------------------------------------
    # Authentication helpers
    # ------------------------------------------------------------------

    def get_current_user(self) -> Dict[str, Any]:
        """Fetch the authenticated user payload."""

        payload = self._make_request("GET", "/v1/auth/me")
        if isinstance(payload, dict):
            return payload
        return {"user": payload}

    # Module endpoints

    def start_module_run(
        self,
        module_id: str,
        *,
        parameters: Optional[Dict[str, Any]] = None,
        mode: str = "interactive",
        compute_profile: Optional[str] = None,
        version: Optional[str] = None,
        navigation_data: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Start a module run via the canonical /v1/module-runs API."""
        payload: Dict[str, Any] = {
            "moduleId": module_id,
            "parameters": parameters or {},
            "mode": mode,
        }
        if compute_profile:
            payload["computeProfile"] = compute_profile
        if version:
            payload["version"] = version
        if navigation_data:
            payload["navigationData"] = navigation_data
        result = self._make_request("POST", "/v1/module-runs", json_data=payload)
        if isinstance(result, dict):
            return self._normalize_module_run_payload(result)
        return result

    def get_module_run(self, run_id: str) -> Dict[str, Any]:
        """Fetch module run status."""
        result = self._make_request("GET", f"/v1/module-runs/{run_id}")
        if isinstance(result, dict):
            return self._normalize_module_run_payload(result)
        return result

    def cancel_module_run(self, run_id: str) -> Dict[str, Any]:
        """Cancel an in-progress module run."""
        result = self._make_request("POST", f"/v1/module-runs/{run_id}/cancel")
        if isinstance(result, dict):
            return self._normalize_module_run_payload(result)
        return result

    def heartbeat_module_run(
        self,
        run_id: str,
        *,
        refresh_urls: Optional[bool] = None,
        progress: Optional[float] = None,
        progress_message: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Record a module run heartbeat and optionally refresh URLs.

        Args:
            run_id: The module run identifier.
            refresh_urls: When ``True``, generate fresh pre-signed URLs.
            progress: Execution progress percentage (0-100).
            progress_message: Human-readable progress description.
        """
        payload: Dict[str, Any] = {}
        if refresh_urls is not None:
            payload["refreshUrls"] = bool(refresh_urls)
        if progress is not None:
            payload["progress"] = max(0.0, min(100.0, float(progress)))
        if progress_message is not None:
            payload["progressMessage"] = str(progress_message)
        result = self._make_request(
            "POST",
            f"/v1/module-runs/{run_id}/heartbeat",
            json_data=payload or None,
        )
        if isinstance(result, dict):
            return self._normalize_module_run_payload(result)
        return result

    def refresh_module_run_urls(self, run_id: str) -> Dict[str, Any]:
        """Refresh pre-signed URLs for module run artifacts."""
        result = self._make_request("POST", f"/v1/module-runs/{run_id}/urls")
        if isinstance(result, dict):
            return self._normalize_module_run_payload(result)
        return result

    def complete_module_run(
        self,
        run_id: str,
        *,
        status: str = "completed",
        summary: Optional[Dict[str, Any]] = None,
        error: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Mark a UI-only module run as completed or failed."""
        payload: Dict[str, Any] = {"status": status}
        if summary is not None:
            payload["summary"] = summary
        if error:
            payload["error"] = error
        result = self._make_request(
            "POST",
            f"/v1/module-runs/{run_id}/complete",
            json_data=payload,
        )
        if isinstance(result, dict):
            return self._normalize_module_run_payload(result)
        return result

    def revoke_module_run(self, run_id: str, *, reason: Optional[str] = None) -> Dict[str, Any]:
        """Revoke module run token access (admin)."""
        payload = {"reason": reason} if reason else None
        result = self._make_request(
            "POST",
            f"/v1/module-runs/{run_id}/revoke",
            json_data=payload,
        )
        if isinstance(result, dict):
            return self._normalize_module_run_payload(result)
        return result

    def get_module_run_logs(
        self,
        run_id: str,
        *,
        limit: Optional[int] = None,
        next_token: Optional[str] = None,
        order: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Return CloudWatch log events for the module run."""
        params: Dict[str, Any] = {}
        if limit is not None:
            params["limit"] = limit
        if next_token:
            params["nextToken"] = next_token
        if order:
            params["order"] = order
        result = self._make_request("GET", f"/v1/module-runs/{run_id}/logs", params=params or None)
        if isinstance(result, dict):
            return self._normalize_module_run_payload(result)
        return result

    def get_module_run_result(self, run_id: str) -> Dict[str, Any]:
        """Fetch module run output JSON."""
        result = self._make_request("GET", f"/v1/module-runs/{run_id}/result")
        if isinstance(result, dict):
            return self._normalize_module_run_payload(result)
        return result

    def list_modules(self, category: Optional[str] = None, limit: int = 100) -> List[Dict[str, Any]]:
        """List available modules"""
        params = {"limit": limit}
        if category:
            params["category"] = category

        result = self._make_request("GET", "/v1/modules", params=params)
        if isinstance(result, dict):
            return result.get("modules", result.get("items", []))
        return result if isinstance(result, list) else []

    def get_module(self, module_id: str) -> Dict[str, Any]:
        """Fetch a specific module by id."""
        return self._make_request("GET", f"/v1/modules/{module_id}")

    def list_module_runs(
        self,
        *,
        module_id: Optional[str] = None,
        status: Optional[str] = None,
        page: int = 1,
        limit: int = 20,
        order: str = "desc",
    ) -> Dict[str, Any]:
        """List module runs with basic filters."""
        params: Dict[str, Any] = {
            "page": page,
            "limit": limit,
            "order": order,
        }
        if module_id:
            params["moduleId"] = module_id
        if status:
            params["status"] = status
        result = self._make_request("GET", "/v1/module-runs", params=params)
        if isinstance(result, dict) and isinstance(result.get("runs"), list):
            normalized = dict(result)
            normalized["runs"] = [
                self._normalize_module_run_payload(run) if isinstance(run, dict) else run for run in result["runs"]
            ]
            return normalized
        return result

    def iter_modules(
        self,
        category: Optional[str] = None,
        page_size: int = 100,
        max_results: Optional[int] = None,
    ) -> Iterator[Dict[str, Any]]:
        """Yield modules across pages without manual pagination."""
        page = 1
        yielded = 0
        limit = max(int(page_size), 1)

        while True:
            params = {"limit": limit, "page": page}
            if category:
                params["category"] = category

            result = self._make_request("GET", "/v1/modules", params=params)
            if isinstance(result, dict):
                items = result.get("modules", result.get("items", []))
            else:
                items = result if isinstance(result, list) else []

            if not items:
                return

            for item in items:
                yield item
                yielded += 1
                if max_results is not None and yielded >= max_results:
                    return

            if len(items) < limit:
                return

            page += 1

    def upload_module(self, module_path: str, module_name: str, version: str) -> Dict[str, Any]:
        """
        Upload module package to the real API endpoint using JSON (base64) payloads,
        with a multipart/form-data fallback for legacy environments.

        Args:
            module_path: Path to module zip file
            module_name: Name of the module
            version: Module version

        Returns:
            Upload response including module_id
        """
        # Validate file
        if not os.path.exists(module_path):
            raise APIError(f"Module package not found: {module_path}")
        file_size = os.path.getsize(module_path)
        if file_size > 50 * 1024 * 1024:  # 50MB limit
            raise APIError(f"Module package too large: {file_size / 1024 / 1024:.1f}MB (max 50MB)")

        # Prepare request
        url = f"{self.base_url}/v1/modules/upload"
        metadata = {
            "name": module_name,
            "version": version,
            "filename": os.path.basename(module_path),
            "size": file_size,
        }

        base_headers = self._headers()
        base_headers.setdefault("Accept", "application/json")
        idem = str(uuid.uuid4())
        base_headers.setdefault("Idempotency-Key", idem)

        for attempt in range(2):
            headers = dict(base_headers)
            headers.pop("Content-Type", None)

            try:
                with open(module_path, "rb") as fh:
                    fh.seek(0)
                    files = {
                        "module": (
                            os.path.basename(module_path),
                            fh,
                            "application/zip",
                        )
                    }
                    form_data = {"metadata": json.dumps(metadata)} if metadata else None

                    # Use configurable connect timeout and longer, configurable read timeout for uploads
                    _timeouts = self.config.get_request_timeouts()
                    _upload_read = self.config.get_upload_read_timeout()
                    resp = self.session.post(
                        url,
                        headers=headers,
                        files=files,
                        data=form_data,
                        timeout=(_timeouts.get("connect", 5.0), _upload_read),
                    )
            except OSError as exc:
                raise APIError(f"Failed to read module package: {exc}") from exc

            if resp.status_code == 401 and attempt == 0:
                if self.auth.refresh_token():
                    base_headers = self._headers()
                    base_headers.setdefault("Accept", "application/json")
                    base_headers.setdefault("Idempotency-Key", idem)
                    continue
                raise APIError(
                    "Authentication expired. Please run 'hla-compass auth login' again",
                    401,
                )

            if resp.status_code >= 400:
                # best-effort parse
                try:
                    err = resp.json()
                except Exception:
                    err = {"message": resp.text}
                raise APIError(
                    parse_api_error(resp, "Upload failed"),
                    resp.status_code,
                    details=err,
                )

            try:
                payload = resp.json()
            except Exception as exc:
                raise APIError("Invalid response from upload endpoint") from exc

            data = payload.get("data", payload)
            module_id = data.get("id") or data.get("module_id")
            if not module_id:
                raise APIError("Upload succeeded but module_id not returned")

            normalized = {"module_id": module_id}
            normalized.update(data)
            return normalized

        raise APIError("Upload failed after token refresh attempt")

    def get_rate_limit_state(self) -> Dict[str, Any]:
        """Return the active rate limit configuration and last wait duration."""
        return dict(self._rate_limit_state)

    def _emit_telemetry(
        self,
        *,
        method: str,
        url: str,
        status_code: Optional[int],
        duration: float,
        attempt: int,
        retries: int,
        params: Optional[Dict[str, Any]],
        idempotency_key: Optional[str],
        request_id: str,
        error: Optional[str] = None,
    ) -> None:
        payload = {
            "method": method,
            "url": url,
            "status_code": status_code,
            "duration_seconds": round(duration, 4),
            "attempt": attempt + 1,
            "max_retries": retries,
            "rate_limit_wait": self._rate_limit_state.get("last_wait_seconds"),
            "params": params or {},
            "idempotency_key": idempotency_key,
            "request_id": request_id,
            "error": error,
        }

        if self._telemetry_callback:
            try:
                self._telemetry_callback(dict(payload))
            except Exception as exc:  # pragma: no cover - user callback
                logger.warning("Telemetry callback raised an exception: %s", exc)

        log_method = logger.info
        if error or (status_code is not None and status_code >= 400):
            log_method = logger.warning

        sanitized = self._sanitize_telemetry_payload(payload)
        log_method("API request telemetry", extra={"api_request": sanitized})

    @staticmethod
    def _sanitize_telemetry_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
        """Return a redacted copy of telemetry suitable for logging."""
        sanitized = dict(payload)

        # Sanitize URL query parameters for sensitive-looking keys
        url = payload.get("url")
        if url and "?" in url:
            from urllib.parse import parse_qs, urlencode, urlparse, urlunparse

            parsed = urlparse(url)
            query = parse_qs(parsed.query)
            sensitive_keys = {
                "token",
                "key",
                "auth",
                "secret",
                "password",
                "apiKey",
                "access_token",
            }
            modified = False
            for k in query:
                if any(s in k.lower() for s in sensitive_keys):
                    query[k] = ["<redacted>"]
                    modified = True
            if modified:
                new_query = urlencode(query, doseq=True)
                sanitized["url"] = urlunparse(parsed._replace(query=new_query))

        params_value = payload.get("params")
        if isinstance(params_value, dict) and params_value:
            sanitized["params"] = {key: "<redacted>" for key in params_value.keys()}
        elif params_value:
            sanitized["params"] = "<redacted>"
        else:
            sanitized["params"] = {}

        if payload.get("idempotency_key"):
            sanitized["idempotency_key"] = "***"

        return sanitized

    def _record_failure(self):
        self._cb_failures += 1
        if self._cb_failures >= self._cb_threshold:
            self._cb_opened_at = time.time()

    def _record_success(self):
        self._cb_failures = 0
        self._cb_opened_at = None

    def register_module(self, module_id: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
        """
        Publish an uploaded module version using the real API endpoint.

        Args:
            module_id: Module ID from upload
            metadata: Module metadata (expects at least 'version')

        Returns:
            Publish response from the API
        """
        json_data = {
            "version": metadata.get("version", "1.0.0"),
            "notes": metadata.get("description", "Module published via SDK"),
        }
        return self._make_request("PUT", f"/v1/modules/{module_id}/publish", json_data=json_data)

    def suspend_module_version(self, module_id: str, version: str) -> Dict[str, Any]:
        """Suspend a published module version."""
        return self._make_request("POST", f"/v1/modules/{module_id}/versions/{version}/suspend")

    def delete_module(self, module_id: str, *, reason: Optional[str] = None) -> Dict[str, Any]:
        """Delete an entire module."""
        payload = {"reason": reason} if reason else None
        return self._make_request("DELETE", f"/v1/modules/{module_id}", json_data=payload)

    def register_container_module(
        self, payload: Dict[str, Any], *, idempotency_key: Optional[str] = None
    ) -> Dict[str, Any]:
        """Submit a container module publish job via the intake pipeline."""

        stable_key = idempotency_key or str(uuid.uuid4())
        return self._make_request(
            "POST",
            "/v1/modules/publish",
            json_data=payload,
            idempotency_key=stable_key,
        )

    def get_publish_config(self) -> Dict[str, Any]:
        """Retrieve registry/UI defaults for publishing modules."""

        return self._make_request("GET", "/v1/modules/publish/config")

    def get_module_publish_status(self, publish_id: str) -> Dict[str, Any]:
        """Retrieve the status of a module publish (intake) job."""

        return self._make_request("GET", f"/v1/modules/publish/{publish_id}")

    def set_telemetry_callback(self, callback: Optional[Callable[[Dict[str, Any]], None]]) -> None:
        """Register a callback that receives request telemetry dictionaries."""
        self._telemetry_callback = callback


class ModulesFacade:
    """Convenience facade for module operations."""

    def __init__(self, api_client: "APIClient"):
        self._api = api_client

    def list(self, category: Optional[str] = None, limit: int = 100) -> List[Dict[str, Any]]:
        """List modules available to the current org."""
        return self._api.list_modules(category=category, limit=limit)

    def get(self, module_id: str) -> Dict[str, Any]:
        """Fetch module metadata by id."""
        return self._api.get_module(module_id)

    def run(
        self,
        module_id: str,
        params: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """Start a module run (params maps to APIClient.start_module_run parameters)."""
        if params is not None and "parameters" not in kwargs:
            kwargs["parameters"] = params
        return self._api.start_module_run(module_id, **kwargs)

    def suspend_version(self, module_id: str, version: str) -> Dict[str, Any]:
        """Suspend a specific module version."""
        return self._api.suspend_module_version(module_id, version)

    def delete(self, module_id: str, *, reason: Optional[str] = None) -> Dict[str, Any]:
        """Delete a module."""
        return self._api.delete_module(module_id, reason=reason)


class RunsFacade:
    """Convenience facade for module run operations."""

    def __init__(self, api_client: "APIClient"):
        self._api = api_client

    def list(
        self,
        *,
        module_id: Optional[str] = None,
        status: Optional[str] = None,
        page: int = 1,
        limit: int = 20,
        order: str = "desc",
    ) -> Dict[str, Any]:
        """List module runs with basic filters."""
        return self._api.list_module_runs(
            module_id=module_id,
            status=status,
            page=page,
            limit=limit,
            order=order,
        )

    def get(self, run_id: str) -> Dict[str, Any]:
        """Fetch module run status."""
        return self._api.get_module_run(run_id)

    def logs(
        self,
        run_id: str,
        *,
        limit: Optional[int] = None,
        next_token: Optional[str] = None,
        order: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Fetch CloudWatch log events for the module run."""
        return self._api.get_module_run_logs(
            run_id,
            limit=limit,
            next_token=next_token,
            order=order,
        )

    def results(self, run_id: str) -> Dict[str, Any]:
        """Fetch module run results."""
        return self._api.get_module_run_result(run_id)

    def cancel(self, run_id: str) -> Dict[str, Any]:
        """Cancel an in-progress module run."""
        return self._api.cancel_module_run(run_id)

    def heartbeat(
        self,
        run_id: str,
        *,
        refresh_urls: Optional[bool] = None,
        progress: Optional[float] = None,
        progress_message: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Record a module run heartbeat."""
        return self._api.heartbeat_module_run(
            run_id,
            refresh_urls=refresh_urls,
            progress=progress,
            progress_message=progress_message,
        )

    def refresh_urls(self, run_id: str) -> Dict[str, Any]:
        """Refresh pre-signed URLs for module run artifacts."""
        return self._api.refresh_module_run_urls(run_id)

    def complete(
        self,
        run_id: str,
        *,
        status: str = "completed",
        summary: Optional[Dict[str, Any]] = None,
        error: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Mark a UI-only module run as completed or failed."""
        return self._api.complete_module_run(
            run_id,
            status=status,
            summary=summary,
            error=error,
        )

    def revoke(self, run_id: str, *, reason: Optional[str] = None) -> Dict[str, Any]:
        """Revoke module run token access (admin)."""
        return self._api.revoke_module_run(run_id, reason=reason)


class DatasetVersionsFacade:
    """Convenience facade for dataset versioning and data-lineage operations.

    Provides methods for creating and browsing immutable dataset snapshots,
    archiving versions, and querying the data-lineage DAG (which datasets
    fed into which runs and what artifacts they produced).

    Access via ``client.dataset_versions``::

        client = APIClient()
        versions = client.dataset_versions.list(status="active")
        lineage  = client.dataset_versions.get_run_lineage(run_id)
    """

    def __init__(self, api_client: "APIClient"):
        self._api = api_client

    # ------------------------------------------------------------------
    # Dataset version CRUD
    # ------------------------------------------------------------------

    def create(
        self,
        name: str,
        sources: List[Dict[str, Any]],
        *,
        description: Optional[str] = None,
        label: Optional[str] = None,
        total_row_count: Optional[int] = None,
        content_hash: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Create an immutable dataset version snapshot.

        Args:
            name: Human-readable unique name within the organisation.
            sources: List of source descriptors, each containing at least
                ``source_table`` and optionally ``source_schema``,
                ``filter_json``, ``row_count``, etc.
            description: Optional longer description of the snapshot.
            label: Short tag, e.g. ``"v1.0"`` or ``"pre-treatment-cohort"``.
            total_row_count: Aggregate row count across all sources.
            content_hash: SHA-256 fingerprint of the canonical content.

        Returns:
            The newly created dataset version record.
        """
        payload: Dict[str, Any] = {
            "name": name,
            "sources": sources,
        }
        if description is not None:
            payload["description"] = description
        if label is not None:
            payload["label"] = label
        if total_row_count is not None:
            payload["totalRowCount"] = total_row_count
        if content_hash is not None:
            payload["contentHash"] = content_hash
        return self._api._make_request("POST", "/v1/dataset-versions", json_data=payload)

    def list(
        self,
        *,
        status: Optional[str] = None,
        limit: int = 20,
        page: int = 1,
    ) -> Dict[str, Any]:
        """List dataset versions for the current organisation.

        Args:
            status: Optional filter by lifecycle status
                (``"active"``, ``"archived"``, ``"deleted"``).
            limit: Maximum number of records per page.
            page: 1-based page number.

        Returns:
            Paginated response containing ``items`` and pagination metadata.
        """
        params: Dict[str, Any] = {"limit": limit, "page": page}
        if status is not None:
            params["status"] = status
        return self._api._make_request("GET", "/v1/dataset-versions", params=params)

    def get(self, version_id: str) -> Dict[str, Any]:
        """Fetch a single dataset version by id.

        Args:
            version_id: UUID of the dataset version.

        Returns:
            The dataset version record including its sources.
        """
        return self._api._make_request("GET", f"/v1/dataset-versions/{version_id}")

    def archive(self, version_id: str) -> Dict[str, Any]:
        """Archive a dataset version (soft-delete).

        This sets the status to ``"archived"`` so the version no longer
        appears in default listings but is still retrievable by id.

        Args:
            version_id: UUID of the dataset version to archive.

        Returns:
            The updated dataset version record.
        """
        return self._api._make_request(
            "PATCH",
            f"/v1/dataset-versions/{version_id}",
            json_data={"status": "archived"},
        )

    # ------------------------------------------------------------------
    # Lineage queries
    # ------------------------------------------------------------------

    def get_run_lineage(self, run_id: str) -> Dict[str, Any]:
        """Retrieve the full lineage graph for a module run.

        Returns both the input dataset versions that fed into the run
        and the output artifacts it produced.

        Args:
            run_id: UUID of the module run.

        Returns:
            Lineage response containing ``inputs`` and ``outputs`` edge lists.
        """
        return self._api._make_request("GET", f"/v1/lineage/run/{run_id}")

    def get_downstream_runs(
        self,
        dataset_version_id: str,
        *,
        limit: int = 20,
    ) -> Dict[str, Any]:
        """List module runs that consumed a given dataset version.

        Args:
            dataset_version_id: UUID of the dataset version.
            limit: Maximum number of results.

        Returns:
            Paginated list of downstream module run references.
        """
        params: Dict[str, Any] = {"limit": limit}
        return self._api._make_request(
            "GET",
            f"/v1/lineage/dataset/{dataset_version_id}/downstream",
            params=params,
        )

    def get_artifact_lineage(self, artifact_uri: str) -> Dict[str, Any]:
        """Trace the lineage of an S3 artifact back to its producing run
        and the upstream dataset versions.

        Args:
            artifact_uri: Canonical S3 URI of the artifact, e.g.
                ``s3://bucket/org/runs/run-id/outputs/result.parquet``.

        Returns:
            Lineage record including the producing run and upstream inputs.
        """
        return self._api._make_request(
            "GET",
            "/v1/lineage/artifact",
            params={"uri": artifact_uri},
        )
