"""Tiny client for the Genomic Intelligence API.

Wraps `requests` with bearer auth, the unified `{data, meta}` /
`{error}` envelope, and a polling helper for async jobs. Drop this file
into your project and `from gi_client import Client`.

Endpoints covered:
  - the six per-task predictors, each its own published operation:
    ``POST /v1/tasks/{promoter,splice,enhancer,chromatin,annotation,
    expression}/predict``
  - the composite ``POST /v1/workflows/find-genes-and-predict-expression``
  - the async job endpoints ``GET /v1/tasks/jobs[/{job_id}]``
  - model discovery ``GET /v1/tasks/{task}/models``

Contract reference: https://api.genomicintelligence.ai/redoc
"""

from __future__ import annotations

import time
from typing import Any, Dict, Mapping, Optional

import requests


#: Minimum ``sequence`` length per task, in bp, measured AFTER whitespace is
#: stripped. This is a MIRROR for local pre-flight only — the authority is the
#: ``minLength`` published on each task's request schema in the live OpenAPI
#: document (``GET /v1/openapi.json``, public and unauthenticated), generated
#: from the service's ``api/core/limits.py``.
#:
#: The two can no longer disagree quietly. ``scripts/check_contract.py`` reads
#: this table straight out of this file and asserts it against
#: ``scripts/contract.json``, a pin of the live document; it runs on every PR
#: and **blocks the deploy**. ``scripts/sync_contract.py --check`` re-verifies
#: the pin against the live API daily. If you need to change a number here,
#: re-pin first and reconcile this table to the contract — never the reverse.
TASK_MIN_BP: Dict[str, int] = {
    "promoter": 300,
    "splice": 100,
    "enhancer": 50,
    "chromatin": 200,
    "annotation": 1_000,
    "expression": 9_198,
    "find-genes-and-predict-expression": 1_000,
}

#: Maximum ``sequence`` length, in bp, for every task (same mirror caveat).
#: Also published per model as ``bio_spec.request_max_bp``.
MAX_BP = 500_000

#: Above this, the composite workflow refuses synchronous delivery with
#: ``413 sync_too_large`` — resend the same body with ``Prefer: respond-async``.
FIND_GENES_SYNC_LIMIT_BP = 50_000


class GIError(RuntimeError):
    """Raised on any non-2xx response from the API.

    Attributes mirror the unified error envelope so callers can switch
    on ``code`` rather than HTTP status alone:
        {"error": {"code": "...", "message": "...",
                   "request_id": "...", "details": ...}}

    ``details`` is carried through verbatim and never parsed. Its shape is
    keyed on the sibling ``code`` and has varied across releases — for
    ``validation_failed`` you may see either a bare ``[{loc, msg, type}, ...]``
    array or ``{"errors": [...]}`` — and splice's response-size cap emits a
    ``validation_failed`` with neither (``{task, record_count,
    maximum_records, sequence_length, threshold}``). Branch on ``code`` and
    treat ``details`` as display-only, and you are correct against all of them.
    """

    def __init__(
        self,
        status: int,
        body: Dict[str, Any],
        headers: Optional[Mapping[str, str]] = None,
    ):
        err = (body or {}).get("error", {}) if isinstance(body, dict) else {}
        self.status = status
        self.code = err.get("code", "http_error")
        self.message = err.get("message", "")
        # The X-Request-Id header is set on every response, including any
        # error whose body omits error.request_id. Fall back to it — support
        # tickets always need a correlation id.
        self.request_id = err.get("request_id") or (headers or {}).get("X-Request-Id")
        self.details = err.get("details")
        rid = self.request_id or "unset"
        super().__init__(f"[{status} {self.code}] {self.message} (request_id={rid})")


