"""
VisualRAG — Multi-Modal AI System
==================================
Stack : YOLOv8n · CLIP ViT-B/32 · FAISS · Zephyr-7B · Gradio 4.40.0
Deploy: HuggingFace Spaces (CPU Basic — free tier)

Pipeline
  Index : Image → YOLOv8 detection → CLIP embedding → FAISS vector store
  Query : Text  → CLIP text embedding → cosine k-NN → LLM answer generation

No monkey-patching needed with gradio 4.40.0 — the schema introspector
bug and starlette TemplateResponse API mismatch only affect 4.44.x.
"""

import json
import os
from datetime import datetime

import faiss
import gradio as gr
import numpy as np
import torch
from huggingface_hub import InferenceClient
from PIL import Image
from transformers import CLIPModel, CLIPProcessor
from ultralytics import YOLO

# ── Model loading (runs once at Space start-up) ────────────────────────────────
print("⏳ Loading CLIP ViT-B/32 ...")
CLIP_MODEL_ID  = "openai/clip-vit-base-patch32"
clip_model     = CLIPModel.from_pretrained(CLIP_MODEL_ID)
clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_ID)
clip_model.eval()

print("⏳ Loading YOLOv8n ...")
yolo = YOLO("yolov8n.pt")           # auto-downloads ~6 MB on first run

print("⏳ Initialising LLM client ...")
# Free HF Serverless Inference — LLM runs on HF servers, not in the Space.
# Add HF_TOKEN as a Space Secret for higher rate limits.
HF_TOKEN = os.environ.get("HF_TOKEN", None)
llm = InferenceClient(model="HuggingFaceH4/zephyr-7b-beta", token=HF_TOKEN)

print("✅ All models ready.")

# ── FAISS vector store (in-memory, session-scoped) ─────────────────────────────
EMBED_DIM   = 512                          # CLIP ViT-B/32 output dimension
faiss_index = faiss.IndexFlatIP(EMBED_DIM) # cosine similarity via L2-normalised dot product
image_store = []                           # parallel list: one dict per indexed image


# ── Embedding helpers ──────────────────────────────────────────────────────────

def embed_image(pil_img: Image.Image) -> np.ndarray:
    """Return L2-normalised 512-d CLIP image embedding (shape 1×512)."""
    inputs = clip_processor(images=pil_img, return_tensors="pt")
    with torch.no_grad():
        features = clip_model.get_image_features(**inputs)
    emb = features.numpy().astype("float32")
    faiss.normalize_L2(emb)
    return emb


def embed_text(text: str) -> np.ndarray:
    """Return L2-normalised 512-d CLIP text embedding (shape 1×512)."""
    inputs = clip_processor(text=[text], return_tensors="pt",
                            padding=True, truncation=True)
    with torch.no_grad():
        features = clip_model.get_text_features(**inputs)
    emb = features.numpy().astype("float32")
    faiss.normalize_L2(emb)
    return emb


# ── Detection pipeline ─────────────────────────────────────────────────────────

def run_detection(pil_img: Image.Image):
    """Run YOLOv8n → return (annotated PIL, detections list, summary string)."""
    results   = yolo(np.array(pil_img))[0]
    annotated = Image.fromarray(results.plot())

    detections = []
    if results.boxes is not None:
        for box in results.boxes:
            detections.append({
                "label":      yolo.names[int(box.cls[0])],
                "confidence": round(float(box.conf[0]), 3),
            })

    counts = {}
    for d in detections:
        counts[d["label"]] = counts.get(d["label"], 0) + 1
    summary = ", ".join(f"{v} {k}" for k, v in counts.items()) or "no objects detected"

    return annotated, detections, summary


# ── Index pipeline ─────────────────────────────────────────────────────────────

