"""
Streamlit app for Image Captioning using BLIP-2 (or fallback BLIP).
Usage:
  - Set environment variable HUGGINGFACE_HUB_TOKEN if you need to access gated models.
  - On a machine without GPU memory, choose the smaller fallback model in the UI.
Notes:
  - BLIP-2 large variants require GPU and potentially accelerate/device_map configuration.
  - If you hit OOM errors, switch to "Salesforce/blip-image-captioning-base".
"""

import io
import os
import streamlit as st
from PIL import Image
import torch

# transformers pipeline import is lazy (import inside function) to speed cold-start of Streamlit UI
st.set_page_config(page_title="BLIP-2 Image Captioner", layout="centered")

# ---- UI ----
st.title("🖼️ Image Captioner — BLIP-2 / BLIP")
st.markdown(
    "Upload an image and get an automatic caption generated by a BLIP/BLIP-2 model. "
    "⚠️ Large BLIP-2 models require GPU memory — use the smaller fallback if needed."
)

col1, col2 = st.columns([3, 1])
with col2:
    st.write("")  # spacer
    model_choice = st.selectbox(
        "Model",
        options=[
            "Salesforce/blip2-opt-2.7b (BLIP-2, large, needs GPU)",
            "Salesforce/blip2-vicuna-7b-instruct (example; large)",
            "Salesforce/blip-image-captioning-base (BLIP fallback - small)"
        ],
        index=2
    )

with col1:
    uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
    st.markdown("Or try the example image below:")
    if st.button("Use example image"):
        # small example image from local bytes; replace with remote if you prefer
        example_path = "example.jpg"
        # create a simple placeholder if not available
        img = Image.new("RGB", (512, 384), color=(200, 200, 220))
        st.session_state["_example_image"] = img
        uploaded_file = None
        uploaded_image = img
    else:
        uploaded_image = None

prompt = st.text_input("Optional instruction/prompt (e.g. 'Describe briefly')", value="")

st.markdown("---")

# ---- Helper: load pipeline ----
@st.cache_resource(ttl=3600)
def load_caption_pipeline(model_id: str):
    """Load the image-to-text pipeline from transformers.
    Note: large BLIP-2 variants require GPU and may need accelerate/device_map settings.
    """
    try:
        from transformers import pipeline
    except Exception as e:
        raise RuntimeError("transformers is required. Install via requirements.txt") from e

    # set device to GPU if available
    if torch.cuda.is_available():
        device = 0
    else:
        device = -1

    # Create pipeline. The "image-to-text" pipeline wraps typical image captioning models
    # The pipeline will return a list of dicts with 'generated_text' (per HF docs).
    pipe = pipeline("image-to-text", model=model_id, device=device)
    return pipe

# ---- Generate caption ----
def generate_caption(image: Image.Image, model_choice_str: str, prompt_text: str):
    # map display string to HF model id
    if "blip-image-captioning-base" in model_choice_str:
        model_id = "Salesforce/blip-image-captioning-base"
    elif "blip2-opt-2.7b" in model_choice_str:
        model_id = "Salesforce/blip2-opt-2.7b"
    else:
        # default fallback
        model_id = "Salesforce/blip-image-captioning-base"

    st.info(f"Loading model: `{model_id}` — this may take a moment (and may need GPU).")
    try:
        pipe = load_caption_pipeline(model_id)
    except Exception as e:
        st.error(f"Failed to load model: {e}")
        st.stop()

    with st.spinner("Generating caption..."):
        # pipeline accepts either local PIL image or path
        input_for_pipe = image
        if prompt_text:
            # Many BLIP variants accept a "prompt" or "text" parameter — we put it into the pipeline call as text
            out = pipe(input_for_pipe, prompt=prompt_text, max_new_tokens=64)
        else:
            out = pipe(input_for_pipe, max_new_tokens=64)

    # pipeline returns list of dicts; use first result
    if isinstance(out, list) and len(out) > 0:
        # common key is "generated_text" for image-to-text pipeline
        caption = out[0].get("generated_text") or out[0].get("caption") or str(out[0])
    else:
        caption = str(out)

    return caption

# ---- Main interaction ----
if uploaded_file is not None:
    try:
        image = Image.open(uploaded_file).convert("RGB")
    except Exception as e:
        st.error(f"Could not open image: {e}")
        st.stop()
    st.image(image, caption="Uploaded image", use_column_width=True)
    if st.button("Generate caption"):
        caption = generate_caption(image, model_choice, prompt)
        st.success("Caption generated!")
        st.markdown("**Caption**")
        st.write(caption)
        # copy/download
        st.download_button("Download caption (.txt)", data=caption, file_name="caption.txt", mime="text/plain")
        st.button("Copy to clipboard")  # streamlit doesn't copy programmatically, so users can copy manually
else:
    if "_example_image" in st.session_state:
        img = st.session_state["_example_image"]
        st.image(img, caption="Example image", use_column_width=True)
        if st.button("Generate caption from example"):
            caption = generate_caption(img, model_choice, prompt)
            st.success("Caption generated!")
            st.write(caption)
            st.download_button("Download caption (.txt)", data=caption, file_name="caption.txt", mime="text/plain")
    else:
        st.info("Upload an image above or click 'Use example image'.")

# ---- Footnote / tips ----
st.markdown("---")
st.markdown(
    "Tips:\n"
    "- If you get an out-of-memory error with BLIP-2 large variants, switch to the smaller `Salesforce/blip-image-captioning-base` model.\n"
    "- To access private or gated models, set `HUGGINGFACE_HUB_TOKEN` as an environment variable before launching this app (or log in to `huggingface-cli`).\n"
)