"""
Unified Chhagan's  Multi-Model Studio
Merged from 3 HuggingFace Spaces with enhanced Document Intelligence.
Features:
  - 🪪 Document Scanner (Front + Back)
  - 🖼️ Image Analysis (with Face Detection, Signature Extraction, Annotated Highlights)
  - 📚 Batch Processing
  - 💬 Chat with Attachments
All 5 Chhagan VL models available across all tabs.
"""

import os
import time
import warnings
from threading import Thread
from typing import Optional, Tuple, Dict, Any, List

import gradio as gr
import spaces
import torch
import numpy as np
from PIL import Image
import cv2

# ✅ Transformers v5 mein AutoModelForVision2Seq → AutoModelForImageTextToText rename hua
# Ye try/except dono versions handle karta hai
try:
    from transformers import AutoModelForImageTextToText as VisionModel
except ImportError:
    from transformers import AutoModelForVision2Seq as VisionModel

from transformers import (
    AutoProcessor,
    BitsAndBytesConfig,
    TextIteratorStreamer,
)
from peft import PeftConfig, PeftModel

from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

# ──────────────────────────────────────────────────────────────
# Suppress warnings
# ──────────────────────────────────────────────────────────────
warnings.filterwarnings('ignore', message='.*meta device.*')
warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=DeprecationWarning)

# ──────────────────────────────────────────────────────────────
# Custom Premium Theme
# ──────────────────────────────────────────────────────────────
colors.deep_indigo = colors.Color(
    name="deep_indigo",
    c50="#EEF2FF", c100="#E0E7FF", c200="#C7D2FE", c300="#A5B4FC",
    c400="#818CF8", c500="#6366F1", c600="#4F46E5", c700="#4338CA",
    c800="#3730A3", c900="#312E81", c950="#1E1B4B",
)
colors.cyber_teal = colors.Color(
    name="cyber_teal",
    c50="#F0FDFA", c100="#CCFBF1", c200="#99F6E4", c300="#5EEAD4",
    c400="#2DD4BF", c500="#14B8A6", c600="#0D9488", c700="#0F766E",
    c800="#115E59", c900="#134E4A", c950="#042F2E",
)

