import gradio as gr
import spaces
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os

# AUTENTICACIÓN CORRECTA PARA GEMMA EN SPACES
hf_token = os.getenv("HF_CLASSIFIER")
if hf_token:
    login(token=hf_token)
    print("Autenticado con HF_CLASSIFIER")
else:
    print("ERROR: HF_CLASSIFIER no encontrado en Secrets")

# Carga modelo y tokenizer
model_id = "google/gemma-3-4b-it"
print(f"Cargando {model_id}...")
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True
)

SYSTEM_PROMPT = """You are an expert assistant that analyzes user messages. 
ALWAYS respond ONLY with valid JSON in this EXACT format:

{
  "summary": "Brief and clear summary of the user's message (max 1-2 sentences)",
  "tags": ["tag1", "tag2", "tag3"],
  "sentiment": {
    "score": -1.0,
    "label": "NEGATIVE | NEUTRAL | POSITIVE",
    "explanation": "Brief reason for the analysis"
  }
}

RULES:
- summary: Capture the main essence
- tags: 3-5 relevant keywords (themes, emotions, entities)
- sentiment: score (-1.0=negative, 0=neutral, +1.0=positive)
- NEVER add text outside the JSON
- JSON must be valid and parseable"""

@spaces.GPU(duration=120)
def generate_response(message, history):
    model.to('cuda')
    
    messages = []
    for user_msg, ai_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if ai_msg:
            messages.append({"role": "assistant", "content": ai_msg})
    messages.append({"role": "user", "content": "System: "+SYSTEM_PROMPT+" User:"+message})
    
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=1512,  # Aumentado para evitar respuestas cortadas (truncadas)
            do_sample=True,
            temperature=0.7,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id  # Añadido para forzar fin completo y evitar truncado
        )
    
    response = tokenizer.decode(
        outputs[0][inputs.input_ids.shape[1]:], 
        skip_special_tokens=True
    ).strip()
    
    model.to('cpu')
    torch.cuda.empty_cache()
    
    return response

# Interfaz estándar genérica con Gradio ChatInterface (sin CSS personalizado ni Blocks, como pide la documentación oficial de Gradio)
demo = gr.ChatInterface(
    fn=generate_response,
    title="Desxtools Crypto News Text Analyzer",
    description="Chatbot estándar con google/gemma-3-4b-it en ZeroGPU. Diseño genérico sin modificaciones.",
    cache_examples=True
)

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