def index_image(image_np, note: str):
    """Detect → embed → store in FAISS. Called by the Index button."""
    if image_np is None:
        return None, "❌ Please upload an image first.", _badge()

    pil_img                        = Image.fromarray(image_np)
    annotated, detections, summary = run_detection(pil_img)
    embedding                      = embed_image(pil_img)

    faiss_index.add(embedding)
    image_store.append({
        "id":         len(image_store),
        "image":      pil_img.copy(),
        "annotated":  annotated,
        "detections": detections,
        "summary":    summary,
        "note":       note.strip() or "—",
        "ts":         datetime.now().strftime("%H:%M:%S"),
    })

    msg = f"✅ Image #{len(image_store) - 1} indexed · Found: {summary}"
    return annotated, msg, _badge()


def _badge() -> str:
    return f"📦 {len(image_store)} image(s) in vector store"


# ── RAG query pipeline ─────────────────────────────────────────────────────────

def query_images(text_query: str, top_k: int):
    """Text → CLIP embed → FAISS k-NN → RAG prompt → Zephyr-7B answer."""
    if not image_store:
        return None, "❌ No images indexed yet — upload images in the 'Detect & Index' tab first.", ""
    if not text_query.strip():
        return None, "❌ Please enter a question.", ""

    query_emb    = embed_text(text_query)
    k            = min(int(top_k), len(image_store))
    scores, idxs = faiss_index.search(query_emb, k)

    hits, ctx_lines = [], []
    for rank, (score, idx) in enumerate(zip(scores[0], idxs[0])):
        if idx < 0:
            continue
        item = image_store[int(idx)]
        hits.append({
            "rank":    rank + 1,
            "img_id":  int(idx),
            "score":   round(float(score), 4),
            "objects": item["summary"],
            "note":    item["note"],
        })
        ctx_lines.append(
            f"[Image #{idx}] objects: {item['summary']} | "
            f"note: {item['note']} | indexed at: {item['ts']} | "
            f"cosine similarity: {score:.3f}"
        )

    context = "\n".join(ctx_lines)
    prompt  = (
        "<|system|>\n"
        "You are a concise visual-AI assistant. "
        "Answer using only the retrieved image context below. "
        "If context is insufficient, say so.\n"
        "<|user|>\n"
        f"Retrieved context:\n{context}\n\n"
        f"Question: {text_query}\n"
        "<|assistant|>\n"
    )

    try:
        answer = llm.text_generation(
            prompt,
            max_new_tokens=300,
            temperature=0.2,
            repetition_penalty=1.1,
            stop_sequences=["<|user|>", "<|system|>"],
        ).strip()
    except Exception as exc:
        answer = f"⚠️ LLM unavailable ({exc}).\n\nRaw retrieval context:\n{context}"

    best_idx   = int(idxs[0][0]) if len(idxs[0]) > 0 and idxs[0][0] >= 0 else None
    best_image = image_store[best_idx]["annotated"] if best_idx is not None else None

    return best_image, answer, json.dumps(hits, indent=2)


# ── Gradio UI ──────────────────────────────────────────────────────────────────

CSS = """
.hero{text-align:center;padding:20px 0 4px}
.hero h1{font-size:28px;margin:0}
.hero p{color:var(--color-subdued);margin:6px 0 0}
.badge-row{display:flex;gap:8px;justify-content:center;flex-wrap:wrap;margin-top:10px}
.badge{background:var(--color-background-secondary);border:1px solid var(--border-color-primary);
       border-radius:20px;padding:3px 12px;font-size:12px;color:var(--color-text-body)}
"""