class PremiumTheme(Soft):
    def __init__(self):
        super().__init__(
            primary_hue=colors.deep_indigo,
            secondary_hue=colors.cyber_teal,
            neutral_hue=colors.slate,
            text_size=sizes.text_lg,
            font=(fonts.GoogleFont("Inter"), "system-ui", "sans-serif"),
            font_mono=(fonts.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"),
        )
        super().set(
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_950",
            body_background_fill="linear-gradient(135deg, *primary_100, *secondary_50, *primary_50)",
            body_background_fill_dark="linear-gradient(135deg, *primary_950, *neutral_900, *primary_900)",
            button_primary_text_color="white",
            button_primary_background_fill="linear-gradient(135deg, *primary_600, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(135deg, *primary_700, *secondary_700)",
            button_primary_background_fill_dark="linear-gradient(135deg, *primary_500, *secondary_500)",
            button_primary_shadow="0 4px 15px rgba(99, 102, 241, 0.35)",
            button_secondary_background_fill="linear-gradient(135deg, *neutral_100, *neutral_200)",
            button_secondary_background_fill_hover="linear-gradient(135deg, *neutral_200, *neutral_300)",
            slider_color="*secondary_500",
            slider_color_dark="*secondary_400",
            block_title_text_weight="600",
            block_border_width="1px",
            block_shadow="0 4px 20px rgba(0,0,0,0.08)",
            block_label_background_fill="*primary_100",
        )

premium_theme = PremiumTheme()

css = """
#app-title h1 {
    font-size: 2.5em !important;
    background: linear-gradient(135deg, #6366F1, #14B8A6);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-weight: 800;
    letter-spacing: -0.02em;
}
#app-subtitle { font-size: 1.1em; color: var(--neutral-500); margin-top: -8px; }
.output-box textarea { font-size: 14px !important; line-height: 1.6 !important; }
.model-selector { border: 2px solid var(--primary-200) !important; border-radius: 12px !important; padding: 8px !important; }
.tab-nav button { font-weight: 600 !important; font-size: 14px !important; }
.face-box { border: 3px solid #22c55e; border-radius: 8px; }
.sig-box  { border: 3px solid #3b82f6; border-radius: 8px; }
"""

# ──────────────────────────────────────────────────────────────
# Device & Constants
# ──────────────────────────────────────────────────────────────
MAX_MAX_NEW_TOKENS    = 4096
DEFAULT_MAX_NEW_TOKENS = 1024
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🚀 Using device: {device}")

# ──────────────────────────────────────────────────────────────
# ALL MODELS  ←── CHANGE 1: Only Chhagan custom models remain
# ──────────────────────────────────────────────────────────────
ALL_MODELS = [
    "Chhagan005/Chhagan_ML-VL-OCR-v1",               # Qwen2.5-VL-3B base
    "Chhagan005/Chhagan-DocVL-Qwen3",                 # Qwen3-VL-2B LoRA adapter
    "Chhagan005/CSM-DocExtract-VL",                   # Qwen3-VL-8B BNB INT4
    "Chhagan005/CSM-DocExtract-VL-HF",                # Qwen3-VL-8B full
    "Chhagan005/CSM-DocExtract-VL-Q4KM-merged-fp16",  # Qwen3-VL-8B merged fp16
]

# ── BNB INT4 quantized models ──
BNB_QUANT_MODELS = {
    "Chhagan005/CSM-DocExtract-VL",
}

# ── LoRA Adapter only (unmerged) models ──
PEFT_ADAPTER_MODELS = {
    "Chhagan005/Chhagan-DocVL-Qwen3",
}

# ── Qwen VL image placeholder ──
QWEN_VL_IMG_TOKEN = "<|vision_start|><|image_pad|><|vision_end|>"

# ──────────────────────────────────────────────────────────────
# Lazy Model Loading — 4 strategies
# ──────────────────────────────────────────────────────────────
_model_cache: Dict[str, Tuple[Any, Any]] = {}

def load_model(model_id: str):
    if model_id in _model_cache:
        return _model_cache[model_id]

    print(f"⏳ Loading model: {model_id}")
    model = None

    with warnings.catch_warnings():
        warnings.filterwarnings('ignore')
        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)

        # ── Strategy 1: PEFT LoRA Adapter (unmerged) ─────────────────────
        if model_id in PEFT_ADAPTER_MODELS:
            try:
                peft_config   = PeftConfig.from_pretrained(model_id)
                base_model_id = peft_config.base_model_name_or_path
                print(f"   └─ PEFT base: {base_model_id}")
                dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
                base_model = VisionModel.from_pretrained(
                    base_model_id,
                    torch_dtype=dtype,
                    device_map="auto",
                    trust_remote_code=True,
                )
                model = PeftModel.from_pretrained(base_model, model_id)
                model = model.merge_and_unload()
                print(f"✅ {model_id} → PEFT merged")
            except Exception as e:
                print(f"⚠️ PEFT failed ({e}), falling back...")
                model = None

        # ── Strategy 2: BNB INT4 explicit config ─────────────────────────
        if model is None and model_id in BNB_QUANT_MODELS:
            try:
                bnb_config = BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_quant_type="nf4",
                    bnb_4bit_use_double_quant=True,
                    bnb_4bit_compute_dtype=torch.bfloat16,
                )
                model = VisionModel.from_pretrained(
                    model_id,
                    quantization_config=bnb_config,
                    device_map="auto",
                    trust_remote_code=True,
                )
                print(f"✅ {model_id} → BNB INT4")
            except Exception as e:
                print(f"⚠️ BNB failed ({e}), falling back...")
                model = None

        # ── Strategy 3: Standard bf16 ─────────────────────────────────────
        if model is None:
            try:
                dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
                model = VisionModel.from_pretrained(
                    model_id,
                    torch_dtype=dtype,
                    device_map="auto",
                    trust_remote_code=True,
                )
                print(f"✅ {model_id} → bf16")
            except Exception as e:
                print(f"⚠️ bf16 failed ({e}), last fallback...")
                model = None

        # ── Strategy 4: Last resort float16 ──────────────────────────────
        if model is None:
            model = VisionModel.from_pretrained(
                model_id,
                torch_dtype=torch.float16,
                device_map="auto",
                trust_remote_code=True,
                low_cpu_mem_usage=True,
            )
            print(f"✅ {model_id} → float16 fallback")

        model.eval()

    _model_cache[model_id] = (processor, model)
    print(f"✅ Model {model_id} ready on {device}")
    return processor, model

