"""Single-process PhotoGuider runtime for the Hugging Face Space.

The release runs each extractor family in its own environment and process (envs/, workers.py).
Here every model lives in one environment and one process, loaded once and kept resident: the
extractor ops run in-process through the same load()/run(args) interface worker_serve.py uses, and
the explainer VLM is loaded once instead of once per call. On a Space (SPACE_ID set) the explainer
runs in bf16 instead of NF4, which the 48 GB ZeroGPU slice holds without bitsandbytes.
"""
import functools
import importlib
import os
import sys
import types
from pathlib import Path

ROOT = Path(os.environ.get("PHOTOGUIDER_ROOT", Path(__file__).resolve().parent)).resolve()
INF = ROOT / "inference"
LIB = INF / "_lib"

# one environment serves every role; require_env() only checks that the interpreter exists
for _role in ("MAIN", "LAION", "FGAESQ", "SAM", "EXPLAINER"):
    os.environ.setdefault(f"CROPPER_ENV_{_role}", sys.executable)
sys.path[:0] = [str(INF), str(LIB), str(LIB / "extractors")]

# FG-IAA ships `utils/` as a namespace package; transparent-background appends a directory holding
# a regular `utils.py` to sys.path, which would shadow it. Bind the name to FG-IAA's package first.
from env_config import asset  # noqa: E402
if "utils" not in sys.modules:
    _u = types.ModuleType("utils")
    _u.__path__ = [str(Path(asset("fgaesq_source", "envs/fgaesq/FGAesQ_Inference")) / "utils")]
    sys.modules["utils"] = _u

import pandas as pd  # noqa: E402

if os.environ.get("SPACE_ID"):
    # ZeroGPU: the Gradio process has no GPU and forks a worker per GPU call. transformers 5 asks
    # torch.cuda.is_current_stream_capturing() during a CPU forward, which initialises the CUDA
    # runtime in that GPU-less parent, and every worker forked afterwards then sees no GPU
    # ("No CUDA GPUs are available"). We never capture CUDA graphs, so answer without asking.
    import transformers.utils.import_utils as _tiu
    _tiu.is_cuda_stream_capturing = lambda: False
import _shared  # noqa: E402
import feature_extractors  # noqa: E402
import feat_saliency as FS  # noqa: E402
import explain_vlm  # noqa: E402
import explain as EX  # noqa: E402
import photoguider  # noqa: E402


class InProcessPool:
    """ExtractorPool interface; ops run in this process with their models resident."""

    def __init__(self, verbose=False):
        self.mods = {}

    def module(self, op):
        if op not in self.mods:
            if op not in feature_extractors.OP_ENV:
                raise ValueError(f"unknown extractor op '{op}'")
            m = importlib.import_module(op)
            m.load()
            self.mods[op] = m
        return self.mods[op]

    def run_op(self, op, args=None):
        self.module(op).run(args)
        return True

    def close(self):                               # models stay resident across requests
        pass


def load_vlm_bf16(adapter_dir):
    """explain_vlm.load_vlm without the NF4 quantization."""
    import torch
    from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
    from peft import PeftModel
    adapter_dir = Path(adapter_dir)
    torch.manual_seed(1234)
    proc = AutoProcessor.from_pretrained(explain_vlm.BASE_VLM, revision=explain_vlm.revision("qwen"))
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        explain_vlm.BASE_VLM, revision=explain_vlm.revision("qwen"), torch_dtype=torch.bfloat16).to("cuda")
    model = PeftModel.from_pretrained(model, str(adapter_dir), is_trainable=False).eval()
    return {"proc": proc, "model": model,
            "bad": explain_vlm.cjk_bad_words_ids(proc.tokenizer, adapter_dir / "cjk_ids_cache.json")}


# ── two speed-ups that leave every number unchanged ──
# 1. The ~1,500 candidate boxes get their hand-crafted features in a per-row Python loop
#    (consultant_features.stage_c_features); the rows are independent, so large pools are split
#    across forked CPU workers and concatenated back in order.
import multiprocessing  # noqa: E402
import infer  # noqa: E402

_WORKERS = int(os.environ.get("PHOTOGUIDER_WORKERS") or max(1, min(8, (os.cpu_count() or 2) - 1)))   # 1 = serial
_serial_features = infer.af_chain_features


def _chunk_features(args):
    dfp, maps = args
    return infer.M02.stage_c_features(dfp.reset_index(drop=True), maps)


def parallel_features(dfp, vmix_2d, v1_2d, mask_np, min_rows=200):
    if len(dfp) < min_rows or _WORKERS == 1:
        return _serial_features(dfp, vmix_2d, v1_2d, mask_np)
    maps = infer.maps_for(dfp.photo_id.iloc[0], vmix_2d, v1_2d, mask_np)
    dfp = dfp.reset_index(drop=True)
    bounds = [round(i * len(dfp) / _WORKERS) for i in range(_WORKERS + 1)]
    # The ZeroGPU worker is a daemonic process, which may not have children; it is forked fresh for
    # every call, so lift that flag here and fork a pool for this call only (numpy only, never CUDA).
    multiprocessing.current_process()._config["daemon"] = False
    with multiprocessing.get_context("fork").Pool(_WORKERS) as pool:
        parts = pool.map(_chunk_features, [(dfp.iloc[a:b], maps) for a, b in zip(bounds, bounds[1:]) if b > a])
    out = pd.concat(parts, ignore_index=True)
    assert len(out) == len(dfp), "parallel candidate features lost rows"
    return out


infer.af_chain_features = parallel_features


# 2. The SAM image encoder sees the same full image several times per request (mean embedding,
#    patch norms, the decoder's set_image, both extractor passes) and the same candidate batches
#    in both passes. Memoise its forward on the exact input tensor; a hit returns the stored output.
def memoise_sam_encoder(sam, max_entries=48):
    import hashlib
    enc, cache = sam.image_encoder, {}
    forward = enc.forward

    def cached_forward(x):
        key = (tuple(x.shape), str(x.dtype), hashlib.sha1(x.detach().cpu().numpy().tobytes()).hexdigest())
        if key not in cache:
            if len(cache) >= max_entries:
                cache.pop(next(iter(cache)))
            cache[key] = forward(x)
        return cache[key]

    enc.forward = cached_forward
    return sam


_get_sam = _shared.get_sam


def get_sam_cached():
    if "sam" not in _shared._CACHE:
        memoise_sam_encoder(_get_sam())
    return _shared._CACHE["sam"]


_shared.get_sam = get_sam_cached

explain_vlm.load_vlm = functools.cache(load_vlm_bf16 if os.environ.get("SPACE_ID") else explain_vlm.load_vlm)
EX.write_explanation = lambda crop_path, prompt: explain_vlm.generate_explanation(crop_path, prompt, str(EX.ADAPTER_DIR))
photoguider.ExtractorPool = InProcessPool


def preload(session, explain=True):
    """Load every model once (extractors, saliency, explainer) so requests only run inference."""
    for op in feature_extractors.OP_ENV:
        session.pool.module(op)
    FS._model()
    import feat_clip_refine
    feat_clip_refine._load()
    if "rem" not in FS._STATE:
        from transparent_background import Remover
        FS._STATE["rem"] = Remover()
    if explain:
        explain_vlm.load_vlm(str(EX.ADAPTER_DIR))