import time
import hashlib

import spaces
import torch
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration

MODEL_VARIANTS = {
    "base": "Salesforce/blip-image-captioning-base",
    "large": "Salesforce/blip-image-captioning-large",
}


_loaded_models = {}   # model_name -> (processor, model), kept on CPU
caption_cache = {}    # cache_key  -> result dict


def _load_model(model_variant):
    """Load a BLIP variant once (on CPU) and reuse it on later requests."""
    if model_variant not in MODEL_VARIANTS:
        raise ValueError(f"model_variant must be one of {list(MODEL_VARIANTS)}")
    model_name = MODEL_VARIANTS[model_variant]

    start_load = time.time()
    if model_name not in _loaded_models:
        processor = BlipProcessor.from_pretrained(model_name)
        model = BlipForConditionalGeneration.from_pretrained(model_name)
        model.eval()
        _loaded_models[model_name] = (processor, model)
    load_time = time.time() - start_load   # ~0 on every request after the first

    processor, model = _loaded_models[model_name]
    return processor, model, load_time, model_name


@spaces.GPU
def run_inference(model, inputs):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        output = model.generate(**inputs)
    return output.to("cpu")


def _cache_key(model_variant, image):
    # Stable across process restarts (unlike the built-in hash(), which is
    # salted per process) and includes the variant so Base/Large cache apart.
    digest = hashlib.sha256(image.tobytes()).hexdigest()
    return f"{model_variant}:{digest}"


def generate_caption(image, model_variant="base", use_cache=True):
    start_total = time.time()

    # ---- CPU: preprocessing ----
    start_preprocess = time.time()
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    image = image.convert("RGB")
    key = _cache_key(model_variant, image)
    preprocess_time = time.time() - start_preprocess

    # ---- CPU: cache check. On a hit, the GPU is NEVER allocated. ----
    if use_cache and key in caption_cache:
        cached = caption_cache[key].copy()
        cached["cache_hit"] = True
        cached["total_time"] = round(time.time() - start_total, 3)
        cached["preprocess_time"] = round(preprocess_time, 3)
        cached["model_load_time"] = 0.0
        cached["inference_time"] = 0.0
        cached["postprocess_time"] = 0.0
        cached["approx_gpu_time"] = 0.0
        return cached

    # ---- CPU: load (or reuse) the model and tokenize the image ----
    processor, model, load_time, model_name = _load_model(model_variant)
    inputs = processor(image, return_tensors="pt")

    start_inference = time.time()
    output = run_inference(model, inputs)
    inference_time = time.time() - start_inference

    # ---- CPU: postprocessing ----
    start_postprocess = time.time()
    caption = processor.decode(output[0], skip_special_tokens=True)
    postprocess_time = time.time() - start_postprocess

    total_time = time.time() - start_total

    result = {
        "caption": caption,
        "total_time": round(total_time, 3),
        "preprocess_time": round(preprocess_time, 3),
        "model_load_time": round(load_time, 3),
        "inference_time": round(inference_time, 3),
        "postprocess_time": round(postprocess_time, 3),
        # NOTE: inference_time is measured around run_inference from the main
        # process, so it also includes the CPU<->GPU transfer and the ZeroGPU
        # scheduling/fork overhead. Treat approx_gpu_time as an UPPER BOUND on
        # true GPU time, not a precise GPU measurement.
        "approx_gpu_time": round(inference_time, 3),
        "cache_hit": False,
        "model_variant": model_variant,
        "model": model_name,
        "use_cache": use_cache,
    }

    if use_cache:
        caption_cache[key] = result.copy()

    return result

_load_model("base")