#!/usr/bin/env python3
"""Hugging Face Space entry point for the cascade human review (private Space sociate/cascade-review).

On start: downloads the web copy of the batch (images, product.json, review.html) from the private dataset BATCH_REPO and
the saved verdicts (reviews.json) from VERDICT_REPO, then runs review_server on port 7860. The Space disk is wiped on
restart, so after the workbook is rebuilt, reviews.json + cascade_report.xlsx are pushed back to VERDICT_REPO (at most
one commit per PUSH_EVERY seconds, plus once on shutdown). Needs the HF_TOKEN secret (read on BATCH_REPO, write on VERDICT_REPO).
"""
from __future__ import annotations

import os
import signal
import sys
import threading
import time
from functools import partial
from http.server import ThreadingHTTPServer
from pathlib import Path

from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, snapshot_download
from huggingface_hub.errors import EntryNotFoundError

import review_server

BATCH_REPO = os.environ.get("BATCH_REPO", "sociate/cascade-review-batch")
VERDICT_REPO = os.environ.get("VERDICT_REPO", "sociate/cascade-review-verdicts")
PUSH_EVERY = float(os.environ.get("PUSH_EVERY", "30"))
BATCH = Path(os.environ.get("BATCH_DIR", Path.home() / "data" / "batch"))
api = HfApi()


class Pusher:
    """Throttled push of reviews.json + cascade_report.xlsx to the verdict dataset."""

    def __init__(self):
        self.files, self.dirty, self.lock = None, threading.Event(), threading.Lock()
        threading.Thread(target=self._loop, daemon=True).start()

    def __call__(self, reviews: Path, xlsx: Path):
        self.files = (reviews, xlsx)
        self.dirty.set()

    def push(self):
        with self.lock:
            if not self.files:
                return
            self.dirty.clear()
            reviews, xlsx = self.files
            try:
                api.create_commit(VERDICT_REPO, repo_type="dataset", commit_message="Update human review verdicts",
                                  operations=[CommitOperationAdd("reviews.json", reviews.read_bytes() if reviews.exists() else b"[]"),
                                              CommitOperationAdd("cascade_report.xlsx", xlsx.read_bytes())])
                print(f"pushed verdicts to {VERDICT_REPO}", flush=True)
            except Exception as e:  # retried on the next change / tick
                self.dirty.set()
                print("push failed:", e, flush=True)

    def _loop(self):
        while True:
            self.dirty.wait()
            self.push()
            time.sleep(PUSH_EVERY)


def main() -> None:
    print(f"downloading {BATCH_REPO} ...", flush=True)
    snapshot_download(BATCH_REPO, repo_type="dataset", local_dir=BATCH, max_workers=16)
    try:
        hf_hub_download(VERDICT_REPO, "reviews.json", repo_type="dataset", local_dir=BATCH, force_download=True)
        print("restored reviews.json", flush=True)
    except EntryNotFoundError:
        print("no saved verdicts yet", flush=True)
    pusher = Pusher()
    review_server.Handler.store = store = review_server.Store(BATCH, debounce=2.0, after_build=pusher)
    store.dirty.set()  # build the workbook once at start (with any restored verdicts)

    def stop(*_):
        if store.excel["pending"]:
            time.sleep(3)  # let a rebuild in flight finish
        pusher.push()
        sys.exit(0)
    signal.signal(signal.SIGTERM, stop)
    srv = ThreadingHTTPServer(("0.0.0.0", 7860), partial(review_server.Handler, directory=str(BATCH)))
    print(f"serving {BATCH} ({len(store.reviews)} saved reviews) on :7860", flush=True)
    srv.serve_forever()


if __name__ == "__main__":
    main()