# ──────────────────────────────────────────────────────────────
# Pre-load default model  ←── CHANGE 2: New default model
# ──────────────────────────────────────────────────────────────
DEFAULT_MODEL = "Chhagan005/CSM-DocExtract-VL-HF"
print(f"⏳ Pre-loading default model: {DEFAULT_MODEL}")
load_model(DEFAULT_MODEL)
print(f"✅ Default model ready!")

# ──────────────────────────────────────────────────────────────
# Universal prepare_inputs
# ──────────────────────────────────────────────────────────────
def prepare_inputs(processor, model, messages: List[Dict]) -> Dict:
    pil_images: List[Image.Image] = []
    flat_messages = []

    for msg in messages:
        role    = msg.get("role", "user")
        content = msg.get("content", "")

        if isinstance(content, list):
            parts = []
            for item in content:
                if not isinstance(item, dict):
                    parts.append(str(item))
                    continue
                t = item.get("type", "")
                if t == "text":
                    parts.append(item.get("text", ""))
                elif t == "image":
                    img = item.get("image")
                    if img is not None and isinstance(img, Image.Image):
                        pil_images.append(img)
                    parts.append(QWEN_VL_IMG_TOKEN)
            flat_messages.append({"role": role, "content": "".join(parts)})
        else:
            flat_messages.append({"role": role, "content": content})

    text = processor.apply_chat_template(
        flat_messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    has_image_proc = hasattr(processor, "image_processor")

    if pil_images and has_image_proc:
        inputs = processor(
            text=[text],
            images=pil_images,
            padding=True,
            return_tensors="pt",
        )
    else:
        inputs = processor(
            text=[text],
            padding=True,
            return_tensors="pt",
        )

    return {k: v.to(model.device) if torch.is_tensor(v) else v for k, v in inputs.items()}

# ──────────────────────────────────────────────────────────────
# Utility
# ──────────────────────────────────────────────────────────────
def ensure_rgb(image: Image.Image) -> Optional[Image.Image]:
    if image is None:
        return None
    if image.mode == "RGBA":
        bg = Image.new("RGB", image.size, (255, 255, 255))
        bg.paste(image, mask=image.split()[3])
        return bg
    elif image.mode != "RGB":
        return image.convert("RGB")
    return image

# ──────────────────────────────────────────────────────────────
# Face Detection, Signature Extraction & Annotation
# ──────────────────────────────────────────────────────────────
def detect_faces(image: Image.Image):
    img_array = np.array(image)
    gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.08, minNeighbors=4, minSize=(40, 40), flags=cv2.CASCADE_SCALE_IMAGE)
    if len(faces) == 0:
        profile_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_profileface.xml')
        faces = profile_cascade.detectMultiScale(gray, scaleFactor=1.08, minNeighbors=4, minSize=(40, 40))
    if len(faces) == 0:
        return None, []

    faces_sorted = sorted(faces, key=lambda f: f[2] * f[3], reverse=True)
    x, y, w, h = faces_sorted[0]
    pad = int(0.2 * max(w, h))
    x1, y1 = max(0, x - pad), max(0, y - pad)
    x2, y2 = min(img_array.shape[1], x + w + pad), min(img_array.shape[0], y + h + pad)

    face_gray = gray[y1:y2, x1:x2]
    if face_gray.size > 0 and cv2.Laplacian(face_gray, cv2.CV_64F).var() < 30:
        if len(faces_sorted) > 1:
            x, y, w, h = faces_sorted[1]
            pad = int(0.2 * max(w, h))
            x1, y1 = max(0, x - pad), max(0, y - pad)
            x2, y2 = min(img_array.shape[1], x + w + pad), min(img_array.shape[0], y + h + pad)
            face_gray2 = gray[y1:y2, x1:x2]
            if face_gray2.size > 0 and cv2.Laplacian(face_gray2, cv2.CV_64F).var() < 30:
                return None, [tuple(f) for f in faces_sorted]
        else:
            return None, [tuple(f) for f in faces_sorted]

    face_crop   = image.crop((x1, y1, x2, y2))
    face_bboxes = [tuple(int(v) for v in f) for f in faces_sorted]
    return face_crop, face_bboxes


