#!/usr/bin/env python3
from __future__ import annotations

import argparse
import datetime
import hashlib
import json
import os
import time
import urllib.error
import urllib.request
from pathlib import Path

ROOT = Path("/data/video-pipeline/HaTramAudio/project/028-Nguoi-Duoc-Goi-Ten-Cuoi-Cung")
PROJECT_ID = "028-Nguoi-Duoc-Goi-Ten-Cuoi-Cung"
RUN_ID = "run-20260722T063856Z-87753fd5"
OWNER = "Levy"
CONFIG = {
    "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",
    },
}
GATES = {"footage": "footage", "final": "final_render", "transcode": "transcode"}
DEPENDENCIES = {
    "footage": {"tts", "layout"},
    "final": {"footage", "intro"},
    "transcode": {"final_render"},
}


def now() -> str:
    return datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z")


def sha(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for block in iter(lambda: f.read(8 * 1024 * 1024), b""):
            h.update(block)
    return h.hexdigest()


def guard() -> None:
    lock = json.loads((ROOT / ".ownership-lock.json").read_text(encoding="utf-8"))
    for key, value in {"project_id": PROJECT_ID, "run_id": RUN_ID, "owner": OWNER}.items():
        if lock.get(key) != value:
            raise RuntimeError(f"ownership mismatch: {key}")


def atomic_json(path: Path, value: dict) -> None:
    guard()
    path.parent.mkdir(parents=True, exist_ok=True)
    temp = path.with_name("." + path.name + ".tmp")
    with temp.open("w", encoding="utf-8") as f:
        json.dump(value, f, ensure_ascii=False, indent=2)
        f.write("\n")
        f.flush()
        os.fsync(f.fileno())
    os.replace(temp, path)


def get_json(url: str) -> dict:
    request = urllib.request.Request(url, headers={"Accept": "application/json"})
    with urllib.request.urlopen(request, timeout=30) as response:
        return json.loads(response.read())


def post_json(url: str, payload: dict) -> tuple[int, dict]:
    request = urllib.request.Request(
        url,
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json", "Accept": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=60) as response:
        return response.status, json.loads(response.read())


def normalize_jobs(value: object) -> list[dict]:
    if isinstance(value, dict) and isinstance(value.get("jobs"), list):
        jobs = value["jobs"]
    elif isinstance(value, list):
        jobs = value
    else:
        raise RuntimeError("unknown list-jobs response schema")
    if not all(isinstance(item, dict) for item in jobs):
        raise RuntimeError("invalid job item schema")
    return jobs


def endpoint_is_idle(base_url: str, jobs_path: str) -> bool:
    health = get_json(base_url + "/health")
    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(get_json(base_url + jobs_path))
    return not any(item.get("status") in {"queued", "running"} for item in jobs)


def select_idle_endpoint(pool: dict, wait_seconds: int = 30, timeout_seconds: int = 14400) -> str:
    deadline = time.monotonic() + timeout_seconds
    while True:
        for base_url in pool["base_urls"]:
            try:
                if endpoint_is_idle(base_url, pool["jobs_path"]):
                    return base_url
            except Exception as exc:
                print(f"probe base_url={base_url} fail_closed={type(exc).__name__}", flush=True)
        if time.monotonic() >= deadline:
            raise TimeoutError("no idle render endpoint before deadline")
        time.sleep(wait_seconds)


def input_hashes(bindings: dict) -> dict[str, str]:
    result = {}
    for key, value in bindings.items():
        if key.startswith("source_") and key.endswith("_sha256") and key != "source_canon_sha256":
            result[key] = value
    if not result:
        raise RuntimeError("request bindings contain no input hashes")
    return result


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("stage", choices=sorted(CONFIG))
    parser.add_argument("--poll-seconds", type=int, default=5)
    parser.add_argument("--max-polls", type=int, default=7200)
    args = parser.parse_args()
    guard()
    manifest = json.loads((ROOT / "script" / "project-manifest.json").read_text(encoding="utf-8"))
    gate = GATES[args.stage]
    if manifest.get("steps", {}).get(gate) != "pending":
        raise RuntimeError(f"remote gate is not pending: {gate}")
    incomplete = {
        dependency: manifest.get("steps", {}).get(dependency, "pending")
        for dependency in DEPENDENCIES[args.stage]
        if manifest.get("steps", {}).get(dependency) != "completed"
    }
    if incomplete:
        raise RuntimeError(f"remote gate dependencies are incomplete: {incomplete}")
    lease_path = ROOT / "script" / "gate-leases" / f"{gate}.json"
    if not lease_path.is_file():
        raise RuntimeError(f"remote gate lease is missing: {gate}")
    lease = json.loads(lease_path.read_text(encoding="utf-8"))
    expires_at = datetime.datetime.fromisoformat(lease["expires_at"].replace("Z", "+00:00"))
    if lease.get("status") != "active" or lease.get("run_id") != RUN_ID or expires_at <= datetime.datetime.now(datetime.timezone.utc):
        raise RuntimeError(f"remote gate lease is not active for current run: {gate}")
    pool = CONFIG[args.stage]
    route = pool["jobs_path"]
    work = ROOT / "work" / args.stage / RUN_ID
    request_path = work / "request.json"
    bindings_path = work / "request-bindings.json"
    intent_path = work / "job-intent.json"
    terminal_path = work / "terminal-response.json"
    payload = json.loads(request_path.read_text(encoding="utf-8"))
    bindings = json.loads(bindings_path.read_text(encoding="utf-8"))
    request_hash = sha(request_path)
    if bindings.get("request_sha256") != request_hash:
        raise RuntimeError("request binding hash mismatch")

    intent = None
    intent_existed = intent_path.exists()
    if intent_path.exists():
        intent = json.loads(intent_path.read_text(encoding="utf-8"))
        if intent.get("request_sha256") != request_hash or intent.get("project_id") != PROJECT_ID or intent.get("run_id") != RUN_ID:
            raise RuntimeError("stale/foreign job intent")
        if intent.get("stage") != args.stage or intent.get("jobs_path") != route:
            raise RuntimeError("intent stage/jobs path mismatch")
        if intent.get("base_url") not in pool["base_urls"]:
            raise RuntimeError("intent base_url is not in the stage pool")
    if intent and intent.get("job_id"):
        base = intent["base_url"]
        job_id = intent["job_id"]
        print(f"resume base_url={base} job_id={job_id}", flush=True)
    elif intent_existed:
        raise RuntimeError("manual_reconciliation_required: intent has base_url but no job_id; inspect only the pinned replica")
    else:
        base = select_idle_endpoint(pool)
        if not endpoint_is_idle(base, route):
            raise RuntimeError("selected endpoint became busy before intent/POST")
        output_path = bindings.get("output")
        if not isinstance(output_path, str) or not output_path.startswith(str(ROOT)):
            raise RuntimeError("invalid output path binding")
        creating = {
            "schema_version": 1,
            "project_id": PROJECT_ID,
            "run_id": RUN_ID,
            "stage": args.stage,
            "status": "creating",
            "verified": False,
            "base_url": base,
            "jobs_path": route,
            "request_path": str(request_path),
            "request_sha256": request_hash,
            "input_sha256": input_hashes(bindings),
            "output_path": output_path,
            "created_at": now(),
        }
        atomic_json(intent_path, creating)
        try:
            status_code, response = post_json(base + route, payload)
        except Exception as exc:
            print(f"ambiguous_post_transport_error base_url={base} error={type(exc).__name__}", flush=True)
            raise RuntimeError("manual_reconciliation_required: POST result ambiguous; do not fallback or resubmit") from exc
        job_id = response.get("job_id") or response.get("id")
        if not job_id:
            print(f"ambiguous_post_response_missing_job_id base_url={base}", flush=True)
            raise RuntimeError("manual_reconciliation_required: POST response lacks job_id; do not fallback or resubmit")
        creating.update({"status": "submitted", "http_status": status_code, "job_id": job_id, "submitted_at": now()})
        atomic_json(intent_path, creating)
        print(f"accepted base_url={base} job_id={job_id}", flush=True)

    status_url = base + route + "/" + str(job_id)
    last = None
    for poll in range(1, args.max_polls + 1):
        try:
            job = get_json(status_url)
        except (urllib.error.URLError, TimeoutError) as exc:
            print(f"poll={poll} transient={type(exc).__name__}", flush=True)
            time.sleep(args.poll_seconds)
            continue
        status = job.get("status")
        if status != last or poll % 12 == 0:
            print(f"poll={poll} status={status}", flush=True)
            last = status
        intent = json.loads(intent_path.read_text(encoding="utf-8"))
        intent.update({"status": status, "last_polled_at": now(), "poll_count": poll})
        atomic_json(intent_path, intent)
        if status in {"completed", "failed", "cancelled"}:
            atomic_json(terminal_path, job)
            intent.update({"status": status, "terminal_response_path": str(terminal_path), "terminal_at": now(), "verified": status == "completed" and job.get("verification", {}).get("verified") is True})
            atomic_json(intent_path, intent)
            print(json.dumps({"job_id": job_id, "status": status, "server_verified": job.get("verification", {}).get("verified")}, ensure_ascii=False), flush=True)
            return 0 if status == "completed" and job.get("verification", {}).get("verified") is True else 1
        time.sleep(args.poll_seconds)
    print(json.dumps({"job_id": job_id, "status": last, "terminal": False, "action": "continue_poll_same_job"}), flush=True)
    return 2


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