import os
import requests
import gradio as gr
from PIL import Image
import torch
from transformers import BlipProcessor, BlipForConditionalGeneration

# 1) Устройство
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("🖥️ Using device:", device)

# 2) Лёгкая BLIP-модель (~240 MiB)
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model     = BlipForConditionalGeneration.from_pretrained(
    "Salesforce/blip-image-captioning-base"
).to(device)

# 3) TMDb API key
TMDB_KEY = os.environ.get("TMDB_API_KEY")
if not TMDB_KEY:
    raise RuntimeError("❌ TMDB_API_KEY is missing in your Secrets")
TMDB_SEARCH_URL = "https://api.themoviedb.org/3/search/movie"

def caption_and_search(image: Image.Image):
    print("➡️ Received request, image size:", image.size)
    # 4) caption
    inputs = processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        out_ids = model.generate(**inputs, max_new_tokens=30)
    caption = processor.decode(out_ids[0], skip_special_tokens=True).strip()
    print("🔖 Caption:", caption)
    # 5) TMDb search
    resp = requests.get(TMDB_SEARCH_URL, params={"api_key": TMDB_KEY, "query": caption}, timeout=10)
    if resp.status_code != 200:
        print("⚠️ TMDb returned", resp.status_code)
        return caption, [{"error": f"TMDb {resp.status_code}"}]
    data = resp.json().get("results", [])[:3]
    results = [{"title": m["title"], "url": f"https://www.themoviedb.org/movie/{m['id']}"} for m in data]
    print("✅ Returning", len(results), "results")
    return caption, results

# 6) Интерфейс с очередью
iface = gr.Interface(
    fn=caption_and_search,
    inputs=gr.Image(type="pil", label="Постер или кадр фильма"),
    outputs=[
        gr.Textbox(label="Auto‑caption"),
        gr.JSON(label="Top‑3 Movies (title + URL)")
    ],
    title="Fast Movie Finder (BLIP‑Base + TMDb)",
    description="Генерирует подпись и ищет топ‑3 фильма в TMDb"
).queue()  # <— включает очередь обработки долгих запросов

if __name__ == "__main__":
    # debug=True позволит вывести стектрейсы ошибок в браузер
    iface.launch(debug=True)