"""A ShockAI Model API host: Python 3.11+, the standard library and
`cryptography`. Run it with

    pip install -r requirements.txt
    PUBLIC_ORIGIN=https://your-host.example.com python3 server.py

It serves your model (model.py, described by manifest.json) under
/shockai/v0, checks the platform's signature on every request but /healthz,
runs each case on a thread of its own, answers a run within 25 seconds or
accepts it and calls back until the platform takes the callback, and keeps
runs for 24 hours under data/. Guide:
https://app.shockai.co/developers/guides/host-your-model

Your key signs your callbacks. While manifest.json has no signing block, the
host makes data/host-key.pem on first start, serves its public half as
signing.callbackKeys and prints the block to copy into manifest.json. Once
manifest.json lists keys, the host serves them as they are, and will not
start without its key or with a key they do not list. To rotate, list both
keys in a new version and connect it, sign with the new key, then drop the
old one in the version after.

Environment: the same as server.mts.
    PUBLIC_ORIGIN       required: the origin the platform calls, as it spells it
    PORT                default 8080
    BASE_PATH           default /shockai/v0
    PLATFORM_KEYS_URL   default https://app.shockai.co/developers/keys/v0.json
    DATA_DIR            default ./data (runs, your key, the request log)
    BUILD_HASH          what is deployed, e.g. a commit; reported with every result
    HOST_SIGNING_KEY    your Ed25519 key as a PKCS#8 PEM; default data/host-key.pem
    SYNC_BUDGET_MS      default 25000: how long a run may take before it is answered 202
    MODEL_DELAY_MS      default 0: how long every run computes first, its thread kept busy, to try the 202 path
"""

import asyncio
import base64
import hashlib
import inspect
import json
import os
import re
import threading
import time
import urllib.error
import urllib.request
import uuid
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import urlsplit

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey

from model import run_model

HERE = Path(__file__).resolve().parent
ENV = os.environ
if not ENV.get("PUBLIC_ORIGIN"):
    raise SystemExit("Set PUBLIC_ORIGIN to the origin the platform calls, e.g. https://your-host.example.com")


def origin_of(url: str) -> str:
    """A URL's origin as the platform signs it (WHATWG URL, as in Node):
    lowercase scheme and host, IPv6 brackets kept, no default port."""
    parts = urlsplit(url)
    scheme = parts.scheme.lower()
    host = parts.hostname or ""  # lowercased, brackets dropped
    if ":" in host:
        host = f"[{host}]"
    port = parts.port
    default = {"https": 443, "http": 80}.get(scheme)
    return f"{scheme}://{host}" + (f":{port}" if port is not None and port != default else "")


ORIGIN = origin_of(ENV["PUBLIC_ORIGIN"])
PORT = int(ENV.get("PORT", "8080"))
BASE_PATH = ENV.get("BASE_PATH", "/shockai/v0").rstrip("/")
PLATFORM_KEYS_URL = ENV.get("PLATFORM_KEYS_URL", "https://app.shockai.co/developers/keys/v0.json")
DATA_DIR = Path(ENV.get("DATA_DIR", HERE / "data"))
BUILD_HASH = ENV.get("BUILD_HASH", "dev")
SYNC_BUDGET_S = int(ENV.get("SYNC_BUDGET_MS", "25000")) / 1000
MODEL_DELAY_S = int(ENV.get("MODEL_DELAY_MS", "0")) / 1000

SCHEME = "shockai-model-api-v0"
TOLERANCE_SECONDS = 300
MAX_BODY_BYTES = 16 * 1024 * 1024
KEEP_RUNS_S = 24 * 3600
CALLBACK_ATTEMPTS = 8
FINAL = {"succeeded", "failed", "cancelled"}
RUN_ID = re.compile(r"[A-Za-z0-9_-]{1,128}")

(DATA_DIR / "runs").mkdir(parents=True, exist_ok=True)
(DATA_DIR / "keys").mkdir(parents=True, exist_ok=True)
lock = threading.RLock()

# ---------------------------------------------------------------------------
# Keys and signatures (the scheme: /developers/guides/signing)
# ---------------------------------------------------------------------------


def _b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def _unb64url(text: str) -> bytes:
    return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))


MANIFEST_FILE = json.loads((HERE / "manifest.json").read_text())
# The keys manifest.json lists. Once it lists any, no key is ever made.
PINNED_KEYS = (MANIFEST_FILE.get("signing") or {}).get("callbackKeys") or []


