"""
app.py — Image Captioning Web App (Dual Model Support: Pretrained vs Fine-Tuned)

FastAPI backend that serves:
- /api/predict: accepts image upload & model_choice ("both", "fine-tuned", "pretrained") -> returns captions
- /api/status: returns loading and device status
- /: serves HTML frontend
"""

import os
import sys
import io
import torch
from pathlib import Path
from PIL import Image
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse
from transformers import (
    BlipProcessor,
    BlipForConditionalGeneration,
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
)

# ============================================================
# CONFIG
# ============================================================
MODEL_NAME = "Salesforce/blip-image-captioning-base"
FINE_TUNED_LOCAL_PATH = "./flickr8k_blip_output/best_model"
FINE_TUNED_HUB_NAME = "Pokzy/flickr8k-finetuned"
TRANSLATION_MODEL_NAME = "facebook/nllb-200-distilled-600M"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(f"[INFO] Using device: {DEVICE}")

# ============================================================
# DUAL MODEL MANAGER & TRANSLATOR
# ============================================================
models = {}
processors = {}
translation_model = None
translation_tokenizer = None


def load_translation_model():
    """Lazy load NLLB translation model for EN -> TH translation."""
    global translation_model, translation_tokenizer
    if translation_model is None:
        print(f"[INFO] Loading NLLB Translation model ({TRANSLATION_MODEL_NAME})...")
        translation_tokenizer = AutoTokenizer.from_pretrained(TRANSLATION_MODEL_NAME)
        translation_model = AutoModelForSeq2SeqLM.from_pretrained(TRANSLATION_MODEL_NAME).to(DEVICE)
        translation_model.eval()
        print("[INFO] NLLB Translation model loaded successfully!")


def translate_to_thai(texts: list[str]) -> list[str]:
    """Translate list of English texts to Thai using NLLB."""
    if not texts:
        return []
    load_translation_model()

    inputs = translation_tokenizer(
        texts, return_tensors="pt", padding=True, truncation=True, max_length=128
    ).to(DEVICE)

    # Convert "tha_Thai" token for NLLB target language
    thai_lang_id = translation_tokenizer.convert_tokens_to_ids("tha_Thai")

    with torch.no_grad():
        output_ids = translation_model.generate(
            **inputs,
            forced_bos_token_id=thai_lang_id,
            max_length=128,
        )

    translations = translation_tokenizer.batch_decode(output_ids, skip_special_tokens=True)
    return translations


def load_models():
    """Preload both Pretrained and Fine-Tuned BLIP models if available."""
    global models, processors

    # 1. Load Pretrained BLIP
    if "pretrained" not in models:
        print(f"[INFO] Loading Pretrained model ({MODEL_NAME})...")
        processors["pretrained"] = BlipProcessor.from_pretrained(MODEL_NAME)
        models["pretrained"] = BlipForConditionalGeneration.from_pretrained(MODEL_NAME).to(DEVICE)
        models["pretrained"].eval()
        print("[INFO] Pretrained model loaded successfully!")

    # 2. Load Fine-Tuned BLIP (Local directory or Hugging Face Model Hub)
    if "fine-tuned" not in models:
        # Priority A: Check local directory
        if os.path.exists(FINE_TUNED_LOCAL_PATH) and os.path.isdir(FINE_TUNED_LOCAL_PATH):
            has_files = any(
                os.path.exists(os.path.join(FINE_TUNED_LOCAL_PATH, f))
                for f in ["pytorch_model.bin", "model.safetensors", "config.json"]
            )
            if has_files:
                try:
                    print(f"[INFO] Loading Fine-Tuned model from local directory {FINE_TUNED_LOCAL_PATH}...")
                    try:
                        processors["fine-tuned"] = BlipProcessor.from_pretrained(FINE_TUNED_LOCAL_PATH)
                    except Exception:
                        processors["fine-tuned"] = processors.get("pretrained", BlipProcessor.from_pretrained(MODEL_NAME))

                    models["fine-tuned"] = BlipForConditionalGeneration.from_pretrained(FINE_TUNED_LOCAL_PATH).to(DEVICE)
                    models["fine-tuned"].eval()
                    print("[INFO] Local Fine-Tuned model loaded successfully!")
                except Exception as e:
                    print(f"[WARN] Could not load local Fine-Tuned model: {e}")

        # Priority B: Download from Hugging Face Hub (Pokzy/flickr8k-finetuned)
        if "fine-tuned" not in models:
            try:
                print(f"[INFO] Attempting to load Fine-Tuned model from Hugging Face Hub ({FINE_TUNED_HUB_NAME})...")
                try:
                    processors["fine-tuned"] = BlipProcessor.from_pretrained(FINE_TUNED_HUB_NAME)
                except Exception:
                    print(f"[INFO] Using base BLIP processor for Fine-Tuned model ({MODEL_NAME})...")
                    processors["fine-tuned"] = processors.get("pretrained", BlipProcessor.from_pretrained(MODEL_NAME))

                models["fine-tuned"] = BlipForConditionalGeneration.from_pretrained(FINE_TUNED_HUB_NAME).to(DEVICE)
                models["fine-tuned"].eval()
                print("[INFO] Fine-Tuned model loaded successfully from Hugging Face Hub!")
            except Exception as e:
                print(f"[WARN] Fine-Tuned model on HF Hub ({FINE_TUNED_HUB_NAME}) not loaded yet: {e}")