def detect_signature(image: Image.Image):
    img_array = np.array(image)
    h, w = img_array.shape[:2]
    search_top = int(h * 0.5)
    gray = cv2.cvtColor(img_array[search_top:, :], cv2.COLOR_RGB2GRAY)
    binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 15, 10)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 3))
    binary = cv2.erode(cv2.dilate(binary, kernel, iterations=2), kernel, iterations=1)
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None, None

    sig_contours = [c for c in contours if cv2.contourArea(c) > 200
                    and 1.0 < cv2.boundingRect(c)[2] / max(cv2.boundingRect(c)[3], 1) < 15
                    and cv2.boundingRect(c)[2] > 40 and cv2.boundingRect(c)[3] > 10]
    if not sig_contours:
        return None, None

    rx, ry, rw, rh = cv2.boundingRect(np.concatenate(sig_contours))
    if rw < 30 or rh < 10:
        return None, None

    pad = 12
    sx1, sy1 = max(0, rx - pad), max(0, search_top + ry - pad)
    sx2, sy2 = min(w, rx + rw + pad), min(h, search_top + ry + rh + pad)
    region_gray = cv2.cvtColor(img_array[sy1:sy2, sx1:sx2], cv2.COLOR_RGB2GRAY)
    _, region_bin = cv2.threshold(region_gray, 120, 255, cv2.THRESH_BINARY_INV)
    ink_density = np.sum(region_bin > 0) / max(region_bin.size, 1)
    if not (0.03 <= ink_density <= 0.6):
        return None, None
    return image.crop((sx1, sy1, sx2, sy2)), (sx1, sy1, sx2, sy2)


