import time
import hashlib

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

MODEL_NAME = "Salesforce/blip-image-captioning-base"

# LAZY: not loaded at startup. These stay None until the first request, then
# live in the MAIN (CPU) process so they persist across later requests.
# Only run_inference runs in the stateless GPU process.
processor = None
model = None

caption_cache = {}


def load_model():
    """Load the model on first use (CPU, main process) and reuse it after."""
    global processor, model
    start_load = time.time()
    if processor is None or model is None:
        processor = BlipProcessor.from_pretrained(MODEL_NAME)
        model = BlipForConditionalGeneration.from_pretrained(MODEL_NAME)
        model.eval()
    load_time = time.time() - start_load   # >0 on the first request, ~0 after
    return load_time


@spaces.GPU
def run_inference(model, inputs):
    """The ONLY function that touches the GPU."""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    print("inference device:", next(model.parameters()).device)  # confirm GPU
    with torch.no_grad():
        output = model.generate(**inputs)
    return output.to("cpu")


def _cache_key(image):
    return hashlib.sha256(image.tobytes()).hexdigest()


def generate_caption(image, 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(image)
    preprocess_time = time.time() - start_preprocess

    # ---- CPU: cache check (no GPU on a hit) ----
    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: lazy load (first request only) + tokenize ----
    load_time = load_model()
    inputs = processor(image, return_tensors="pt")

    # ---- GPU: the only step that holds the GPU ----
    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),
        "approx_gpu_time": round(inference_time, 3),
        "cache_hit": False,
        "model": MODEL_NAME,
    }

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

    return result