def _load_host_key() -> Ed25519PrivateKey:
    if ENV.get("HOST_SIGNING_KEY"):
        pem = ENV["HOST_SIGNING_KEY"].replace("\\n", "\n").encode()
    else:
        path = DATA_DIR / "host-key.pem"
        if not path.exists():
            if PINNED_KEYS:
                raise SystemExit(
                    f"manifest.json lists signing.callbackKeys, but {path} is missing. Restore it, or set HOST_SIGNING_KEY. "
                    "A new key needs a new manifest version that lists it."
                )
            key = Ed25519PrivateKey.generate()
            path.write_bytes(
                key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption())
            )
            path.chmod(0o600)
        pem = path.read_bytes()
    return serialization.load_pem_private_key(pem, password=None)


def thumbprint(jwk: dict) -> str:
    members = json.dumps({"crv": jwk.get("crv"), "kty": jwk.get("kty"), "x": jwk.get("x")}, separators=(",", ":"))
    return _b64url(hashlib.sha256(members.encode()).digest())


HOST_KEY = _load_host_key()
HOST_JWK = {
    "kty": "OKP",
    "crv": "Ed25519",
    "x": _b64url(HOST_KEY.public_key().public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)),
}
HOST_KID = thumbprint(HOST_JWK)
if not PINNED_KEYS:
    pin = json.dumps({"signing": {"callbackKeys": [HOST_JWK]}}, separators=(",", ":"))
    print(f"manifest.json lists no signing key, so the host serves its own. To pin it, add to manifest.json: {pin}", flush=True)
elif not any(isinstance(key, dict) and thumbprint(key) == HOST_KID for key in PINNED_KEYS):
    raise SystemExit(
        f"The host's key (kid {HOST_KID}) is not among manifest.json's signing.callbackKeys. "
        "Sign with a key it lists, or list this one in a new version."
    )


def signature_base(timestamp: str, method: str, target: str, idempotency_key: str, body: bytes) -> bytes:
    lines = [SCHEME, timestamp, method.upper(), target, idempotency_key, hashlib.sha256(body).hexdigest()]
    return "\n".join(lines).encode()


platform_keys: list = []
keys_fetched_at = 0.0
keys_lock = threading.Lock()


def refresh_platform_keys() -> None:
    """Fetch the platform's keys, at most once a minute. A request that
    arrives while a fetch is on its way waits for it, rather than being
    refused."""
    global platform_keys, keys_fetched_at
    with keys_lock:
        if time.time() - keys_fetched_at < 60:
            return
        try:
            with urllib.request.urlopen(PLATFORM_KEYS_URL, timeout=10) as response:
                platform_keys = json.load(response)["keys"]
        except Exception:  # noqa: BLE001 - keep the keys we had; the next unknown key id asks again
            pass
        finally:
            keys_fetched_at = time.time()


def signature_problem(headers, method: str, target: str, body: bytes):
    """Why a request's signature fails, or None when it holds."""
    parts = [part.strip() for part in (headers.get("ShockAI-Signature") or "").split(",")]

    def field(name):
        values = [part[len(name) + 1 :] for part in parts if part.startswith(name + "=")]
        return values[0] if len(values) == 1 else ""

    timestamp, kid, signature = field("t"), field("kid"), field("ed25519")
    if not re.fullmatch(r"[0-9]{1,12}", timestamp) or not re.fullmatch(r"[A-Za-z0-9_-]{86}", signature):
        return "bad_signature"
    if abs(time.time() - int(timestamp)) > TOLERANCE_SECONDS:
        return "expired_signature"
    jwk = next((key for key in platform_keys if thumbprint(key) == kid), None)
    if jwk is None:
        refresh_platform_keys()
        jwk = next((key for key in platform_keys if thumbprint(key) == kid), None)
    if jwk is None:
        return "unknown_key"
    base = signature_base(timestamp, method, ORIGIN + target, headers.get("Idempotency-Key") or "", body)
    try:
        Ed25519PublicKey.from_public_bytes(_unb64url(jwk["x"])).verify(_unb64url(signature), base)
    except InvalidSignature:
        return "bad_signature"
    return None