def generate_captions_for_model(model_key: str, image: Image.Image, num_captions: int = 5):
    """Generate captions using a specific loaded model with sampling for creative variation."""
    if model_key not in models:
        raise ValueError(f"Model '{model_key}' is not loaded or available.")

    proc = processors[model_key]
    mdl = models[model_key]

    inputs = proc(images=image, return_tensors="pt").to(DEVICE)

    with torch.no_grad():
        output_ids = mdl.generate(
            **inputs,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            num_return_sequences=num_captions,
            max_length=50,
        )

    captions = proc.batch_decode(output_ids, skip_special_tokens=True)
    return captions


# ============================================================
# FASTAPI APP
# ============================================================
app = FastAPI(title="Image Captioning Web App (Pretrained vs Fine-Tuned)", version="2.0")

# Serve static files (frontend)
app.mount("/static", StaticFiles(directory="static"), name="static")


@app.get("/", response_class=HTMLResponse)
async def read_root():
    """Serve the main HTML page."""
    index_path = Path("static/index.html")
    if index_path.exists():
        return index_path.read_text(encoding="utf-8")
    return "<h1>Image Captioning Web App</h1><p>Frontend not found. Make sure static/index.html exists.</p>"


@app.get("/api/status")
async def status():
    """Check status of loaded models."""
    available = list(models.keys())
    return {
        "status": "ready" if len(available) > 0 else "loading",
        "available_models": available,
        "device": str(DEVICE),
        "has_finetuned": "fine-tuned" in models,
        "translation_loaded": translation_model is not None,
    }


@app.post("/api/predict")
async def predict(
    file: UploadFile = File(...),
    model_choice: str = Form("both"),  # "both", "fine-tuned", "pretrained"
    language: str = Form("both"),      # "en", "th", "both"
):
    """
    Accept an image file upload and optional model choice & language.
    Returns generated captions for chosen model(s).
    """
    load_models()

    try:
        image_bytes = await file.read()
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    except Exception as e:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid image file: {str(e)}. Please upload a valid image file (JPEG, PNG, WebP, etc.)."
        )

    try:
        results = {}

        def process_model(key: str, label: str):
            en_captions = generate_captions_for_model(key, image)
            data = {"label": label, "captions": en_captions}
            if language in ["th", "both"]:
                data["captions_th"] = translate_to_thai(en_captions)
            results[key] = data

        # 1. Fine-tuned model inference (with graceful fallback to pretrained if fine-tuned model weights not present)
        if model_choice in ["both", "fine-tuned"]:
            if "fine-tuned" in models:
                process_model("fine-tuned", "Fine-Tuned BLIP (Flickr8k)")
            else:
                # Graceful fallback: use pretrained if fine-tuned is missing
                process_model("pretrained", "Pretrained BLIP (Base — Fine-Tuned Model Loading/Pending)")

        # 2. Pretrained model inference
        if model_choice in ["both", "pretrained"]:
            if "pretrained" in models and "pretrained" not in results:
                process_model("pretrained", "Pretrained BLIP (Base)")

        return {
            "results": results,
            "model_choice": model_choice,
            "language": language,
            "device": str(DEVICE),
            "filename": file.filename,
        }

    except HTTPException as he:
        raise he
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")


# ============================================================
# MAIN ENTRY
# ============================================================
if __name__ == "__main__":
    import uvicorn

    print("[INFO] Pre-loading models (Pretrained & Fine-Tuned)...")
    load_models()

    port = int(os.environ.get("PORT", "7860"))
    print(f"[INFO] Starting server at http://0.0.0.0:{port}")
    print("[INFO] Press Ctrl+C to stop\n")

    uvicorn.run(app, host="0.0.0.0", port=port)