class Client:
    """Thin synchronous client.

    >>> c = Client(api_key="gi_…")
    >>> r = c.predict("promoter", sequence="ACGT" * 500, sequence_name="demo")
    >>> r["meta"]["inference_time_ms"]
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.genomicintelligence.ai",
        timeout: float = 120.0,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._session = requests.Session()
        self._session.headers.update(
            {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
                "Accept": "application/json",
            }
        )

    # ------------------------------------------------------------------ helpers

    def _check(self, resp: requests.Response) -> Dict[str, Any]:
        try:
            body = resp.json()
        except ValueError:
            # A response arrived but is not JSON — usually an edge proxy. Use
            # http_error: error.code is a closed enum published in the schema,
            # so a synthesized value outside it is one no caller can match on.
            # Client-origin errors carry no request_id (the server never saw
            # this one), which is how you tell them apart from server errors.
            body = {"error": {"code": "http_error", "message": resp.text[:200]}}
        if not resp.ok:
            raise GIError(resp.status_code, body, resp.headers)
        return body

    # ----------------------------------------------------------------- requests

    def health(self) -> Dict[str, Any]:
        r = self._session.get(f"{self.base_url}/health", timeout=self.timeout)
        return self._check(r)

    def list_models(self, task: str) -> Dict[str, Any]:
        """``GET /v1/tasks/{task}/models`` — public; the key is sent but not needed.

        Returns a FLAT ``{task, default_model, models: [...]}`` object, not
        the ``{data, meta}`` envelope. Each model carries ``bio_spec`` with
        ``request_max_bp`` (the enforced ceiling), ``context_window_bp``
        (the model's own window — compare your sequence length against it to
        know whether the model saw real sequence or padding; null for
        annotation and expression) and ``trained_window_bp`` (fixed
        receptive field where the model has one; null for sliding-window
        models). Expression models also carry ``recommended_flank_bp``.
        Never hard-code a window, flank or model id: use the model your
        user named, else ``default_model``, and pass it as ``model`` on
        every call of a run; take F = the largest ``recommended_flank_bp``
        among the listed models (one fetch serves any of them), never below
        4,599 and with 2F within ``request_max_bp``. Put the TSS at offset F
        of the gene-sense sequence and send ``tss_index = F``. With ``tss``
        the 0-based genomic TSS position, fetch ``[tss - F, tss + F)`` on
        the plus strand, or ``(tss - F, tss + F]`` reverse-complemented on
        the minus strand.
        ``request_max_bp`` is the only cap; the window fields describe what
        the model reads, not what the route accepts.

        An unknown ``task`` is ``404 not_found``.
        """
        r = self._session.get(
            f"{self.base_url}/v1/tasks/{task}/models", timeout=self.timeout
        )
        return self._check(r)

    @staticmethod
    def _body(
        sequence: str,
        sequence_name: str,
        model: Optional[str],
        options: Optional[Dict[str, Any]],
        tss_index: Optional[int] = None,
    ) -> Dict[str, Any]:
        body: Dict[str, Any] = {"sequence": sequence, "sequence_name": sequence_name}
        if model is not None:
            body["model"] = model
        if options is not None:
            body["options"] = options
        if tss_index is not None:
            body["tss_index"] = tss_index
        return body

    def predict(
        self,
        task: str,
        sequence: str,
        sequence_name: str = "sequence",
        model: Optional[str] = None,
        options: Optional[Dict[str, Any]] = None,
        tss_index: Optional[int] = None,
    ) -> Dict[str, Any]:
        """Synchronous prediction. Returns the full ``{data, meta}`` body.

        ``task`` is one of ``promoter``, ``splice``, ``enhancer``,
        ``chromatin``, ``annotation``, ``expression`` — each a separately
        published operation at the same URL this method builds. An
        unrecognised task is ``404 not_found``, not a validation error.

        ``model`` is a TOP-LEVEL field, not ``options["model"]``: every
        task's ``options`` object is closed (``additionalProperties:
        false``), so an unknown key — ``model`` included — is rejected with
        ``422 validation_failed`` at ``loc ["body","options","<key>"]``.
        The accepted option keys differ per task; see the request schema in
        ReDoc (``PromoterOptions``, ``SpliceOptions``, …). ``expression``
        additionally *requires* ``options["description"]``.

        ``sequence`` must be within this task's bounds (see
        :data:`TASK_MIN_BP` / :data:`MAX_BP`, measured after whitespace is
        stripped); outside them is ``422 validation_failed``, never ``413``.
        A sequence above the floor but shorter than the selected model's
        ``bio_spec.context_window_bp`` (from :meth:`list_models`) is still
        accepted and scored — against a window padded out to that size.

        ``tss_index`` applies to the ``expression`` task only: the 0-based
        offset of the TSS into the whitespace-stripped ``sequence``. Every
        expression model takes at least 9,198 bp with the TSS at least
        4,599 bp from each end, i.e. ``4599 <= tss_index <= len(sequence) -
        4599``; ``tss_index`` is required unless ``sequence`` is exactly
        9,198 bp. 9,198 bp is a floor, not a target: longer TSS-centred
        sequence is better, and the model list says how much (the rule and
        strand handling are in :meth:`list_models`). There is no
        ``tes_index``. The client does not
        validate any of this — the API returns ``422 validation_failed``
        (the ``tss_index`` checks at ``loc ["body"]``, so switch on
        ``error.code``). A window dominated by ``N`` or low-complexity
        repeats can also be refused with ``422``: each model needs roughly
        2-3.3 kb of real sequence upstream of the TSS and 2.6 kb
        downstream.
        """
        body = self._body(sequence, sequence_name, model, options, tss_index)
        r = self._session.post(
            f"{self.base_url}/v1/tasks/{task}/predict",
            json=body,
            timeout=self.timeout,
        )
        return self._check(r)

    def submit_async(
        self,
        task: str,
        sequence: str,
        sequence_name: str = "sequence",
        model: Optional[str] = None,
        options: Optional[Dict[str, Any]] = None,
        tss_index: Optional[int] = None,
    ) -> str:
        """Submit a task in async mode. Returns the ``job_id``.

        ``Prefer: respond-async`` is a declared header parameter on all six
        predict operations and on the composite workflow. Async delivery is
        JSON-only: combining it with a text ``?format=`` is a ``400``.

        ``tss_index`` is the ``expression``-only TSS offset; see :meth:`predict`.
        """
        body = self._body(sequence, sequence_name, model, options, tss_index)
        r = self._session.post(
            f"{self.base_url}/v1/tasks/{task}/predict",
            headers={"Prefer": "respond-async"},
            json=body,
            timeout=self.timeout,
        )
        body = self._check(r)
        # 202 envelope is {data: {job_id, status, links}, meta: {...}} —
        # uniform with the {data, meta} shape every other successful
        # response on the inference URL produces.
        return body["data"]["job_id"]

    # --------------------------------------------------- composite workflow

    def find_genes_and_predict_expression(
        self,
        sequence: str,
        description: str,
        sequence_name: str = "sequence",
        options: Optional[Dict[str, Any]] = None,
        async_: bool = False,
    ) -> Dict[str, Any]:
        """``POST /v1/workflows/find-genes-and-predict-expression``.

        Runs annotation over the region, cuts a window around each found
        gene's TSS, and scores expression for each — one round trip. Each
        gene is scored with the API's default expression model, using as
        much flank around each TSS as the submitted region holds, so longer
        regions give better predictions.

        ``sequence`` is 1,000–500,000 bp (the annotation stage's floor).
        ``description`` is the experimental context applied to every gene;
        it is required (missing or empty is ``422 validation_failed``) and
        is passed as ``options["description"]``. Extra ``options`` may add
        ``annotation_model``, ``expression_model`` and
        ``shift_coordinates``; the object is closed, so any
        other key is a ``422``.

        Above :data:`FIND_GENES_SYNC_LIMIT_BP` (50,000 bp) synchronous
        delivery is refused with ``413 sync_too_large`` carrying
        ``details = {sequence_length, threshold}`` — resend the identical
        body with ``async_=True``.

        Returns the ``{data, meta}`` body; with ``async_=True`` that body is
        the ``202`` acceptance ``{data: {job_id, status, links}, meta}``,
        so poll :meth:`wait_for_job` on ``body["data"]["job_id"]``.

        ``meta.task_specific_counts`` is
        ``{genes_found, genes_predicted, genes_skipped}`` with
        ``genes_predicted + genes_skipped == genes_found``; per-gene causes
        are in ``data.expression_predictions[].skip_reason``.
        """
        opts: Dict[str, Any] = dict(options or {})
        opts["description"] = description
        body = self._body(sequence, sequence_name, None, opts)
        headers = {"Prefer": "respond-async"} if async_ else None
        r = self._session.post(
            f"{self.base_url}/v1/workflows/find-genes-and-predict-expression",
            headers=headers,
            json=body,
            timeout=self.timeout,
        )
        return self._check(r)

    def get_job(self, job_id: str) -> requests.Response:
        """Single poll. The caller inspects ``status_code`` to discriminate."""
        return self._session.get(
            f"{self.base_url}/v1/tasks/jobs/{job_id}", timeout=self.timeout
        )

    def wait_for_job(
        self,
        job_id: str,
        poll_interval: float = 2.0,
        max_wait: float = 30 * 60,
        on_progress=None,
    ) -> Dict[str, Any]:
        """Poll until terminal. Returns ``{data, meta}`` on success, raises ``GIError``."""
        deadline = time.monotonic() + max_wait
        while True:
            r = self.get_job(job_id)
            if r.status_code == 200:
                return r.json()
            if r.status_code == 202:
                if on_progress is not None:
                    try:
                        # 202 poll body is also {data: {progress: {...}}, meta}.
                        on_progress((r.json().get("data") or {}).get("progress") or {})
                    except Exception:
                        pass
                if time.monotonic() > deadline:
                    raise TimeoutError(f"job {job_id} did not finish within {max_wait}s")
                time.sleep(poll_interval)
                continue
            # Terminal error
            try:
                body = r.json()
            except ValueError:
                # http_error for the same reason as _check above: the code must
                # come from the published enum.
                body = {"error": {"code": "http_error", "message": r.text[:200]}}
            raise GIError(r.status_code, body, r.headers)

    def list_jobs(self, limit: int = 50) -> Dict[str, Any]:
        r = self._session.get(
            f"{self.base_url}/v1/tasks/jobs",
            params={"limit": limit},
            timeout=self.timeout,
        )
        return self._check(r)