with gr.Blocks(title="VisualRAG", theme=gr.themes.Soft(primary_hue="blue"), css=CSS) as demo:

    gr.HTML("""
    <div class="hero">
      <h1>🔍 VisualRAG</h1>
      <p>Multi-Modal AI · Object Detection + Visual Embeddings + RAG + LLM</p>
      <div class="badge-row">
        <span class="badge">YOLOv8</span>
        <span class="badge">CLIP ViT-B/32</span>
        <span class="badge">FAISS</span>
        <span class="badge">Zephyr-7B</span>
        <span class="badge">Gradio 4.40.0</span>
      </div>
    </div>
    """)

    with gr.Tabs():

        # ── TAB 1: Detect & Index ──────────────────────────────────────────────
        with gr.Tab("📤 Detect & Index"):
            gr.Markdown(
                "Upload any image. YOLOv8n detects objects, then CLIP ViT-B/32 "
                "encodes it into a 512-d embedding stored in FAISS for later retrieval."
            )
            with gr.Row():
                with gr.Column(scale=1):
                    img_in    = gr.Image(label="Upload image", type="numpy")
                    note_in   = gr.Textbox(label="Context note (optional)",
                                           placeholder="e.g. 'Warehouse camera, aisle 3'")
                    index_btn = gr.Button("🔍 Detect & Index", variant="primary")
                with gr.Column(scale=1):
                    det_out    = gr.Image(label="Detection result")
                    status_out = gr.Textbox(label="Status", interactive=False)
                    badge_out  = gr.Textbox(label="Vector store", interactive=False,
                                            value=_badge())

            index_btn.click(
                fn=index_image,
                inputs=[img_in, note_in],
                outputs=[det_out, status_out, badge_out],
            )

        # ── TAB 2: Query (RAG) ─────────────────────────────────────────────────
        with gr.Tab("💬 Query (RAG)"):
            gr.Markdown(
                "Ask any question about your indexed images. CLIP embeds the query, "
                "FAISS retrieves the most similar images by cosine similarity, "
                "and Zephyr-7B generates a grounded answer."
            )
            with gr.Row():
                with gr.Column(scale=1):
                    query_in  = gr.Textbox(
                        label="Your question",
                        placeholder="e.g. 'How many people are visible?' or 'Are there any vehicles?'",
                        lines=3,
                    )
                    topk_sl   = gr.Slider(minimum=1, maximum=5, value=3, step=1,
                                          label="Top-K images to retrieve")
                    query_btn = gr.Button("🔎 Search & Generate Answer", variant="primary")
                with gr.Column(scale=1):
                    match_img = gr.Image(label="Best matching image")
                    llm_out   = gr.Textbox(label="AI Answer (RAG-grounded)",
                                           lines=6, interactive=False)

            hits_out = gr.Textbox(label="Retrieval scores", interactive=False, lines=8)

            query_btn.click(
                fn=query_images,
                inputs=[query_in, topk_sl],
                outputs=[match_img, llm_out, hits_out],
            )

        # ── TAB 3: How it works ────────────────────────────────────────────────
        with gr.Tab("🏗️ How it works"):
            gr.Markdown("""
## System overview

### Index pipeline
```
Image → YOLOv8n detection (objects + confidence scores)
      → CLIP ViT-B/32 image encoder  →  512-d embedding
      → L2 normalisation
      → FAISS IndexFlatIP  (cosine similarity store)
```

### Query / RAG pipeline
```
Text query → CLIP text encoder  →  512-d query embedding
           → L2 normalisation
           → FAISS k-NN search (cosine similarity, top-K)
           → RAG prompt = query + retrieved context
           → Zephyr-7B-β  (HF Serverless Inference API)
           → Natural language answer
```

## Stack

| Component | Technology |
|---|---|
| Object detection | YOLOv8n (Ultralytics) |
| Visual embedding | CLIP ViT-B/32 (OpenAI via HF) |
| Vector index | FAISS IndexFlatIP (cosine sim) |
| LLM | Zephyr-7B-β (HF Serverless API) |
| UI | Gradio 4.40.0 |

## Why gradio 4.40.0

Version 4.44.1 has three cascading runtime bugs on HF Spaces: a schema
introspector TypeError, a non-existent gradio_client pin, and a starlette
TemplateResponse API mismatch that causes a Jinja2 `unhashable type: dict`
crash. Version 4.40.0 is widely deployed and has none of these issues.
            """)

    gr.HTML("""
    <div style="text-align:center;padding:14px 0 4px;color:var(--color-subdued);font-size:12px">
      VisualRAG · YOLOv8 + CLIP + FAISS + LLM · HuggingFace Spaces
    </div>
    """)

demo.launch(server_name="0.0.0.0")