"""Hugging Face entry point: custom studio UI, Gradio queue, free ZeroGPU."""

import os

os.environ["PIXTEXT_GRADIO"] = "1"
os.environ.setdefault("PIXTEXT_CPU_THREADS", "1")
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")

# Import before torch so ZeroGPU can enable its CUDA staging mode.
import spaces

# Only gpu_refine needs a GPU; keep validation and session state in the web process.
spaces.disable_gradio_auto_wrap()

import gradio as gr
import hashlib
import json
import time
import uuid
from pathlib import Path

import app as studio

from zerogpu_attention import enable_sdpa

enable_sdpa()


def load_models():
    import torch
    from pixtext import PixTextPipeline
    from huggingface_hub import snapshot_download

    torch.set_num_threads(int(os.environ["PIXTEXT_CPU_THREADS"]))
    if not torch.cuda.is_available() and not spaces.config.Config.zero_gpu:
        return {}
    pipelines = {}
    for variant, name in (("meanflow", "PixText-MeanFlow"), ("teacher", "PixText-Teacher")):
        if os.environ.get("PIXTEXT_MODEL_ROOT"):
            package = Path(os.environ["PIXTEXT_MODEL_ROOT"]) / name
        else:
            package = Path(snapshot_download("ATH-MaaS/PixText", allow_patterns=[f"models/{name}/*"],
                                             local_dir=studio.RUNTIME / "model-packages")) / "models" / name
        pipeline = PixTextPipeline.from_pretrained(package, device="cpu", dtype=torch.bfloat16)
        pipeline.model = pipeline.model.to("cuda")
        pipeline.device = torch.device("cuda")
        pipelines[variant] = pipeline
    return pipelines


# Stage both models before serving requests, as required by ZeroGPU.
PIPELINES = load_models()


def estimate_duration(directory: str, size: list[int], jid: str, payload: str):
    import math
    request = json.loads(payload)
    count = sum(row.get("enabled", True) for row in request["regions"])
    return max(20, min(90, 10 + math.ceil(request["steps"] * max(1, count / 8))))


@spaces.GPU(duration=estimate_duration)
def gpu_refine(directory: str, size: list[int], jid: str, payload: str):
    print("PixText GPU worker: starting", flush=True)
    started = time.perf_counter()
    request = studio.RefineRequest.model_validate_json(payload)
    studio.PIPELINE = PIPELINES.get(request.variant)
    studio.PIPELINE_VARIANT = request.variant if studio.PIPELINE is not None else None
    progress = gr.Progress()

    def report(**values):
        message = values.get("message", "Refining")
        print("PixText GPU:", message, flush=True)
        progress(values.get("progress", 0), desc=message)

    result = studio.refine({"directory": Path(directory), "size": size}, jid, request, report)
    result["elapsed"] = round(time.perf_counter() - started, 2)
    print("PixText GPU worker: complete", result["elapsed"], flush=True)
    return result


@studio.app.api(name="refine_image", concurrency_limit=1, concurrency_id="pixtext-gpu")
def refine_image(sid: str, payload: str) -> dict:
    request = studio.RefineRequest.model_validate_json(payload)
    session = studio.validate_refine(sid, request)
    canonical = request.model_dump_json()
    key = hashlib.sha256(canonical.encode()).hexdigest()
    with studio.LOCK:
        previous = session.get("last_gradio_result")
        if previous and previous["key"] == key:
            job = studio.JOBS.get(previous["jid"])
            if job and job["status"] == "done":
                return {"job_id": job["id"], "result": job["result"]}
        jid = uuid.uuid4().hex
        studio.JOBS[jid] = {"id": jid, "session": sid, "kind": "refine", "status": "running",
                            "progress": 0, "message": "Waiting for GPU", "created": time.time()}
        session["job"] = jid
    try:
        result = gpu_refine(str(session["directory"]), list(session["size"]), jid, canonical)
        studio.update_job(jid, status="done", progress=1, message="Complete", result=result)
        session["last_gradio_result"] = {"key": key, "jid": jid}
        return {"job_id": jid, "result": result}
    except Exception:
        studio.update_job(jid, status="error", message="Refinement could not finish. Please retry.")
        raise


if __name__ == "__main__":
    studio.app.launch(server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
                      server_port=int(os.environ.get("PIXTEXT_PORT", "7860")), show_error=False, ssr_mode=False)