"""Genomic Intelligence API — partner quickstart.

Hits every public task endpoint with a real, biologically meaningful
sequence drawn from the bundled ``sequences/`` directory (the same
fixtures the service uses for its golden numeric-regression tests).
Sync calls are sized so each task completes in a few seconds on a warm
GPU; ``annotation`` is the one outlier and is intentionally exercised
both sync and async.

    pip install -r requirements.txt
    export GI_API_KEY=gi_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
    export GI_BASE_URL=https://api.genomicintelligence.ai   # optional
    python quickstart.py
"""

from __future__ import annotations

import os
import pathlib
import sys

import requests

from gi_client import Client, GIError


SEQ_DIR = pathlib.Path(__file__).parent / "sequences"


def load_fasta(filename: str) -> tuple[str, str]:
    """Return (sequence_name, sequence) from a single-record FASTA file."""
    text = (SEQ_DIR / filename).read_text()
    lines = text.splitlines()
    header = lines[0].lstrip(">").strip()
    sequence = "".join(line.strip() for line in lines[1:] if line.strip())
    return header, sequence.upper()


def _section(title: str) -> None:
    print(f"\n=== {title} ===")


def _summary(label: str, body: dict) -> None:
    meta = body.get("meta", {})
    counts = meta.get("task_specific_counts", {})
    print(
        f"{label:<22} model={meta.get('model'):<32} "
        f"{meta.get('inference_time_ms', '?')} ms  counts={counts}"
    )


def main() -> int:
    api_key = os.environ.get("GI_API_KEY")
    if not api_key:
        print("ERROR: set GI_API_KEY (your gi_… bearer key)", file=sys.stderr)
        return 2
    base_url = os.environ.get("GI_BASE_URL", "https://api.genomicintelligence.ai")

    client = Client(api_key=api_key, base_url=base_url)

    _section("Health")
    print(client.health())

    # All sequences are real biological inputs from the service's golden
    # fixtures — see client/sequences/. Each FASTA carries
    # provenance in its header (gene symbol, coordinates, assembly).
    tp53_name,    tp53_seq    = load_fasta("promoter_tp53.fa")            # human, ~26 kb, gene-sense (TP53 is - strand)
    hbb_name,     hbb_seq     = load_fasta("splice_hbb.fa")               # human, ~4 kb
    eve_name,     eve_seq     = load_fasta("enhancer_eve.fa")             # drosophila, ~10 kb, eve developmental enhancer (+ strand)
    chr19_name,   chr19_seq   = load_fasta("chromatin_active_promoter_chr19.fa")  # human, ~40 kb
    hbb_tss_name, hbb_tss_seq = load_fasta("expression_hbb_k562.fa")      # human, 9,198 bp (the minimum), TSS-centered

    _section("Sync inference (real biological sequences)")
    try:
        _summary("promoter (TP53)",
                 client.predict("promoter",  tp53_seq,  tp53_name))
        _summary("splice   (HBB)",
                 client.predict("splice",    hbb_seq,   hbb_name))
        # Enhancer model is Drosophila-trained — use a fly enhancer.
        _summary("enhancer (eve)",
                 client.predict("enhancer",  eve_seq,   eve_name))
        _summary("chromatin (chr19)",
                 client.predict("chromatin", chr19_seq, chr19_name))
    except GIError as exc:
        print(f"sync task failed: {exc}", file=sys.stderr)
        return 1
    except requests.RequestException as exc:
        # Transport-level: DNS, TLS, connection refused, read timeout.
        # Routine for a hosted service and not a request problem.
        print(f"sync task failed: network error: {type(exc).__name__}: {exc}", file=sys.stderr)
        return 1

    _section("Expression (HBB in K562, TSS-centered)")
    # Every expression model takes at least 9,198 bp with the TSS at least
    # 4,599 bp from each end. Anything shorter is a hard 422 -- no padding, no
    # truncation. The bundled fixture is that minimum for HBB, centered on its
    # canonical TSS, so the quickstart runs offline of any genome fetch; with
    # the K562 cell-type description this should report HIGH expression
    # (HBB is highly expressed in K562 erythroleukemia cells).
    #
    # At exactly 9,198 bp, tss_index is optional and defaults to 4599. The
    # floor is not a target: for real use, longer TSS-centred sequence is
    # better, and the model list says how much. From
    # client.list_models("expression") take the model your user named, else
    # default_model (pass it as model= on every call of a run), and
    # F = the largest bio_spec.recommended_flank_bp among the listed models,
    # never below 4,599 and with 2F within request_max_bp. Never hard-code a
    # window, flank or model id. With tss the 0-based genomic TSS position,
    # fetch [tss - F, tss + F) on the plus strand, or (tss - F, tss + F] and
    # reverse-complement it on the minus strand, so the TSS sits at offset F.
    # Then pass
    # tss_index=<0-based TSS offset into the whitespace-stripped sequence>; it
    # is REQUIRED at any other length and must satisfy
    # 4599 <= tss_index <= len(sequence) - 4599. Note "whitespace-stripped":
    # an index counted over line-wrapped FASTA characters is off by the newline
    # count and can silently score the wrong region, so assert on the
    # scored_window echoed back below.
    try:
        body = client.predict(
            "expression",
            sequence=hbb_tss_seq,
            sequence_name=hbb_tss_name,
            options={
                "description": (
                    "assay term name is polyA plus RNA-seq. "
                    "biosample summary is Homo sapiens K562."
                ),
            },
        )
        pred = body.get("data", {}).get("prediction", {})
        counts = body.get("meta", {}).get("task_specific_counts", {})
        print(
            f"expression: {pred.get('expression_log_tpm')} log(TPM+1)  "
            f"({pred.get('expression_tpm')} TPM)"
        )
        # Windowing provenance: confirm the server scored the window you meant.
        print(
            f"  tss_index={counts.get('tss_index')}  "
            f"scored_window={counts.get('scored_window')}"
        )
    except GIError as exc:
        print(f"expression failed: {exc}", file=sys.stderr)
    except requests.RequestException as exc:
        print(f"expression failed: network error: {type(exc).__name__}: {exc}", file=sys.stderr)

    _section("Async job — annotation on TP53 (~26 kb)")
    # Annotation is the slowest atomic task and the one where async
    # actually matters. TP53 is well-annotated; expect at least one
    # transcript in the response.
    try:
        job_id = client.submit_async(
            "annotation",
            sequence=tp53_seq,
            sequence_name=tp53_name,
        )
        print(f"submitted job_id={job_id}")

        def progress(p):
            pct = p.get("current_percent")
            msg = p.get("message", "")
            print(f"  {pct:>3}% {msg}")

        body = client.wait_for_job(job_id, poll_interval=2.0, on_progress=progress)
        meta = body.get("meta", {})
        counts = meta.get("task_specific_counts", {})
        transcripts = body.get("data", {}).get("transcripts", []) or []
        print(
            f"done — counts={counts}  "
            f"transcripts={len(transcripts)}  "
            f"total_time_ms={meta.get('inference_time_ms')}"
        )
    except GIError as exc:
        print(f"annotation failed: {exc}", file=sys.stderr)
        return 1
    except requests.RequestException as exc:
        print(f"annotation failed: network error: {type(exc).__name__}: {exc}", file=sys.stderr)
        return 1
    except TimeoutError as exc:
        # wait_for_job gives up rather than polling forever.
        print(f"annotation timed out: {exc}", file=sys.stderr)
        return 1

    _section("Recent jobs")
    print(client.list_jobs(limit=5))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())