def create_annotated_image(image: Image.Image, face_bboxes: list, sig_bbox: Optional[tuple]):
    img_array = np.array(image).copy()
    for i, (x, y, w, h) in enumerate(face_bboxes):
        cv2.rectangle(img_array, (x, y), (x + w, y + h), (34, 197, 94), 3)
        label = f"Face #{i+1}"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
        cv2.rectangle(img_array, (x, y - th - 10), (x + tw + 6, y), (34, 197, 94), -1)
        cv2.putText(img_array, label, (x + 3, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    if sig_bbox:
        x1, y1, x2, y2 = sig_bbox
        cv2.rectangle(img_array, (x1, y1), (x2, y2), (59, 130, 246), 3)
        label = "Signature"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
        cv2.rectangle(img_array, (x1, y1 - th - 10), (x1 + tw + 6, y1), (59, 130, 246), -1)
        cv2.putText(img_array, label, (x1 + 3, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    return Image.fromarray(img_array)


def run_visual_extraction(image: Optional[Image.Image]):
    if image is None:
        return None, None, None, "_Upload an image to detect visual elements._"
    image = ensure_rgb(image)
    detections = []
    face_crop, face_bboxes = detect_faces(image)
    if face_crop is not None:
        detections.append(f"✅ **Face detected** — {len(face_bboxes)} face(s) found")
    elif face_bboxes:
        detections.append(f"⚠️ **Face found but blurry** — {len(face_bboxes)} detected")
    else:
        detections.append("❌ **No face detected**")
    sig_crop, sig_bbox = detect_signature(image)
    detections.append("✅ **Signature detected**" if sig_crop else "ℹ️ **No signature detected**")
    annotated = create_annotated_image(image, face_bboxes, sig_bbox)
    detections.append(f"\n🎯 **Annotated** — {len(face_bboxes)} face box(es)" + (" + 1 sig" if sig_bbox else ""))
    return face_crop, sig_crop, annotated, "### 🔍 Detection Results\n\n" + "\n\n".join(detections)


# ──────────────────────────────────────────────────────────────
# TAB 1: Document Scanner
# ──────────────────────────────────────────────────────────────
@spaces.GPU(duration=180)
def generate_document_scan(
    model_name, front_image, back_image, prompt,
    max_new_tokens=1024, temperature=0.6, top_p=0.9, top_k=50, repetition_penalty=1.2,
):
    if front_image is None and back_image is None:
        yield "⚠️ Please upload at least one image.", "⚠️ Please upload at least one image."
        return
    if not prompt.strip():
        prompt = "Analyze this document. Extract all text and key details and provide a structured summary."
    try:
        processor, model = load_model(model_name)
    except Exception as e:
        yield f"❌ {e}", f"❌ {e}"; return

    content = []
    if front_image is not None:
        front_image = ensure_rgb(front_image)
        content += [{"type": "text", "text": "**[FRONT SIDE]**"}, {"type": "image", "image": front_image}]
    if back_image is not None:
        back_image = ensure_rgb(back_image)
        content += [{"type": "text", "text": "**[BACK SIDE]**"}, {"type": "image", "image": back_image}]
    content.append({"type": "text", "text": prompt})

    inputs   = prepare_inputs(processor, model, [{"role": "user", "content": content}])
    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
    Thread(target=model.generate, kwargs={**inputs, "streamer": streamer,
        "max_new_tokens": max_new_tokens, "do_sample": temperature > 0,
        "temperature": temperature, "top_p": top_p, "top_k": top_k}).start()
    buffer = ""
    for t in streamer:
        buffer += t; time.sleep(0.01); yield buffer, buffer


# ──────────────────────────────────────────────────────────────
# TAB 2: Image Analysis
# ──────────────────────────────────────────────────────────────
@spaces.GPU(duration=180)
def generate_image_analysis(
    model_name, text, image,
    max_new_tokens=1024, temperature=0.6, top_p=0.9, top_k=50, repetition_penalty=1.2,
):
    if image is None:
        yield "⚠️ Please upload an image.", "⚠️ Please upload an image."; return
    if not text.strip():
        text = "Describe this image in detail."
    try:
        processor, model = load_model(model_name)
    except Exception as e:
        yield f"❌ {e}", f"❌ {e}"; return

    image    = ensure_rgb(image)
    messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text}]}]
    inputs   = prepare_inputs(processor, model, messages)
    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
    Thread(target=model.generate, kwargs={**inputs, "streamer": streamer,
        "max_new_tokens": max_new_tokens, "do_sample": temperature > 0,
        "temperature": temperature, "top_p": top_p, "top_k": top_k}).start()
    buffer = ""
    for t in streamer:
        buffer += t; time.sleep(0.01); yield buffer, buffer


