"""Minimal client for the Kilobaser One-XT REST API.

Everything the other example scripts do goes through this file. It is deliberately
small and dependency-light -- `requests` and the standard library -- so you can read
it end to end and then write your own.

    from kb_client import Kilobaser, KilobaserError

    kb = Kilobaser("kilobaser.lab.example.org", "apibot", "secret", verify=False)
    kb.login()
    kb.queue_oligo("LAMP-042_RPP30-F3", "GACCTGCTAGATCGTACG")
    for entry in kb.queue():
        print(entry["priority"], entry["title"])

Four things about this API surprise people, and this file exists mostly to absorb
them:

1. The device uses a **session cookie**, not a token, and the session lives in the
   daemon's memory -- so it disappears when the daemon restarts. Every call here
   re-logs-in once and retries when that happens.
2. Errors come back **double-encoded**: the body is `{"error": "<a JSON string>"}`.
   `KilobaserError` unwraps it so you get a real code and message. A few responses put a
   plain string there instead, so the unwrapping has to be forgiving.
3. `POST /processRuns/queue` **does not return the id** of what you queued. Use
   `queue_oligo(...)`, which reads the queue back and returns the created entry.
4. The queue is **rebuilt by the machine** whenever anything changes. Entry ids are
   stable, but priorities are renumbered and unplannable entries are dropped. Treat
   the queue you read (or the `processRunQueue` event) as the truth, never your own
   record of what you submitted.
"""

import json
import time

import requests


class KilobaserError(RuntimeError):
    """An error response from the device, with the code unwrapped.

    The device answers `{"error": "{\\"code\\": \\"2-11-12\\", ...}"}` -- a JSON
    string inside a JSON object. `code` is the stable identifier to branch on;
    `message` is human-facing text that may be reworded between releases.
    """

    def __init__(self, status, code, message, raw=None):
        super().__init__(f"[{code}] {message} (HTTP {status})")
        self.status = status
        self.code = code
        self.message = message
        self.raw = raw

    @classmethod
    def from_response(cls, response):
        raw = response.text
        code, message = "", raw.strip()
        try:
            envelope = response.json().get("error", "")
            if isinstance(envelope, str):
                try:
                    inner = json.loads(envelope)
                    code = inner.get("code", "")
                    message = inner.get("message", envelope)
                except json.JSONDecodeError:
                    message = envelope  # some errors are a plain string
        except ValueError:
            pass
        return cls(response.status_code, code, message, raw)


# Session problems. The device answers 400 (not 401) when the cookie is missing or
# expired, so status code alone cannot tell you to re-authenticate -- the code can.
AUTH_ERROR_CODES = {"2-1-1", "2-1-2", "2-1-3", "2-1-4"}