def sign_outgoing(method: str, url: str, body: bytes) -> str:
    timestamp = str(int(time.time()))
    parts = urlsplit(url)
    # The origin, then the path and query as given.
    target = origin_of(url) + (parts.path or "/") + (f"?{parts.query}" if parts.query else "")
    signature = HOST_KEY.sign(signature_base(timestamp, method, target, "", body))
    return f"t={timestamp},kid={HOST_KID},ed25519={_b64url(signature)}"


# ---------------------------------------------------------------------------
# The manifest: manifest.json, with your build, and your public key unless it
# lists its own
# ---------------------------------------------------------------------------

MANIFEST = {
    **MANIFEST_FILE,
    "signing": MANIFEST_FILE["signing"] if PINNED_KEYS else {"callbackKeys": [HOST_JWK]},
    "build": {"hash": BUILD_HASH},
}

# ---------------------------------------------------------------------------
# Runs, kept on disk for 24 hours
# ---------------------------------------------------------------------------


def now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def run_path(run_id: str) -> Path:
    return DATA_DIR / "runs" / f"{run_id}.json"


def key_path(key: str) -> Path:
    return DATA_DIR / "keys" / f"{hashlib.sha256(key.encode()).hexdigest()}.json"


def save(stored: dict) -> None:
    # Write then rename, so a crash never leaves half a file.
    path = run_path(stored["run"]["id"])
    temporary = path.with_suffix(".json.tmp")
    temporary.write_text(json.dumps(stored))
    temporary.replace(path)


def load(run_id: str):
    if not RUN_ID.fullmatch(run_id) or not run_path(run_id).exists():
        return None
    return json.loads(run_path(run_id).read_text())


def sweep() -> None:
    # Under the lock, so a file being written is never taken for a crash's.
    with lock:
        for directory in ("runs", "keys"):
            for path in (DATA_DIR / directory).iterdir():
                try:
                    record = json.loads(path.read_text())
                    at = (record.get("run") or {}).get("updatedAt") or record.get("at") or ""
                    age = time.time() - datetime.strptime(at, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc).timestamp()
                except Exception:  # noqa: BLE001 - a write cut off by a crash: nothing to keep
                    age = KEEP_RUNS_S + 1
                if age > KEEP_RUNS_S:
                    path.unlink(missing_ok=True)


class _NoRedirects(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, *args, **kwargs):
        return None  # a 3xx is a failed delivery, never followed


_callbacks = urllib.request.build_opener(_NoRedirects)


def _mark_delivered(run_id: str) -> None:
    with lock:
        current = load(run_id)
        if current:
            save({**current, "delivered": True})


def deliver(stored: dict) -> None:
    """Send a run's callback until the platform takes it. A callback the host
    could not deliver is sent again on its next start, so it arrives at least
    once; a 4xx other than 408 or 429 is a refusal, and not sent again."""
    run, callback_url = stored["run"], stored["callbackUrl"]
    # The event is the run reaching its state, the same on every delivery,
    # so the platform can tell a repeat.
    body = json.dumps(
        {"id": f"evt_{run['id']}.{run['status']}", "createdAt": run["updatedAt"], "type": f"run.{run['status']}", "data": {"run": run}}
    ).encode()
    for attempt in range(CALLBACK_ATTEMPTS):
        if attempt > 0:
            time.sleep(2 ** (attempt - 1))
        # Signed afresh for every attempt, so a retry carries a current time.
        request = urllib.request.Request(
            callback_url,
            data=body,
            method="POST",
            headers={"content-type": "application/json", "shockai-signature": sign_outgoing("POST", callback_url, body)},
        )
        status = 0
        try:
            with _callbacks.open(request, timeout=30) as response:
                status = response.status
        except urllib.error.HTTPError as error:
            status = error.code
        except Exception:  # noqa: BLE001 - try again; the platform also polls GET /runs/{id}
            pass
        if 200 <= status < 300:
            return _mark_delivered(run["id"])
        if 400 <= status < 500 and status not in (408, 429):
            print("callback refused", run["id"], status, flush=True)
            return _mark_delivered(run["id"])
    print(f"callback not delivered after {CALLBACK_ATTEMPTS} attempts; the next start tries again", run["id"], flush=True)


def deliver_later(stored: dict) -> None:
    """Deliver a run's callback on a thread of its own, if it owes one:
    answered 202, to its request or a repeat of it, and not yet taken."""
    if stored.get("answeredAsync") and stored.get("callbackUrl") and not stored.get("delivered"):
        threading.Thread(target=deliver, args=(stored,), daemon=True).start()