# ──────────────────────────────────────────────────────────────
# TAB 3: Batch Processing
# ──────────────────────────────────────────────────────────────
def process_batch_images(
    files: List, prompts_text: str, model_name: str,
    max_new_tokens: int, temperature: float, top_p: float, top_k: int, seed: int,
):
    if not files:            return "⚠️ Please upload images."
    if not prompts_text.strip(): return "⚠️ Please enter prompts."

    prompts = [p.strip() for p in prompts_text.split('\n') if p.strip()]
    if len(prompts) == 1:        prompts = prompts * len(files)
    elif len(prompts) != len(files): return f"⚠️ Prompt count ({len(prompts)}) ≠ image count ({len(files)})."

    try:
        processor, model = load_model(model_name)
    except Exception as e:
        return f"❌ {e}"

    results = []
    for idx, (file, prompt) in enumerate(zip(files, prompts), 1):
        try:
            image_path = file.name if hasattr(file, 'name') else file
            image = Image.open(image_path).convert("RGB")
            if seed != -1: torch.manual_seed(seed + idx - 1)
            messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt}]}]
            inputs   = prepare_inputs(processor, model, messages)
            with torch.no_grad():
                generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens,
                    temperature=temperature, top_p=top_p, top_k=top_k, do_sample=temperature > 0)
            trimmed = [o[len(i):] for i, o in zip(inputs['input_ids'], generated_ids)]
            result  = processor.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        except Exception as e:
            result = f"Error: {e}"
        results += [f"═══ Image {idx}: {os.path.basename(str(image_path))} ═══",
                    f"📝 Prompt: {prompt}", f"📄 Result: {result}\n"]
    return "\n".join(results)


# ──────────────────────────────────────────────────────────────
# TAB 4: Chat
# ──────────────────────────────────────────────────────────────
@spaces.GPU(duration=120)
def process_chat_message(
    message: str, image: Optional[Image.Image],
    history: List[Dict[str, Any]], model_name: str,
) -> str:
    try:
        processor, model = load_model(model_name)
    except Exception as e:
        return f"❌ {e}"

    content = []
    if image is not None:
        image = ensure_rgb(image)
        content.append({"type": "image", "image": image})
    if message:
        content.append({"type": "text", "text": message})

    messages = [{"role": m["role"], "content": m["content"]}
                for m in history if m.get("role") in ("user", "assistant")]
    if content:
        messages.append({"role": "user", "content": content})

    inputs = prepare_inputs(processor, model, messages)
    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=1024,
            temperature=0.7, do_sample=True, top_p=0.95)
    trimmed = [o[len(i):] for i, o in zip(inputs['input_ids'], generated_ids)]
    return processor.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]


def chat_fn(message: Dict[str, Any], history: List[Dict[str, Any]], model_name: str):
    text  = message.get("text", "")
    files = message.get("files", [])
    image = None
    if files:
        try: image = ensure_rgb(Image.open(files[0]))
        except Exception as e: print(f"Image load error: {e}")

    try:
        response = process_chat_message(text, image, history, model_name)
    except Exception as e:
        response = f"❌ Error: {e}"

    history.append({"role": "user",      "content": f"{text}\n📎 [Image attached]" if image else text})
    history.append({"role": "assistant", "content": response})
    return "", history


def retry_fn(history, model_name):
    if not history or len(history) < 2: return "", history
    last    = history[-2]
    history = history[:-2]
    text    = last.get("content", "").replace("\n📎 [Image attached]", "").replace("📎 [Image attached]", "")
    return chat_fn({"text": text}, history, model_name)

def undo_fn(history):  return history[:-2] if len(history) >= 2 else []
def clear_fn():        return "", []


