"""Supervise independent UI and evaluation processes on the Space."""
from __future__ import annotations

import json
import hashlib
import os
from pathlib import Path
import signal
import shutil
import subprocess
import sys
import time

ROOT = Path(__file__).resolve().parents[1]


def ensure_bootstrap(root=ROOT, *, downloader=None):
    archive = root / "bootstrap/seed.tar.gz"
    if archive.exists():
        return
    from huggingface_hub import hf_hub_download
    expected = json.loads((root / "bootstrap/manifest.json").read_text())["sha256"]
    download = downloader or hf_hub_download
    source = Path(download(os.environ["HF_STATE_REPO"], f"bootstrap/{expected}.tar.gz",
                           repo_type="dataset", token=os.environ["HF_TOKEN"]))
    if hashlib.sha256(source.read_bytes()).hexdigest() != expected:
        raise ValueError("Remote bootstrap checksum mismatch")
    shutil.copy2(source, archive)


def main():
    os.chdir(ROOT)
    ensure_bootstrap()
    subprocess.run([sys.executable, "scripts/bootstrap_data.py"], check=True)
    (ROOT / ".cloud-state").mkdir(exist_ok=True)
    fallback = ROOT / ".cloud-state/ui-results"
    shutil.rmtree(fallback, ignore_errors=True)
    shutil.copytree(ROOT / "space/results", fallback)
    os.environ["TSFM_RESULTS_PATH"] = str(fallback)
    revision = json.loads((ROOT / "deployment.json").read_text())["source_revision"]
    os.environ["LIVEHOUSE_SOURCE_REVISION"] = revision
    processes = {}
    enabled = os.getenv("TSFM_EVALUATOR_ENABLED", "0") == "1"
    commands = {"web": [sys.executable, "cloud/web.py"]}
    if enabled:
        commands["worker"] = [sys.executable, "cloud/worker.py"]
    if os.getenv("LIVEHOUSE_ACCEPTANCE_ID"):
        commands["acceptance"] = [sys.executable, "cloud/acceptance.py"]
    completed = set()
    stopping = False

    def shutdown(*_):
        nonlocal stopping
        stopping = True
        for proc in processes.values():
            if proc.poll() is None:
                proc.terminate()

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)
    restart_at = {}
    try:
        while not stopping:
            for name, command in commands.items():
                proc = processes.get(name)
                if proc is not None and proc.poll() is not None:
                    if name == "acceptance":
                        print(f"Acceptance process finished with {proc.returncode}; no automatic rerun", flush=True)
                        completed.add(name)
                        del processes[name]
                        continue
                    print(f"{name} exited with {proc.returncode}; restart after 60 seconds", flush=True)
                    restart_at[name] = time.monotonic() + 60
                    del processes[name]
                if name not in processes and name not in completed and time.monotonic() >= restart_at.get(name, 0):
                    processes[name] = subprocess.Popen(command, cwd=ROOT)
            status = {"updated_at": time.time(), "source_revision": revision,
                      "evaluator_enabled": enabled,
                      "worker_running": "worker" in processes and processes["worker"].poll() is None}
            path = ROOT / ".cloud-state/supervisor.json"
            temporary = path.with_suffix(".tmp")
            temporary.write_text(json.dumps(status))
            temporary.replace(path)
            time.sleep(5)
    finally:
        shutdown()
        for proc in processes.values():
            try:
                proc.wait(timeout=35)
            except subprocess.TimeoutExpired:
                proc.kill()
                proc.wait()


if __name__ == "__main__":
    main()