"""Parallel workers against the hosted TauOpt service with no package installed.

The package-free twin of ``parallel_workers.py``: N threads share one
experiment created with ``parallel_bandwidth`` = N, each looping POST
suggestion -> evaluate -> POST observation over plain HTTP (the standard
library only). A 429 ``bandwidth_full`` is the EXPECTED backpressure signal:
the shared helper waits its ``Retry-After`` and re-asks (the worker sleeps 0.5
s itself when the header is missing). Any other 429 that carries
``Retry-After`` (a quota, the beta lane's rate limit) and a 503
``experiment_busy`` are waited out the same way (a quota 429 without a header,
5 s); ``experiment_quota_exceeded`` carries none and is raised. Every wait ends
at ``--deadline-s`` seconds after the script started (every call hands the
helper what is left of it as its wait budget): the worker then exits nonzero
naming the error (a bandwidth that never clears on a reused experiment means
stale open suggestions hold the slots;
``DELETE /v1/experiments/{id}/suggestions`` frees them). A reused finite
experiment that runs dry (409 ``space_exhausted``) stops its workers normally.
Exits nonzero on any bandwidth violation or duplicate assignment anywhere in
the run (dedup forbids repeats).

    export TAUOPT_BASE_URL="https://api.tauopt.com"
    export TAUOPT_TOKEN="tau_<the-value-the-tokens-page-showed>"
    python parallel_workers_rest.py --workers 4

Without the two variables it prints one line and exits 0. ``--experiment-id``
reuses an existing experiment (a batch job array shares one id) instead of
creating one; an experiment created here is archived at the end unless
``--keep``, a reused one never is.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import sys
import threading
import time
import urllib.error
import urllib.request
from collections import Counter


# --- rest helper (identical in every *_rest.py example; tests/test_rest_examples.py pins it) ---
class ApiError(RuntimeError):
    """A non-2xx answer: the service's one error shape,
    {"error": {"type", "message", "details"}}."""

    def __init__(self, status: int, body: dict, headers) -> None:
        err = body.get("error") or {}
        self.status = status
        self.error_type = err.get("type", "http_error")
        self.details = err.get("details")
        ra = headers.get("Retry-After")
        self.retry_after = float(ra) if ra and ra.replace(".", "", 1).isdigit() else None
        super().__init__(f"[{status}] {self.error_type}: {err.get('message', '')}")


class _NoRedirect(urllib.request.HTTPRedirectHandler):
    """Never follow a redirect: a 301 from http:// to https:// would turn a POST into a GET."""

    def redirect_request(self, *args, **kwargs):
        return None


class Api:
    """The hosted REST API over the standard library: a bearer token, JSON in and out."""

    def __init__(
        self, base_url: str, token: str, timeout_s: float = 120.0, wait_s: float = 600.0
    ) -> None:
        self._base = base_url.rstrip("/")
        self._token = token
        self._timeout = timeout_s
        self._wait = wait_s  # the budget per call for the waits a 429 or a 503 names
        self._opener = urllib.request.build_opener(_NoRedirect)
        self.last_exhaustion: dict | None = None  # the 409 space_exhausted details ask() last saw

    def call(
        self, method: str, path: str, body: dict | None = None, wait_s: float | None = None
    ) -> dict | None:
        """One request. A non-2xx answer raises ApiError (a 429 carries retry_after); 204
        returns None; a body that is not JSON raises ApiError whatever the status (the
        service answers JSON or 204; anything else is a proxy's page).

        A 429 that names a wait in Retry-After (``bandwidth_full``, a quota, the rate limit)
        and a 503 ``experiment_busy`` are waited out, the named interval (at least 1 s) in
        sleeps of at most 30 s, and re-sent, within a budget per call: the constructor's
        ``wait_s`` unless the call passes its own (a script with a deadline passes what is
        left of it). A wait that would end past the budget raises the answer instead, with
        no further request. ``experiment_quota_exceeded`` carries no Retry-After and raises
        at once. A failure before the answer is retried twice (see ``_send``); a timeout
        after the request was sent is raised as is.
        """
        budget = self._wait if wait_s is None else wait_s
        started = time.monotonic()
        while True:
            try:
                return self._send(method, path, body)
            except ApiError as e:
                busy = e.status == 503 and e.error_type == "experiment_busy"
                if e.retry_after is None or not (e.status == 429 or busy):
                    raise
                left = max(e.retry_after, 1.0)  # a Retry-After of 0 is never a tight loop
                if time.monotonic() - started + left > budget:
                    raise
                while left > 0:
                    time.sleep(min(left, 30.0))
                    left -= 30.0

    def _send(self, method: str, path: str, body: dict | None) -> dict | None:
        req = urllib.request.Request(
            self._base + path,
            data=None if body is None else json.dumps(body).encode(),
            method=method,
            headers={
                "Authorization": f"Bearer {self._token}",
                "Content-Type": "application/json",
                "Accept": "application/json",
            },
        )
        for attempt in range(3):
            try:
                with self._opener.open(req, timeout=self._timeout) as r:
                    status, raw, headers = r.status, r.read(), r.headers
                break
            except urllib.error.HTTPError as e:  # every non-2xx answer, the 3xx included
                status, raw, headers = e.code, e.read(), e.headers
                break
            except urllib.error.URLError:
                # a failure before the answer: DNS, connect, the TLS handshake, or the send
                # itself (urllib wraps all of them in URLError; whether any byte reached the
                # service is unknowable here, so a resend of a landed POST is possible in
                # principle) -- retried twice; a timeout while WAITING for the answer is a
                # bare TimeoutError, not caught here, because the request landed: list the
                # open suggestions first
                if attempt == 2:
                    raise
                time.sleep(1.0 + attempt)
        if status == 204:
            return None
        try:
            data = json.loads(raw)
        except ValueError:
            data = {"error": {"type": "http_error", "message": raw[:200].decode(errors="replace")}}
            raise ApiError(status, data, headers) from None  # a 2xx that is not JSON too
        if status >= 300:
            raise ApiError(status, data, headers)
        return data