# ══════════════════════════════════════════════════════════════
# ████  GRADIO UI  ████
# ══════════════════════════════════════════════════════════════
with gr.Blocks(title="Chhagan's Multi-Model Studio") as demo:

    gr.Markdown("# 🧠 Chhagan's  Multi-Model Studio", elem_id="app-title")
    gr.Markdown(
        # ←── CHANGE 3: Updated subtitle (removed Qwen official mention)
        "_Custom CSM/Chhagan VL models — "
        "Document Scanning • Image Analysis (Face + Signature) • Batch Processing • Chat_",
        elem_id="app-subtitle",
    )

    with gr.Row():
        model_choice = gr.Dropdown(choices=ALL_MODELS, value=DEFAULT_MODEL,
            label="🤖 Select Model",
            # ←── CHANGE 4: Updated info text
            info="5 custom CSM/Chhagan models",
            scale=3, elem_classes="model-selector")

    with gr.Accordion("⚙️ Advanced Generation Parameters", open=False):
        with gr.Row():
            max_new_tokens     = gr.Slider(64, MAX_MAX_NEW_TOKENS, DEFAULT_MAX_NEW_TOKENS, step=64, label="Max New Tokens")
            temperature        = gr.Slider(0.1, 2.0, 0.6,  step=0.1,  label="Temperature")
        with gr.Row():
            top_p              = gr.Slider(0.05, 1.0, 0.9,  step=0.05, label="Top-p")
            top_k              = gr.Slider(1, 1000, 50,      step=1,    label="Top-k")
        with gr.Row():
            repetition_penalty = gr.Slider(1.0, 2.0, 1.2,  step=0.05, label="Repetition Penalty")
            seed_number        = gr.Number(value=-1, label="Seed (-1 = random)", precision=0)

    with gr.Tabs():

        # ─── TAB 1 ───
        with gr.TabItem("🪪 Document Scanner"):
            gr.Markdown("### Scan Front & Back of Documents\nFace profiles and signatures are **auto-detected** on front image upload.")
            with gr.Row():
                with gr.Column(scale=1):
                    doc_front_image = gr.Image(type="pil", label="📄 Front Side", height=280)
                    doc_back_image  = gr.Image(type="pil", label="📄 Back Side",  height=280)
                    doc_prompt = gr.Textbox(label="Custom Prompt (optional)", lines=3,
                        placeholder="e.g., Extract all text, MRZ data, name, DOB, ID number...",
                        value="Analyze this document carefully. Extract ALL text, key fields (Name, DOB, ID Number, Expiry, MRZ etc.) and provide a structured summary.")
                    doc_submit = gr.Button("🔍 Scan Document", variant="primary", size="lg")
                with gr.Column(scale=1):
                    doc_output = gr.Textbox(label="📋 Scan Results", lines=18, elem_classes="output-box")
                    with gr.Accordion("📑 Formatted Result (Markdown)", open=False):
                        doc_md_output = gr.Markdown()
            gr.Markdown("---\n### 🔍 Visual Element Detection _(auto-detected on front image upload)_")
            with gr.Row():
                doc_face_output      = gr.Image(label="👤 Detected Face",      height=220, elem_classes="face-box")
                doc_sig_output       = gr.Image(label="✍️ Detected Signature", height=220, elem_classes="sig-box")
                doc_annotated_output = gr.Image(label="🎯 Annotated Image",    height=220)
            doc_detection_summary = gr.Markdown("_Upload a front side image to detect visual elements._")
            doc_front_image.change(fn=run_visual_extraction, inputs=[doc_front_image],
                outputs=[doc_face_output, doc_sig_output, doc_annotated_output, doc_detection_summary])
            doc_submit.click(fn=generate_document_scan,
                inputs=[model_choice, doc_front_image, doc_back_image, doc_prompt,
                        max_new_tokens, temperature, top_p, top_k, repetition_penalty],
                outputs=[doc_output, doc_md_output])

        # ─── TAB 2 ───
        with gr.TabItem("🖼️ Image Analysis"):
            gr.Markdown("### Smart Image Analysis\nAuto-detects **face profiles**, **signatures**, and **annotations**.")
            with gr.Row():
                with gr.Column(scale=1):
                    img_upload = gr.Image(type="pil", label="Upload Image", height=320)
                    img_query  = gr.Textbox(label="Query / Prompt", lines=2,
                        placeholder="What do you see? / Extract all text / Describe in detail...")
                    img_submit = gr.Button("🚀 Analyze with Model", variant="primary")
                with gr.Column(scale=1):
                    img_output = gr.Textbox(label="📄 Model Analysis", lines=12, elem_classes="output-box")
                    with gr.Accordion("📑 Formatted (Markdown)", open=False):
                        img_md_output = gr.Markdown()
            gr.Markdown("---\n### 🔍 Visual Element Detection _(auto-detected on upload)_")
            with gr.Row():
                face_output      = gr.Image(label="👤 Detected Face",      height=220, elem_classes="face-box")
                sig_output       = gr.Image(label="✍️ Detected Signature", height=220, elem_classes="sig-box")
                annotated_output = gr.Image(label="🎯 Annotated Image",    height=220)
            detection_summary = gr.Markdown("_Upload an image to detect visual elements._")
            img_upload.change(fn=run_visual_extraction, inputs=[img_upload],
                outputs=[face_output, sig_output, annotated_output, detection_summary])
            img_submit.click(fn=generate_image_analysis,
                inputs=[model_choice, img_query, img_upload, max_new_tokens,
                        temperature, top_p, top_k, repetition_penalty],
                outputs=[img_output, img_md_output])

        # ─── TAB 3 ───
        with gr.TabItem("📚 Batch Processing"):
            gr.Markdown("### Process Multiple Images at Once")
            with gr.Row():
                with gr.Column(scale=1):
                    batch_images  = gr.File(file_count="multiple", label="Upload Images", file_types=["image"])
                    batch_prompts = gr.Textbox(label="Prompts (one per line)", lines=5,
                        placeholder="Describe this image\nExtract all text...",
                        info="One prompt for all OR one per image")
                    batch_submit  = gr.Button("🚀 Process Batch", variant="primary")
                with gr.Column(scale=1):
                    batch_output = gr.Textbox(label="Batch Results", lines=20, elem_classes="output-box")
            batch_submit.click(fn=process_batch_images,
                inputs=[batch_images, batch_prompts, model_choice,
                        max_new_tokens, temperature, top_p, top_k, seed_number],
                outputs=batch_output)

        # ─── TAB 4 ───
        with gr.TabItem("💬 Chat"):
            gr.Markdown("### Multi-Turn Chat with Image Attachments")
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("**💡 Tips:**\n- Upload an image and ask questions\n- Multi-turn memory\n- Switch models anytime")
                with gr.Column(scale=3):
                    chatbot  = gr.Chatbot(label="Chat", height=450, value=[])
                    chat_msg = gr.MultimodalTextbox(label="Message",
                        placeholder="Type a message or upload an image...",
                        file_types=["image"], submit_btn=True, stop_btn=False)
                    with gr.Row():
                        retry_btn = gr.Button("🔄 Retry", variant="secondary", size="sm")
                        undo_btn  = gr.Button("↩️ Undo",  variant="secondary", size="sm")
                        clear_btn = gr.Button("🗑️ Clear", variant="secondary", size="sm")
            chat_msg.submit(chat_fn,  [chat_msg, chatbot, model_choice], [chat_msg, chatbot], queue=True)
            retry_btn.click(retry_fn, [chatbot, model_choice],           [chat_msg, chatbot], queue=True)
            undo_btn.click( undo_fn,  [chatbot],                         [chatbot],           queue=False)
            clear_btn.click(clear_fn, outputs=[chat_msg, chatbot],                            queue=False)

    gr.Markdown(
        # ←── CHANGE 5: Updated footer (removed Qwen official models, 10→5)
        "---\n**🧠 Chhagan's  Multi-Model Studio** • 5 Models\n\n"
        "CSM-DocExtract-VL (BNB-INT4) • CSM-DocExtract-VL-HF • "
        "CSM-DocExtract-VL-Q4KM-merged-fp16 • Chhagan_ML-VL-OCR-v1 • Chhagan-DocVL-Qwen3\n\n"
        "_Built with ❤️ using Gradio_"
    )

if __name__ == "__main__":
    demo.queue(max_size=50).launch(
        server_name="0.0.0.0", server_port=7860,
        share=False, show_error=True, ssr_mode=False,
        theme=premium_theme, css=css,
    )