class Kilobaser:
    def __init__(self, host, username, password, verify=True, timeout=30):
        """`verify` accepts what requests accepts: True, False, or a CA bundle path.

        Devices ship with a self-signed certificate, so `verify=False` is the usual
        starting point. Export the certificate and pass its path instead as soon as
        you move past experimenting -- see the guide's authentication chapter.
        """
        self.base = f"https://{host}/api"
        self.username = username
        self.password = password
        self.timeout = timeout
        self.session = requests.Session()
        self.session.verify = verify
        self.user = None
        if verify is False:
            # Otherwise urllib3 warns on every single request. Silencing it is only
            # reasonable because we have already decided to trust this connection;
            # do not copy this line into code that talks to anything else.
            requests.packages.urllib3.disable_warnings(
                requests.packages.urllib3.exceptions.InsecureRequestWarning)

    # ---------------------------------------------------------------- plumbing

    def login(self):
        response = self.session.post(
            f"{self.base}/login",
            json={"username": self.username, "password": self.password},
            timeout=self.timeout,
        )
        if not response.ok:
            raise KilobaserError.from_response(response)
        body = response.json()
        self.user = body["user"]
        return body

    def request(self, method, path, body=None, _retried=False):
        """One HTTP call, with error unwrapping and a single re-login retry.

        Note the explicit Content-Type: the device rejects any request that carries
        a body without it, with 415 and a message that does not mention the header.
        """
        headers = {"Content-Type": "application/json"} if body is not None else {}
        response = self.session.request(
            method,
            f"{self.base}{path}",
            data=json.dumps(body) if body is not None else None,
            headers=headers,
            timeout=self.timeout,
        )
        if response.ok:
            if not response.content:
                return None
            try:
                return response.json()
            except ValueError:
                return response.text  # e.g. the batch import, see import_fasta

        # A device that is still booting answers 503 and tells you how long to wait.
        # Worth honouring: it is the first thing you meet after a restart, before the
        # session has even had a chance to be rejected.
        if response.status_code == 503 and not _retried:
            time.sleep(float(response.headers.get("Retry-After", 2)))
            return self.request(method, path, body, _retried=True)

        error = KilobaserError.from_response(response)
        if error.code in AUTH_ERROR_CODES and not _retried:
            self.login()
            return self.request(method, path, body, _retried=True)
        raise error

    def get(self, path):
        return self.request("GET", path)

    # ------------------------------------------------------------------- state

    def init(self):
        """The whole client-visible state in one call.

        This is the right way to start: it returns status, the queue, past runs,
        active errors, cartridge and chip settings and your session together, so a
        client can build its picture without a burst of requests. Use it to prime,
        then keep up to date with the event stream.
        """
        return self.get("/init")

    def status(self):
        return self.get("/status")

    def queue(self):
        """The queue, priority-ordered, as the machine currently plans it.

        Read via /init rather than GET /processRuns/queue: that endpoint returns
        entries in storage order and pages at 15, whereas init gives you the
        planned order, which is what the machine will actually do.
        """
        return self.init().get("processRunQueue") or []

    def cartridge_chip_settings(self):
        """Cartridge and chip catalogue, including the compatibility matrix.

        `operations[cartridgeType][chipKind]` exists exactly when that pair can be
        run. Read it instead of hardcoding a table -- the set of available cartridges
        and modifications differs between devices and grows between releases.
        """
        return self.init()["ccSettings"]

    # ------------------------------------------------------------------- queue

    def queue_oligo(self, title, sequence, cartridge_type="2", chip_kind="2",
                    priority=None, process_type="synthesis"):
        """Queue one oligo and return the created entry, including its id.

        The POST itself answers with the request echoed back and no id, so this
        reads the queue afterwards and picks out the new entry. That read-back is
        also your confirmation that the machine could plan the run at all.
        """
        before = {entry["id"] for entry in self.queue()}
        body = {
            "title": title,
            "processType": process_type,
            "cartridgeType": cartridge_type,
            "chipKind": chip_kind,
            "answers": {"sequence": sequence},
        }
        if priority is not None:
            body["priority"] = priority
        self.request("POST", "/processRuns/queue", body)
        for entry in self.queue():
            if entry["id"] not in before:
                return entry
        raise RuntimeError(
            f"queued {title!r} but it is not in the queue -- the machine could not "
            "plan it against its current state"
        )

    def import_fasta(self, content, cartridge_type="2", chip_kind="2", priority=0):
        """Queue every record of a FASTA file in one call.

        One cartridge/chip pair applies to the whole file, so a set of oligos with
        different modifications needs one call per pair. Each record's header line
        becomes the entry's title.

        The response body is a series of concatenated JSON objects rather than an
        array, so it cannot be parsed as JSON; it is ignored here. Read the queue
        afterwards to see what was created.
        """
        before = {entry["id"] for entry in self.queue()}
        self.request("POST", "/processRuns/queueImportBrowser", {
            "content": content,
            "processRun": {
                "processType": "synthesis",
                "cartridgeType": cartridge_type,
                "chipKind": chip_kind,
                "priority": priority,
            },
        })
        return [entry for entry in self.queue() if entry["id"] not in before]

    def update_queued(self, entry, **fields):
        """Change a queued entry. Pass the entry you read, plus what you changed.

        Setting `priority` moves the entry; the machine then renumbers the whole
        queue from 1, so the number you set is a sort key, not an address.
        """
        return self.request("PUT", f"/processRuns/queue/{entry['id']}",
                            {**entry, **fields})

    def delete_queued(self, entry):
        return self.request("DELETE", f"/processRuns/queue/{entry['id']}")

    # -------------------------------------------------------------------- runs

    def start(self, entry):
        """Start a queued entry. Only the first entry in the queue can be started.

        Anything else fails with 1-1-9, and starting while a run is in progress
        fails with 1-1-8. To run something sooner, move it to the front first.
        """
        return self.request("POST", "/processRuns", {"id": entry["id"]})

    def answer(self, run_id, answer_id, value="true"):
        """Answer one of the machine's checkpoint questions.

        Confirmations must be the string "true", not a boolean. Send one answer at
        a time and wait for it to leave `pendingAnswers` before sending the next.
        """
        return self.request("PUT", f"/processRuns/{run_id}/control",
                            {"answers": {answer_id: value}})

    def cancel(self, confirm=True):
        """Cancel the running process. Two steps: take the lease, then decide.

        `GET /control/cancel` reserves the cancel and starts a countdown that is
        visible to every client as `cancelLeaseExpiresAt`; the second call commits
        or releases it.
        """
        self.get("/control/cancel")
        return self.request(
            "PUT", "/control/cancel/confirm" if confirm else "/control/cancel/abort")

    def handle_error(self, error, cancel_run=False):
        return self.request("PUT", f"/errors/handle/{error['id']}",
                            {"cancelRun": cancel_run, "error": error})

    # ------------------------------------------------------------------ events

    def subscribe(self):
        """Yield `(event_name, payload)` from the device, forever.

        The stream carries no heartbeat, so a quiet machine simply sends nothing --
        silence does not mean the connection died, but it does mean a proxy with an
        idle timeout may close it without either end noticing. This reconnects with
        backoff and, because the server sends no event ids and so cannot replay what
        you missed, re-primes from /init after every connect. The first thing you
        receive is therefore always a synthetic ("init", snapshot).

        Note that a browser cannot consume this stream cross-origin: unlike the rest
        of the API, /api/subscribe sends no CORS headers. Run this in your backend.
        """
        delay = 1.0
        while True:
            try:
                response = self.session.get(f"{self.base}/subscribe", stream=True,
                                            timeout=(self.timeout, None))
                if not response.ok:
                    error = KilobaserError.from_response(response)
                    if error.code in AUTH_ERROR_CODES:
                        self.login()
                        continue
                    raise error
                delay = 1.0
                yield "init", self.init()
                yield from _parse_sse(response)
            except (requests.RequestException, KilobaserError) as exc:
                self.last_stream_error = exc
            time.sleep(delay)
            delay = min(delay * 2, 30.0)

    def watch(self):
        """Yield `(event_name, state)` where `state` is a live mirror of the device.

        This is the piece worth copying. The two collection events do not mean the
        same thing and must not be applied the same way: `processRunQueue` carries
        the entire queue and replaces it, while `processRuns` carries one run and
        has to be merged in by id. Getting that backwards produces a client whose
        queue slowly fills with entries the machine has already discarded.
        """
        state = {"status": None, "queue": [], "runs": {}, "errors": []}
        for name, data in self.subscribe():
            if name == "init":
                state["status"] = data.get("status")
                state["queue"] = data.get("processRunQueue") or []
                state["runs"] = {r["id"]: r for r in data.get("processRuns") or []}
                state["errors"] = data.get("errors") or []
            elif name == "status":
                state["status"] = data
            elif name == "processRunQueue":
                state["queue"] = data or []          # whole queue, replace
            elif name == "processRuns":
                state["runs"][data["id"]] = data     # single run, merge by id
            elif name == "errors":
                state["errors"] = data or []
            yield name, state

    def wait_for(self, predicate, timeout=3600):
        """Block until `predicate(state)` is true, then return the state.

        Everything an integration waits for is expressible this way -- a checkpoint
        appearing, a run finishing, the queue draining -- which is why this is the
        only waiting primitive the examples use.
        """
        deadline = time.monotonic() + timeout
        for _, state in self.watch():
            if predicate(state):
                return state
            if time.monotonic() > deadline:
                raise TimeoutError("timed out waiting for the device")
        raise RuntimeError("event stream ended")


def _parse_sse(response):
    """Turn a text/event-stream body into (event, payload) pairs.

    The wire format is a few lines per message and a blank line to end it. The
    device never sets `id:` or `retry:`, so this only needs `event:` and `data:`.
    """
    event, data = None, []
    for raw in response.iter_lines(decode_unicode=True):
        if raw is None:
            continue
        line = raw.rstrip("\r")
        if line == "":
            if event and data:
                try:
                    yield event, json.loads("\n".join(data))
                except json.JSONDecodeError:
                    pass
            event, data = None, []
        elif line.startswith(":"):
            continue  # comment, used by some servers as a keep-alive
        elif line.startswith("event:"):
            event = line[len("event:"):].strip()
        elif line.startswith("data:"):
            data.append(line[len("data:"):].lstrip())


def checkpoint(state):
    """The machine is waiting for a person, and what it is waiting for.

    Two conditions have to hold together: the mode is `check`, and there is at
    least one pending answer. Pending answers can linger for a moment while the
    machine is still working, so testing the list alone will make you answer too
    early.
    """
    status = state.get("status") or {}
    pending = status.get("pendingAnswers") or []
    mode = (status.get("currentState") or {}).get("mode")
    if mode == "check" and pending:
        return pending
    return None