def finish(stored: dict, result: dict) -> dict:
    with lock:
        current = load(stored["run"]["id"]) or stored
        if current["run"]["status"] in FINAL:
            return current  # cancelled meanwhile
        done = {**current, "run": {**current["run"], **result, "updatedAt": now()}}
        save(done)
        return done


def execute(stored: dict, case: dict) -> dict:
    """Run the model on this thread, the run's own, so a model that computes
    for minutes holds up neither /healthz nor the 202. run_model may be async."""
    started_at = now()
    with lock:
        current = load(stored["run"]["id"]) or stored
        if current["run"]["status"] in FINAL:
            return current  # cancelled before it began
        save({**current, "run": {**current["run"], "status": "running", "updatedAt": started_at}})
    try:
        until = time.monotonic() + MODEL_DELAY_S
        while time.monotonic() < until:
            pass  # MODEL_DELAY_MS: a model computing
        answer = run_model(case)
        if inspect.iscoroutine(answer):
            answer = asyncio.run(answer)
        done = finish(
            stored,
            {
                "status": "succeeded",
                "result": {
                    "specVersion": "0.1",
                    "caseId": case["caseId"],
                    "outputs": answer["outputs"],
                    "unavailable": answer["unavailable"],
                    "warnings": answer["warnings"],
                    "provenance": {
                        "modelId": MANIFEST["id"],
                        "modelVersion": MANIFEST["version"],
                        "caseId": case["caseId"],
                        "startedAt": started_at,
                        "finishedAt": now(),
                        "inputsUsed": answer["inputsUsed"],
                        "buildHash": BUILD_HASH,
                    },
                },
            },
        )
    except Exception as error:  # noqa: BLE001 - a model failure is the run's, not the host's
        print("run failed", stored["run"]["id"], repr(error), flush=True)
        done = finish(
            stored,
            {"status": "failed", "error": {"type": "provider_error", "code": "model_error", "message": "The model failed on this case."}},
        )
    # Answered 202, to its request or a repeat of it: call back.
    deliver_later(done)
    return done


