#!/usr/bin/env python3
"""Fail-closed render-pool helpers for Huyen An Audio project runners."""

from __future__ import annotations

import hashlib
import json
import os
import time
import urllib.request
from pathlib import Path
from typing import Any, Callable

RENDER_POOLS = {
    "footage": {
        "base_urls": [
            "http://192.168.1.104:8022",
            "http://192.168.1.104:8025",
            "http://192.168.1.104:8028",
            "http://192.168.1.104:8031",
        ],
        "jobs_path": "/v1/footage-render-huyenan/jobs",
    },
    "final": {
        "base_urls": [
            "http://192.168.1.104:8023",
            "http://192.168.1.104:8026",
            "http://192.168.1.104:8029",
            "http://192.168.1.104:8032",
        ],
        "jobs_path": "/v1/huyenan-render-final/jobs",
    },
    "transcode": {
        "base_urls": [
            "http://192.168.1.104:8024",
            "http://192.168.1.104:8027",
            "http://192.168.1.104:8030",
            "http://192.168.1.104:8033",
        ],
        "jobs_path": "/v1/huyenan-transcode-upload/jobs",
    },
}


def get_json(url: str, timeout: float = 15.0) -> Any:
    with urllib.request.urlopen(url, timeout=timeout) as response:
        return json.load(response)


def normalize_jobs(response: Any) -> list[dict[str, Any]]:
    if isinstance(response, dict) and isinstance(response.get("jobs"), list):
        rows = response["jobs"]
    elif isinstance(response, list):
        rows = response
    else:
        raise RuntimeError("Unknown list-jobs response schema")
    if not all(isinstance(row, dict) for row in rows):
        raise RuntimeError("Invalid item in list-jobs response")
    return rows


def endpoint_is_idle(
    base_url: str,
    jobs_path: str,
    fetch: Callable[[str, float], Any] = get_json,
    timeout: float = 15.0,
) -> bool:
    health = fetch(base_url + "/health", timeout)
    if not isinstance(health, dict):
        return False
    if health.get("status") != "ok" or health.get("h264_nvenc") is not True:
        return False
    running = health.get("running")
    if not isinstance(running, list) or running:
        return False
    jobs = normalize_jobs(fetch(base_url + jobs_path, timeout))
    return not any(row.get("status") in {"queued", "running"} for row in jobs)


def select_idle_endpoint(
    stage: str,
    fetch: Callable[[str, float], Any] = get_json,
    wait_timeout: float = 4 * 60 * 60,
    interval: float = 30.0,
) -> tuple[str, str]:
    pool = RENDER_POOLS.get(stage)
    if pool is None:
        raise ValueError(f"Unknown render stage: {stage}")
    deadline = time.monotonic() + wait_timeout
    while True:
        for base_url in pool["base_urls"]:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError("No idle render endpoint before deadline")
            try:
                if endpoint_is_idle(
                    base_url,
                    pool["jobs_path"],
                    fetch=fetch,
                    timeout=min(15.0, remaining),
                ):
                    return base_url, pool["jobs_path"]
            except Exception:
                # Transport errors and unknown schemas fail closed for this replica.
                continue
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise TimeoutError("No idle render endpoint before deadline")
        time.sleep(min(interval, remaining))


def canonical_payload_sha256(payload: dict[str, Any]) -> str:
    encoded = json.dumps(
        payload,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def atomic_write_json(path: Path, data: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(path.name + ".part")
    temporary.write_text(
        json.dumps(data, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    os.replace(temporary, path)


def write_creating_intent(
    path: Path,
    *,
    project_id: str,
    run_id: str,
    stage: str,
    base_url: str,
    jobs_path: str,
    payload: dict[str, Any],
    input_sha256: dict[str, str],
    output_path: str,
    created_at: str,
) -> dict[str, Any]:
    if path.exists():
        raise RuntimeError("Render intent already exists; reconcile it before selection")
    pool = RENDER_POOLS.get(stage)
    if pool is None or base_url not in pool["base_urls"] or jobs_path != pool["jobs_path"]:
        raise RuntimeError("Selected endpoint does not belong to the requested stage")
    intent = {
        "schema_version": 1,
        "project_id": project_id,
        "run_id": run_id,
        "stage": stage,
        "status": "creating",
        "base_url": base_url,
        "jobs_path": jobs_path,
        "request_sha256": canonical_payload_sha256(payload),
        "input_sha256": input_sha256,
        "output_path": output_path,
        "created_at": created_at,
    }
    atomic_write_json(path, intent)
    return intent


def mark_submitted(path: Path, job_id: str, submitted_at: str) -> dict[str, Any]:
    intent = json.loads(path.read_text(encoding="utf-8"))
    if intent.get("status") != "creating" or not intent.get("base_url"):
        raise RuntimeError("Intent is not in creating state")
    intent.update({"status": "submitted", "job_id": job_id, "submitted_at": submitted_at})
    atomic_write_json(path, intent)
    return intent


def pinned_job_url(intent: dict[str, Any]) -> str:
    if not intent.get("base_url") or not intent.get("jobs_path") or not intent.get("job_id"):
        raise RuntimeError("Pinned base_url + jobs_path + job_id are required")
    return intent["base_url"].rstrip("/") + intent["jobs_path"] + "/" + intent["job_id"]


def recheck_selected_endpoint(
    base_url: str,
    jobs_path: str,
    fetch: Callable[[str, float], Any] = get_json,
) -> None:
    if not endpoint_is_idle(base_url, jobs_path, fetch=fetch):
        raise RuntimeError("Selected render endpoint is no longer idle")


if __name__ == "__main__":
    raise SystemExit("Import this helper from a project runner; it does not submit jobs itself")