import gradio as gr
import torch
from PIL import Image
from transformers import AutoProcessor, Florence2ForConditionalGeneration
from deep_translator import GoogleTranslator
import json
import os
import tempfile

# ── Model ─────────────────────────────────────────────────────────────────────
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
print(f"Loading model on {device}...")

MODEL_ID = "florence-community/Florence-2-base"

model = Florence2ForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch_dtype,
).to(device).eval()

processor = AutoProcessor.from_pretrained(MODEL_ID)
print("Model ready!")

translator = GoogleTranslator(source="en", target="km")


# ── Caption ───────────────────────────────────────────────────────────────────
def caption_one(img: Image.Image):
    img = img.convert("RGB")
    inputs = processor(
        text="<MORE_DETAILED_CAPTION>",
        images=img,
        return_tensors="pt"
    ).to(device, torch_dtype)

    with torch.no_grad():
        ids = model.generate(
            input_ids=inputs["input_ids"],
            pixel_values=inputs["pixel_values"],
            max_new_tokens=1024,
            early_stopping=False,
            do_sample=False,
            num_beams=3,
        )

    raw = processor.batch_decode(ids, skip_special_tokens=False)[0]
    parsed = processor.post_process_generation(
        raw,
        task="<MORE_DETAILED_CAPTION>",
        image_size=(img.width, img.height),
    )
    english = parsed["<MORE_DETAILED_CAPTION>"].strip()
    try:
        khmer = translator.translate(english)
    except Exception:
        khmer = english
    return english, khmer


def process_batch(files, prefix: str, progress=gr.Progress(track_tqdm=True)):
    if not files:
        yield "[]", None, [], "0 / 0"
        return

    prefix = prefix.strip().rstrip("/")
    results, rows = [], []
    total = len(files)

    # Write to a fixed temp file path so we can update it incrementally
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".json", mode="w", encoding="utf-8")
    tmp_path = tmp.name
    tmp.close()

    for i, fp in enumerate(progress.tqdm(files, desc="Captioning")):
        path = fp if isinstance(fp, str) else fp.name
        try:
            img = Image.open(path).convert("RGB")
        except Exception:
            continue

        filename = os.path.basename(path)
        key = f"{prefix}/{filename}" if prefix else filename

        en, kh = caption_one(img)
        results.append({"images": key, "caption_kh": kh, "caption_en": en})
        rows.append([key, kh, en])

        # Update JSON file and UI after every image
        json_text = json.dumps(results, ensure_ascii=False, indent=2)
        with open(tmp_path, "w", encoding="utf-8") as f:
            f.write(json_text)

        status = f"✅ {i+1} / {total} done"
        yield json_text, gr.update(value=tmp_path, visible=True), rows, status

    status = f"🎉 All {total} images captioned!"
    yield json_text, gr.update(value=tmp_path, visible=True), rows, status


# ── UI ────────────────────────────────────────────────────────────────────────
CSS = """
#title { text-align: center; }
#json_box textarea { font-family: 'Courier New', monospace; font-size: 12px; }
#status { text-align: center; font-size: 1.1em; font-weight: bold; }
"""

with gr.Blocks(elem_id="main") as demo:
    gr.Markdown(
        """
# 🏃 Images Captioning Text
Upload **any number of images** → get captions in **Khmer (ខ្មែរ)** + **English**.  
Results update **live** as each image is processed. Powered by **Florence-2-base** — free, no API key.
        """,
        elem_id="title",
    )

    with gr.Row():
        with gr.Column(scale=1):
            upload = gr.File(
                label="📁 Upload Images (no limit)",
                file_count="multiple",
                file_types=["image"],
            )
            prefix_box = gr.Textbox(
                label="Dataset path prefix",
                value="datasets_train/Train/images",
            )
            run_btn = gr.Button("✨ Generate Captions", variant="primary", size="lg")
            status_box = gr.Textbox(
                label="Progress",
                value="",
                interactive=False,
                elem_id="status",
            )

        with gr.Column(scale=2):
            with gr.Tab("📋 JSON Output"):
                json_out = gr.Textbox(
                    label="JSON (updates live)",
                    lines=20,
                    max_lines=60,
                    elem_id="json_box",
                )
                dl_file = gr.File(label="⬇️ Download captions.json", visible=False)
            with gr.Tab("🗂️ Table View"):
                table_out = gr.Dataframe(
                    headers=["Image Path", "Caption (ខ្មែរ)", "Caption (English)"],
                    datatype=["str", "str", "str"],
                    wrap=True,
                )

    gr.Markdown("---\n**Output fields:** `images` · `caption_kh` · `caption_en`")

    run_btn.click(
        fn=process_batch,
        inputs=[upload, prefix_box],
        outputs=[json_out, dl_file, table_out, status_box],
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="violet"), css=CSS, debug=True)