def ask(
    api: Api, exp_id, body: dict | None = None, wait_s: float | None = None
) -> dict | None:
    """One suggestion, or None on proven exhaustion (a completion, not an error;
    the 409's details stay on ``api.last_exhaustion``).

    ``body`` is the optional SuggestionIn (a ``config``, ``max_wait_s``). A 429 that
    names a wait and a 503 ``experiment_busy`` are waited out by ``Api.call`` (``wait_s``
    is this call's budget); ``experiment_quota_exceeded`` carries no wait: archive an
    experiment first.
    """
    try:
        return api.call(
            "POST", f"/v1/experiments/{exp_id}/suggestions", body or {}, wait_s=wait_s
        )
    except ApiError as e:
        if e.status == 409 and e.error_type == "space_exhausted":
            api.last_exhaustion = e.details
            return None
        raise


def archive(api: Api, exp_id, wait_s: float | None = None) -> None:
    """DELETE the experiment; a 404 means it is already archived (the call is idempotent).
    ``wait_s`` is this call's budget for a wait the answer names."""
    try:
        api.call("DELETE", f"/v1/experiments/{exp_id}", wait_s=wait_s)
    except ApiError as e:
        if e.error_type != "not_found":
            raise


# --- end rest helper ---


def branin(x1: float, x2: float) -> float:
    a, b, c = 1.0, 5.1 / (4 * math.pi**2), 5 / math.pi
    r, s, t = 6.0, 10.0, 1 / (8 * math.pi)
    return a * (x2 - b * x1**2 + c * x1 - r) ** 2 + s * (1 - t) * math.cos(x1) + s


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.strip().split("\n\n")[0])
    ap.add_argument(
        "--workers", type=int, default=4,
        help="parallel workers, one open suggestion each (default %(default)s)",
    )
    ap.add_argument(
        "--observations-per-worker", type=int, default=6,
        help="observations each worker reports before it stops (default %(default)s)",
    )
    ap.add_argument(
        "--experiment-id",
        default=None,
        help="reuse this experiment (a batch job array shares one id) instead of creating one",
    )
    ap.add_argument(
        "--deadline-s",
        type=float,
        default=600.0,
        help="seconds after start past which a 429 wait (quota or bandwidth) exits nonzero"
             " (default %(default)s)",
    )
    ap.add_argument(
        "--keep", action="store_true", help="leave a created experiment active at the end"
    )
    args = ap.parse_args()
    base_url, token = os.environ.get("TAUOPT_BASE_URL"), os.environ.get("TAUOPT_TOKEN")
    if not base_url or not token:
        pair = (("TAUOPT_BASE_URL", base_url), ("TAUOPT_TOKEN", token))
        missing = "/".join(name for name, value in pair if not value)
        print(f"{missing} not set; skipping"
              " (export TAUOPT_BASE_URL=https://api.tauopt.com and a TAUOPT_TOKEN minted at app.tauopt.com/settings/tokens)")
        return 0

    started = time.monotonic()

    def budget() -> float:
        """What is left of --deadline-s: the wait budget every call hands the helper."""
        return max(0.0, args.deadline_s - (time.monotonic() - started))

    api = Api(base_url, token, wait_s=args.deadline_s)
    if args.experiment_id:
        exp = api.call("GET", f"/v1/experiments/{args.experiment_id}", wait_s=budget())
        print(f"reusing experiment {exp['id']} (parallel_bandwidth={exp['parallel_bandwidth']})")
    else:
        exp = api.call(
            "POST",
            "/v1/experiments",
            {
                "name": f"branin-parallel-{args.workers}w",
                "parameters": [
                    {"name": "x1", "type": "double", "bounds": [-5.0, 10.0]},
                    {"name": "x2", "type": "double", "bounds": [0.0, 15.0]},
                ],
                "metric": {"name": "branin", "objective": "minimize"},
                "parallel_bandwidth": args.workers,
            },
            wait_s=budget(),
        )
    exp_id = exp["id"]
    failures: dict[int, str] = {}
    try:
        if not args.experiment_id:  # inside the try: the print can fail (a closed stdout)
            print(f"created experiment {exp_id} (parallel_bandwidth={args.workers})")
        bandwidth = exp["parallel_bandwidth"]
        before = (exp.get("progress") or {}).get("observation_count", 0)

        lock = threading.Lock()
        seen: Counter[tuple] = Counter()
        in_flight = 0
        bandwidth_violations = 0
        completed = 0

        def wait_or_expire(seconds: float, e: ApiError, cause: str) -> None:
            """Sleep at most the remaining deadline; past it, exit nonzero naming the cause."""
            left = args.deadline_s - (time.monotonic() - started)
            if left > 0:
                time.sleep(min(seconds, left))
                left = args.deadline_s - (time.monotonic() - started)
            if left <= 0:
                raise RuntimeError(
                    f"deadline of {args.deadline_s:.0f}s passed while {cause} ({e})"
                ) from e

        def next_suggestion(worker_api: Api) -> dict | None:
            """Ask until a suggestion arrives (None on exhaustion). The helper waits out a 429 that
            names a Retry-After (and a 503 experiment_busy) under what is left of --deadline-s; a
            429 that reaches here would outlive the deadline or carries no header, and the
            branches below report it."""
            while True:
                try:
                    return worker_api.call(
                        "POST", f"/v1/experiments/{exp_id}/suggestions", {}, wait_s=budget()
                    )
                except ApiError as e:
                    if e.status == 409 and e.error_type == "space_exhausted":
                        return None
                    if e.status != 429 or e.error_type == "experiment_quota_exceeded":
                        raise
                    if e.error_type == "bandwidth_full":
                        # every slot is held by a running evaluation: wait and re-ask
                        wait_or_expire(
                            e.retry_after or 0.5,
                            e,
                            "every slot stayed full (stale open suggestions on a reused "
                            "experiment? DELETE /v1/experiments/{id}/suggestions frees them)",
                        )
                        continue
                    # a quota 429 (quota_exceeded, quota_unverifiable, rate_limited): wait it out
                    wait_or_expire(e.retry_after or 5.0, e, "waiting on a quota")

        def worker(wid: int) -> None:
            nonlocal in_flight, bandwidth_violations, completed
            worker_api = Api(base_url, token, wait_s=args.deadline_s)  # one session per thread
            for _ in range(args.observations_per_worker):
                sug = next_suggestion(worker_api)
                if sug is None:
                    print(f"[w{wid}] space exhausted; stopping")
                    break
                with lock:
                    seen[tuple(sorted(sug["assignments"].items()))] += 1
                    in_flight += 1
                    if in_flight > bandwidth:  # the server exceeded parallel_bandwidth
                        bandwidth_violations += 1
                value = branin(sug["assignments"]["x1"], sug["assignments"]["x2"])
                time.sleep(0.05)  # simulate evaluation work while holding the slot
                # the observe and the decrement share one critical section: the server
                # frees the slot when the observe returns, so a peer's ask landing between
                # the two would count one slot too many locally
                with lock:
                    worker_api.call(
                        "POST",
                        f"/v1/experiments/{exp_id}/observations",
                        {"suggestion_id": sug["id"], "value": value},
                        wait_s=budget(),
                    )
                    in_flight -= 1
                    completed += 1
                print(f"[w{wid}] {sug['provenance']['strategy']:>7} value={value:.4f}")

        def run(wid: int) -> None:
            try:
                worker(wid)
            except Exception as e:  # noqa: BLE001 -- the main thread reports it and exits nonzero
                failures[wid] = f"{type(e).__name__}: {e}"

        threads = [threading.Thread(target=run, args=(i,)) for i in range(args.workers)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        if failures:
            for wid, msg in sorted(failures.items()):
                print(f"[w{wid}] failed: {msg}", file=sys.stderr)
            return 1
        duplicates = {k: v for k, v in seen.items() if v > 1}
        final = api.call("GET", f"/v1/experiments/{exp_id}", wait_s=budget())
        count = (final.get("progress") or {}).get("observation_count", 0)
        print(f"\ncompleted: {completed} observations from this process;"
              f" the experiment holds {count}")
        assert count >= before + completed
        assert bandwidth_violations == 0
        assert not duplicates, (
            f"duplicate assignments in the run (dedup forbids repeats): {duplicates}"
        )
        print("zero bandwidth violations; all assignments in the run distinct")
    finally:
        if not args.experiment_id and not args.keep:
            failing = sys.exc_info()[1]  # the error above, if any, stays the one reported
            try:
                archive(api, exp_id, wait_s=budget())
                print(f"experiment {exp_id} archived")
            except Exception as e:  # noqa: BLE001 -- never hide the error above
                if failing is None and not failures:
                    raise
                print(f"experiment {exp_id} not archived: {e}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())