# ---------------------------------------------------------------------------
# HTTP
# ---------------------------------------------------------------------------


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, *args) -> None:
        pass  # the JSONL log below replaces the default

    def _send(self, status: int, body) -> None:
        data = json.dumps(body).encode()
        self.send_response(status)
        self.send_header("content-type", "application/json")
        self.send_header("content-length", str(len(data)))
        if self.close_connection:
            self.send_header("connection", "close")
        self.end_headers()
        self.wfile.write(data)
        # What was asked and how it went; never a body.
        line = {"at": now(), "method": self.command, "path": self.path.split("?")[0], "status": status, "ms": int((time.time() - self._started) * 1000)}
        with lock, open(DATA_DIR / "requests.jsonl", "a") as log:
            log.write(json.dumps(line) + "\n")

    def _error(self, status: int, kind: str, code: str, message: str) -> None:
        self._send(status, {"error": {"type": kind, "code": code, "message": message}})

    def do_GET(self) -> None:
        self._handle()

    def do_POST(self) -> None:
        self._handle()

    def do_DELETE(self) -> None:
        self._handle()

    def _handle(self) -> None:
        self._started = time.time()
        try:
            path = self.path.split("?")[0]
            if not path.startswith(BASE_PATH + "/"):
                return self._error(404, "invalid_request", "not_found", "No such endpoint.")
            route = path[len(BASE_PATH) :]
            if route == "/healthz" and self.command == "GET":
                return self._send(200, {"ok": True})

            # A body comes with a Content-Length, of 16 MiB at most. Anything
            # else is refused before it is read, and the connection closed.
            lengths = [value.strip() for value in self.headers.get_all("content-length") or []]
            if self.headers.get("transfer-encoding"):
                self.close_connection = True
                return self._error(411, "invalid_request", "length_required", "Send the body with a Content-Length.")
            if len(lengths) > 1 or not re.fullmatch(r"[0-9]{1,20}", lengths[0] if lengths else "0"):
                self.close_connection = True
                return self._error(400, "invalid_request", "bad_content_length", "The Content-Length is not a number.")
            length = int(lengths[0]) if lengths else 0
            if length > MAX_BODY_BYTES:
                self.close_connection = True
                return self._error(413, "invalid_request", "too_large", "The body is over 16 MiB.")
            body = self.rfile.read(length) if length else b""
            problem = signature_problem(self.headers, self.command, self.path, body)
            if problem:
                return self._error(401, "authentication", problem, "The signature check failed.")

            if route == "/manifest" and self.command == "GET":
                return self._send(200, MANIFEST)
            if route == "/runs" and self.command == "POST":
                return self._create_run(body)

            match = re.fullmatch(r"/runs/([A-Za-z0-9_-]{1,128})", route)
            stored = load(match.group(1)) if match else None
            if match and stored is None:
                return self._error(404, "invalid_request", "run_not_found", "No such run.")
            if stored and self.command == "GET":
                return self._send(200, stored["run"])
            if stored and self.command == "DELETE":
                with lock:
                    stored = load(stored["run"]["id"])
                    if stored["run"]["status"] not in FINAL:
                        stored = {**stored, "run": {**stored["run"], "status": "cancelled", "updatedAt": now()}}
                        save(stored)
                return self._send(200, stored["run"])
            return self._error(404, "invalid_request", "not_found", "No such endpoint.")
        except Exception as error:  # noqa: BLE001
            print(repr(error), flush=True)
            self._error(500, "provider_error", "internal", "Something went wrong on the host.")

    def _create_run(self, body: bytes) -> None:
        idempotency_key = self.headers.get("Idempotency-Key") or ""
        if not idempotency_key:
            return self._error(400, "invalid_request", "missing_idempotency_key", "Send an Idempotency-Key.")

        with lock:
            # A repeated key is the original run, never a second one. Answered
            # 202, it calls back when it ends, whatever its first request is told.
            if key_path(idempotency_key).exists():
                stored = load(json.loads(key_path(idempotency_key).read_text())["runId"])
                if stored and stored["run"]["status"] in FINAL:
                    return self._send(200, stored["run"])
                if stored:
                    save({**stored, "answeredAsync": True})
                    return self._send(202, stored["run"])
            try:
                parsed = json.loads(body)
            except ValueError:
                return self._error(400, "invalid_request", "invalid_json", "The body is not JSON.")
            case = parsed.get("case") if isinstance(parsed, dict) else None
            if not isinstance(case, dict) or not isinstance(case.get("caseId"), str):
                return self._error(400, "invalid_request", "missing_case", "The body carries no case.")

            at = now()
            stored = {
                "run": {
                    "id": f"run_{uuid.uuid4()}",
                    "status": "queued",
                    "modelId": MANIFEST["id"],
                    "modelVersion": MANIFEST["version"],
                    "caseId": case["caseId"],
                    "createdAt": at,
                    "updatedAt": at,
                },
                "idempotencyKey": idempotency_key,
                "callbackUrl": parsed.get("callbackUrl") if isinstance(parsed.get("callbackUrl"), str) else None,
                "answeredAsync": False,
            }
            save(stored)
            key_path(idempotency_key).write_text(json.dumps({"runId": stored["run"]["id"], "at": at}))

        outcome = {}
        worker = threading.Thread(target=lambda: outcome.setdefault("done", execute(stored, case)), daemon=True)
        worker.start()
        worker.join(SYNC_BUDGET_S)
        if "done" in outcome:
            return self._send(200, outcome["done"]["run"])

        # Too slow to answer now: accept it, and call back when it is done.
        with lock:
            current = load(stored["run"]["id"]) or stored
            if current["run"]["status"] in FINAL:
                return self._send(200, current["run"])
            save({**current, "answeredAsync": True})
        self._send(202, current["run"])


def main() -> None:
    for path in (DATA_DIR / "runs").glob("*.json"):
        stored = load(path.stem)
        if not stored:
            continue
        # A run the host was working on when it stopped will not finish: say so.
        if stored["run"]["status"] not in FINAL:
            stored = finish(
                stored,
                {"status": "failed", "error": {"type": "provider_error", "code": "host_restarted", "message": "The host restarted during the run."}},
            )
        # And a callback the platform has not yet taken goes again.
        deliver_later(stored)
    sweep()

    def sweep_hourly():
        while True:
            time.sleep(3600)
            sweep()

    threading.Thread(target=sweep_hourly, daemon=True).start()
    refresh_platform_keys()
    server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler)
    print(f"listening on {server.server_address[1]}", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()