import gradio as gr
from PIL import Image
import numpy as np
import torch
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
from peft import PeftModel

system_prompt = (
    "A conversation between User and Assistant. The user asks a question, and the Assistant solves it. "
    "El assistant es un experto sobre Colombia. Primero razona en mente y luego da la respuesta. "
    "El razonamiento y la respuesta van en <think></think> y <answer></answer>."
)

MODEL_ID   = "Qwen/Qwen2.5-VL-3B-Instruct"
ADAPTER_ID = "Factral/qwen2.5vl-3b-colombia-finetuned"

processor = AutoProcessor.from_pretrained(MODEL_ID)

has_gpu   = torch.cuda.is_available()
attn_impl = "flash_attention_2" if has_gpu else "eager"
model     = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation=attn_impl,
    device_map="auto",
)
model = PeftModel.from_pretrained(model, ADAPTER_ID).merge_and_unload()
model.eval().to(torch.device("cuda" if has_gpu else "cpu"))

example_imgs = [
    ("6.png",   "Shakira"),
    ("163.png", "Tienda esquinera"),
    ("img_71_2.png", "Comida colombiana"),
    ("img_98.png",   "Oso de anteojos"),
]

def cargar_imagen(path: str) -> Image.Image:
    return Image.open(path)

CSS_CUSTOM = """
/* Galería horizontal con miniaturas */
#galeria-scroll {
    overflow-x: auto;
    overflow-y: hidden;
    padding: 4px;
    scrollbar-width: thin;
}
#galeria-scroll .gallery { flex-wrap: nowrap !important; }
#galeria-scroll .gallery-item {
    flex: 0 0 auto !important;
    width: 90px !important;
    height: 90px !important;
    margin-right: 6px;
}
#galeria-scroll .gallery-item img { object-fit: cover; }

/* Texto blanco y sin halo azul al enfocar */
input, textarea { color: #fff !important; }
input::placeholder, textarea::placeholder { color: #ddd !important; }
label { color: #fff !important; }

.gr-text-input:focus-within,
.gr-text-area:focus-within,
.gr-input:focus-within {
    outline: none !important;
    box-shadow: none !important;
    border-color: #888 !important;   /* gris neutro opcional */
}

/* Por si quedaba algo en el propio input/textarea */
input:focus, textarea:focus,
input:focus-visible, textarea:focus-visible {
    outline: none !important;
    box-shadow: none !important;
    border-color: #888 !important;

"""

with gr.Blocks(theme="lone17/kotaemon", css=CSS_CUSTOM) as demo:
    # título
    gr.Markdown(
        """
        <h1>🇨🇴 
            <span style='color:gold;'>Bacan</span><span style='color:blue;'>oResp</span><span style='color:red;'>onder</span>
        </h1>
        <p>Sube o elige una imagen, haz una pregunta y obtén una respuesta con contexto local.</p>
        """
    )

    # motivación / ideas futuras en dos columnas
    with gr.Row():
        with gr.Column():
            gr.Markdown(
                """
                #### 📌 Motivación del proyecto  
                BacanoResponder permite a los usuarios colombianos obtener información contextual de sus imágenes.  
                <br/>

                #### 🌟 Impacto  
                Difunde cultura local y apoya a estudiantes, turistas y creadores de contenido.  

                #### 👥 Equipo  
                • Fabian Perez  
                • Henry Mantilla  
                • Andrea Parra  
                • Juan Calderón  
                • Semillero de Investigación del que somos parte [SemilleroCV](https://semillerocv.github.io/)  
                """
            )
        with gr.Column():
            gr.Markdown(
                """
                #### 🚀 Ideas futuras  
                - 📈 Escalar el dataset  
                - 🎤 Soporte de voz en dialectos regionales  
                - 🌐 Traducción automática  
                - 🗺️ Más dialectos/costumbres  
                - 🔄 Retroalimentación continua  
                - 🗺️ Mapas turísticos  

                #### 🤖 Modelos utilizados  
                - *Qwen2.5-VL-3B-Instruct*  
                - Dataset: [QuestionAnswer-ImgsColombia](https://huggingface.co/datasets/4nd/QuestionAnswer-ImgsColombia)  
                """
            )

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            pregunta = gr.Textbox(
                label="❓ Pregunta sobre tu imagen",
                placeholder="¿Qué muestra esta imagen?",
                lines=2,
            )

            galeria = gr.Gallery(
                label="📁 Elige una imagen de ejemplo",
                value=[img for img, _ in example_imgs],
                columns=3,
                height="384px",
                allow_preview=True,
                show_label=True,
                elem_id="galeria-scroll",
            )

        with gr.Column(scale=1):
            imagen_mostrada = gr.Image(
                label="🖼 Imagen seleccionada o subida",
                type="numpy",
                height=256,
            )
            respuesta = gr.Textbox(
                label="🧠 Respuesta",
                interactive=False,
                lines=4,
            )
            btn_procesar = gr.Button("🔍 Procesar")

    def seleccionar_imagen(evt: gr.SelectData):
        path = example_imgs[evt.index][0]
        return np.array(cargar_imagen(path))
    galeria.select(fn=seleccionar_imagen, inputs=None, outputs=imagen_mostrada)

    def responder(img, pregunta_text):
        if img is None or pregunta_text.strip() == "":
            return "Por favor sube una imagen y escribe una pregunta."

        if isinstance(img, np.ndarray):
            img = Image.fromarray(img.astype("uint8"))

        messages = [
            {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
            {"role": "user",
             "content": [
                 {"type": "text",  "text": pregunta_text},
                 {"type": "image", "image": img},
             ]},
        ]
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        ).to(model.device)

        with torch.no_grad():
            out_ids = model.generate(**inputs, max_new_tokens=512, top_p=1.0, do_sample=True, temperature=0.9)
        trimmed = [o[len(i):] for i, o in zip(inputs.input_ids, out_ids)]
        return processor.batch_decode(trimmed, skip_special_tokens=True)[0]

    btn_procesar.click(responder, inputs=[imagen_mostrada, pregunta], outputs=respuesta)

if __name__ == "__main__":
    demo.launch()