import os
import sys
import torch
import pickle
import time
import gc
import asyncio
import aiohttp
from typing import Optional, Dict, List
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from huggingface_hub import snapshot_download
import uvicorn
import wikipedia
from duckduckgo_search import DDGS

# ======================
# CONFIGURACIÓN DE DISPOSITIVO
# ======================
if torch.cuda.is_available():
    DEVICE = "cuda"
    print("✅ GPU NVIDIA detectada. Usando CUDA.")
else:
    DEVICE = "cpu"
    print("⚠️ GPU no detectada. Usando CPU (puede ser más lento).")

if DEVICE == "cpu":
    torch.set_num_threads(max(1, os.cpu_count() // 2))

torch.set_grad_enabled(False)

MODEL_REPO = "TeszenAI/MTPw"

# ======================
# MOTOR DE BÚSQUEDA WEB
# ======================
class WebSearchEngine:
    """Motor de búsqueda integrado con Wikipedia y DuckDuckGo"""
    
    def __init__(self):
        self.ddgs = DDGS()
        wikipedia.set_lang('es')
    
    async def search_wikipedia(self, query: str, sentences: int = 4) -> Optional[Dict]:
        """Buscar en Wikipedia"""
        try:
            search_results = wikipedia.search(query, results=3)
            if not search_results:
                return None
            
            page = wikipedia.page(search_results[0], auto_suggest=False)
            summary = wikipedia.summary(search_results[0], sentences=sentences)
            
            return {
                "source": "Wikipedia",
                "title": page.title,
                "url": page.url,
                "summary": summary,
                "success": True
            }
        except wikipedia.exceptions.DisambiguationError as e:
            try:
                page = wikipedia.page(e.options[0], auto_suggest=False)
                summary = wikipedia.summary(e.options[0], sentences=sentences)
                return {
                    "source": "Wikipedia",
                    "title": page.title,
                    "url": page.url,
                    "summary": summary,
                    "success": True
                }
            except:
                return None
        except:
            return None
    
    async def search_duckduckgo(self, query: str, max_results: int = 5) -> List[Dict]:
        """Buscar en DuckDuckGo"""
        try:
            results = []
            ddg_results = self.ddgs.text(query, max_results=max_results)
            
            for r in ddg_results:
                results.append({
                    "title": r.get("title", ""),
                    "url": r.get("href", ""),
                    "snippet": r.get("body", "")
                })
            
            return results
        except Exception as e:
            print(f"Error en DuckDuckGo: {e}")
            return []
    
    async def search(self, query: str) -> Dict:
        """Búsqueda combinada"""
        wiki_task = self.search_wikipedia(query)
        ddg_task = self.search_duckduckgo(query)
        
        wiki_result, ddg_results = await asyncio.gather(wiki_task, ddg_task)
        
        return {
            "query": query,
            "wikipedia": wiki_result,
            "web_results": ddg_results,
            "timestamp": time.time()
        }

# Inicializar motor de búsqueda
search_engine = WebSearchEngine()

# ======================
# DESCARGA Y CARGA DEL MODELO
# ======================
print(f"📦 Descargando modelo desde {MODEL_REPO}...")
repo_path = snapshot_download(
    repo_id=MODEL_REPO,
    repo_type="model",
    local_dir="mtptz_repo"
)

sys.path.insert(0, repo_path)

from model import MTPMiniModel
from tokenizer import MTPTokenizer

print("🔧 Cargando tensores y configuración...")
with open(os.path.join(repo_path, "mtp_mini.pkl"), "rb") as f:
    model_data = pickle.load(f)

tokenizer = MTPTokenizer(os.path.join(repo_path, "mtp_tokenizer.model"))
VOCAB_SIZE = tokenizer.sp.get_piece_size()
config = model_data["config"]

use_swiglu = config["model"].get("use_swiglu", False)

print(f"🧠 Inicializando modelo...")
print(f"   → Vocabulario: {VOCAB_SIZE}")
print(f"   → Dimensión: {config['model']['d_model']}")
print(f"   → Capas: {config['model']['n_layers']}")
print(f"   → Cabezas: {config['model']['n_heads']}")
print(f"   → SwiGLU: {'✓' if use_swiglu else '✗'}")

model = MTPMiniModel(
    vocab_size=VOCAB_SIZE,
    d_model=config["model"]["d_model"],
    n_layers=config["model"]["n_layers"],
    n_heads=config["model"]["n_heads"],
    d_ff=config["model"]["d_ff"],
    max_seq_len=config["model"]["max_seq_len"],
    dropout=0.0,
    use_swiglu=use_swiglu
)

model.load_state_dict(model_data["model_state_dict"])
model.eval()

if DEVICE == "cpu":
    print("⚡ Aplicando cuantización dinámica para CPU...")
    model = torch.quantization.quantize_dynamic(
        model, 
        {torch.nn.Linear}, 
        dtype=torch.qint8
    )

model.to(DEVICE)

param_count = sum(p.numel() for p in model.parameters())
print(f"✅ Modelo cargado: {param_count:,} parámetros ({param_count/1e6:.1f}M)")
print(f"🔍 Motor de búsqueda web inicializado (Wikipedia + DuckDuckGo)")

# ======================
# API CONFIG
# ======================
app = FastAPI(
    title="MTP-3.5 Enhanced API",
    description="API mejorada con capacidades de búsqueda web (Wikipedia + DuckDuckGo)",
    version="3.5-web"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

class PromptRequest(BaseModel):
    text: str = Field(..., max_length=2000, description="Texto de entrada")
    max_tokens: int = Field(default=150, ge=10, le=300, description="Tokens máximos a generar")
    temperature: float = Field(default=0.7, ge=0.1, le=2.0, description="Temperatura de muestreo")
    top_k: int = Field(default=40, ge=1, le=100, description="Top-k sampling")
    top_p: float = Field(default=0.92, ge=0.1, le=1.0, description="Top-p (nucleus) sampling")
    repetition_penalty: float = Field(default=1.15, ge=1.0, le=2.0, description="Penalización por repetición")
    min_length: int = Field(default=20, ge=5, le=100, description="Longitud mínima de respuesta")
    use_web_search: bool = Field(default=False, description="Activar búsqueda web")

class SearchRequest(BaseModel):
    query: str = Field(..., max_length=500, description="Consulta de búsqueda")

def build_prompt(user_input: str, web_context: Optional[str] = None) -> str:
    """Construye el prompt con contexto web opcional"""
    if web_context:
        # Formato especial para búsqueda web con instrucciones claras
        return f"""### Instrucción del sistema:
Eres un asistente que busca información en internet. Debes resumir la información encontrada de forma clara y útil.

### Información encontrada en la web:
{web_context}

### Pregunta del usuario:
{user_input}

### Respuesta:
Hola, encontré esto en la web:

"""
    return f"### Instrucción:\n{user_input}\n\n### Respuesta:\n"

def format_search_results(search_data: Dict) -> str:
    """Formatea resultados de búsqueda para el contexto del modelo"""
    context_parts = []
    
    if search_data.get("wikipedia") and search_data["wikipedia"].get("success"):
        wiki = search_data["wikipedia"]
        context_parts.append(f"[Wikipedia - {wiki['title']}]\n{wiki['summary']}\nFuente: {wiki['url']}")
    
    if search_data.get("web_results"):
        for i, result in enumerate(search_data["web_results"][:4], 1):
            snippet = result['snippet'][:300].strip()
            context_parts.append(f"[Resultado {i}: {result['title']}]\n{snippet}\nFuente: {result['url']}")
    
    return "\n\n".join(context_parts) if context_parts else ""

# ======================
# ⚡ GESTIÓN DE CARGA
# ======================
ACTIVE_REQUESTS = 0
MAX_CONCURRENT_REQUESTS = 3

@app.post("/search")
async def web_search(req: SearchRequest):
    """Endpoint de búsqueda web"""
    try:
        search_results = await search_engine.search(req.query)
        
        formatted_context = format_search_results(search_results)
        
        return {
            "query": req.query,
            "results": search_results,
            "formatted_context": formatted_context,
            "has_results": bool(formatted_context),
            "sources_used": []
        }
    except Exception as e:
        print(f"Error en búsqueda: {e}")
        return {
            "query": req.query,
            "error": str(e),
            "has_results": False
        }

@app.post("/generate")
async def generate(req: PromptRequest):
    """Endpoint principal con búsqueda web integrada"""
    global ACTIVE_REQUESTS
    
    if ACTIVE_REQUESTS >= MAX_CONCURRENT_REQUESTS:
        return {
            "reply": "El servidor está ocupado. Por favor, intenta de nuevo en unos segundos.",
            "error": "too_many_requests",
            "active_requests": ACTIVE_REQUESTS
        }
    
    ACTIVE_REQUESTS += 1
    
    dyn_max_tokens = req.max_tokens
    dyn_temperature = req.temperature
    
    if ACTIVE_REQUESTS > 1:
        print(f"⚠️ Carga alta ({ACTIVE_REQUESTS} requests). Ajustando parámetros.")
        dyn_max_tokens = min(dyn_max_tokens, 120)
        dyn_temperature = max(0.6, dyn_temperature * 0.95)

    user_input = req.text.strip()
    if not user_input:
        ACTIVE_REQUESTS -= 1
        return {"reply": "", "tokens_generated": 0}

    web_context = ""
    search_results = None
    
    # Realizar búsqueda web si está activada
    if req.use_web_search:
        try:
            # Extraer la consulta de búsqueda del mensaje del usuario
            search_query = user_input
            
            # Si el mensaje es una pregunta larga, extraer palabras clave
            if len(user_input.split()) > 8:
                # Usar las primeras palabras más relevantes
                words = user_input.lower().split()
                # Filtrar palabras comunes
                stop_words = {'qué', 'cuál', 'cómo', 'dónde', 'cuándo', 'por', 'para', 'el', 'la', 'los', 'las', 'un', 'una', 'es', 'sobre', 'me', 'puedes', 'explicar', 'decir', 'información'}
                keywords = [w for w in words if w not in stop_words][:5]
                search_query = ' '.join(keywords)
            
            search_results = await search_engine.search(search_query)
            web_context = format_search_results(search_results)
            
            if web_context:
                print(f"🔍 Búsqueda web completada para: '{search_query}'")
                print(f"   Contexto agregado: {len(web_context)} caracteres")
        except Exception as e:
            print(f"Error en búsqueda web: {e}")

    full_prompt = build_prompt(user_input, web_context if web_context else None)
    tokens = [tokenizer.bos_id()] + tokenizer.encode(full_prompt)
    input_ids = torch.tensor([tokens], device=DEVICE)

    try:
        start_time = time.time()
        
        with torch.no_grad():
            output_ids = model.generate(
                input_ids,
                max_new_tokens=dyn_max_tokens,
                temperature=dyn_temperature,
                top_k=req.top_k,
                top_p=req.top_p,
                repetition_penalty=req.repetition_penalty,
                min_length=req.min_length,
                eos_token_id=tokenizer.eos_id()
            )

        gen_tokens = output_ids[0, len(tokens):].tolist()
        
        safe_tokens = []
        for t in gen_tokens:
            if 0 <= t < VOCAB_SIZE and t != tokenizer.eos_id():
                safe_tokens.append(t)
            elif t == tokenizer.eos_id():
                break
        
        response = tokenizer.decode(safe_tokens).strip()
        
        if "###" in response:
            response = response.split("###")[0].strip()
        
        if response.endswith(("...", ". . .", "…")):
            response = response.rstrip(".")
        
        generation_time = time.time() - start_time
        tokens_per_second = len(safe_tokens) / generation_time if generation_time > 0 else 0

        result = {
            "reply": response,
            "tokens_generated": len(safe_tokens),
            "generation_time": round(generation_time, 2),
            "tokens_per_second": round(tokens_per_second, 1),
            "model": "MTP-3.5",
            "device": DEVICE,
            "web_search_used": req.use_web_search
        }
        
        if req.use_web_search and search_results:
            sources = []
            if search_results.get("wikipedia") and search_results["wikipedia"].get("success"):
                sources.append({
                    "type": "wikipedia",
                    "title": search_results["wikipedia"]["title"],
                    "url": search_results["wikipedia"]["url"]
                })
            
            if search_results.get("web_results"):
                for r in search_results["web_results"][:3]:
                    sources.append({
                        "type": "web",
                        "title": r["title"],
                        "url": r["url"]
                    })
            
            result["sources"] = sources
            result["search_query"] = user_input

        return result
    
    except Exception as e:
        print(f"❌ Error durante generación: {e}")
        import traceback
        traceback.print_exc()
        return {
            "reply": "Lo siento, ocurrió un error al procesar tu solicitud.",
            "error": str(e)
        }
    
    finally:
        ACTIVE_REQUESTS -= 1
        if DEVICE == "cuda":
            torch.cuda.empty_cache()
        gc.collect()

# ======================
# 📊 ENDPOINTS DE INFORMACIÓN
# ======================
@app.get("/health")
def health_check():
    """Check del estado del servicio"""
    memory_info = {}
    if DEVICE == "cuda":
        memory_info = {
            "gpu_memory_allocated_mb": round(torch.cuda.memory_allocated() / 1024**2, 2),
            "gpu_memory_reserved_mb": round(torch.cuda.memory_reserved() / 1024**2, 2)
        }
    
    return {
        "status": "healthy",
        "model": "MTP-3.5-Web",
        "device": DEVICE,
        "active_requests": ACTIVE_REQUESTS,
        "max_concurrent_requests": MAX_CONCURRENT_REQUESTS,
        "vocab_size": VOCAB_SIZE,
        "parameters": sum(p.numel() for p in model.parameters()),
        "web_search_enabled": True,
        **memory_info
    }

@app.get("/info")
def model_info():
    """Información detallada del modelo"""
    improvements = [
        "RoPE (Rotary Position Embedding)",
        "RMSNorm (Root Mean Square Normalization)",
        "Label Smoothing (0.1)",
        "Repetition Penalty",
        "Early Stopping",
        "EOS Loss Weight",
        "Length Control",
        "Gradient Accumulation",
        "Web Search Integration (Wikipedia + DuckDuckGo)"
    ]
    
    if config["model"].get("use_swiglu", False):
        improvements.append("SwiGLU Activation")
    
    return {
        "model_name": "MTP-3.5-Web",
        "version": "3.5-web",
        "architecture": {
            "d_model": config["model"]["d_model"],
            "n_layers": config["model"]["n_layers"],
            "n_heads": config["model"]["n_heads"],
            "d_ff": config["model"]["d_ff"],
            "max_seq_len": config["model"]["max_seq_len"],
            "vocab_size": VOCAB_SIZE,
            "use_swiglu": config["model"].get("use_swiglu", False),
            "dropout": config["model"]["dropout"]
        },
        "parameters": sum(p.numel() for p in model.parameters()),
        "parameters_human": f"{sum(p.numel() for p in model.parameters())/1e6:.1f}M",
        "device": DEVICE,
        "improvements": improvements,
        "web_search": {
            "enabled": True,
            "sources": ["Wikipedia (ES)", "DuckDuckGo"]
        }
    }

# ======================
# 🎨 INTERFAZ WEB MEJORADA CON BÚSQUEDA
# ======================
@app.get("/", response_class=HTMLResponse)
def chat_ui():
    return """
<!DOCTYPE html>
<html lang="es">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no">
<title>MTP 3.5 Web - Chat Interface</title>
<link rel="preconnect" href="https://fonts.googleapis.com">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600&display=swap" rel="stylesheet">
<style>
:root {
    --bg-color: #131314;
    --surface-color: #1E1F20;
    --accent-color: #4a9eff;
    --text-primary: #e3e3e3;
    --text-secondary: #9aa0a6;
    --user-bubble: #282a2c;
    --success-color: #34a853;
    --warning-color: #fbbc04;
    --error-color: #ea4335;
    --search-active: #ff6b35;
    --logo-url: url('https://i.postimg.cc/yxS54PF3/IMG-3082.jpg');
}
* { 
    box-sizing: border-box; 
    outline: none; 
    -webkit-tap-highlight-color: transparent; 
}
body {
    margin: 0;
    background-color: var(--bg-color);
    font-family: 'Inter', sans-serif;
    color: var(--text-primary);
    height: 100dvh;
    display: flex;
    flex-direction: column;
    overflow: hidden;
}
header {
    padding: 12px 20px;
    display: flex;
    align-items: center;
    justify-content: space-between;
    background: rgba(19, 19, 20, 0.85);
    backdrop-filter: blur(12px);
    position: fixed;
    top: 0;
    width: 100%;
    z-index: 50;
    border-bottom: 1px solid rgba(255,255,255,0.05);
}
.brand-wrapper {
    display: flex;
    align-items: center;
    gap: 12px;
    cursor: pointer;
}
.brand-logo {
    width: 32px;
    height: 32px;
    border-radius: 50%;
    background-image: var(--logo-url);
    background-size: cover;
    background-position: center;
    border: 1px solid rgba(255,255,255,0.1);
}
.brand-text {
    font-weight: 500;
    font-size: 1.05rem;
    display: flex;
    align-items: center;
    gap: 8px;
}
.version-badge {
    font-size: 0.75rem;
    background: rgba(74, 158, 255, 0.15);
    color: #8ab4f8;
    padding: 2px 8px;
    border-radius: 12px;
    font-weight: 600;
}
.web-badge {
    font-size: 0.7rem;
    background: rgba(255, 107, 53, 0.15);
    color: #ff8a65;
    padding: 2px 6px;
    border-radius: 8px;
    font-weight: 600;
}
.status-indicator {
    width: 8px;
    height: 8px;
    border-radius: 50%;
    background: var(--success-color);
    animation: pulse 2s infinite;
}
@keyframes pulse {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.5; }
}
.chat-scroll {
    flex: 1;
    overflow-y: auto;
    padding: 80px 20px 40px 20px;
    display: flex;
    flex-direction: column;
    gap: 30px;
    max-width: 850px;
    margin: 0 auto;
    width: 100%;
    scroll-behavior: smooth;
}
.msg-row {
    display: flex;
    gap: 16px;
    width: 100%;
    opacity: 0;
    transform: translateY(10px);
    animation: slideUpFade 0.4s cubic-bezier(0.2, 0.8, 0.2, 1) forwards;
}
.msg-row.user { justify-content: flex-end; }
.msg-row.bot { justify-content: flex-start; align-items: flex-start; }
.msg-content {
    line-height: 1.6;
    font-size: 1rem;
    word-wrap: break-word;
    max-width: 85%;
}
.user .msg-content {
    background-color: var(--user-bubble);
    padding: 10px 18px;
    border-radius: 18px;
    border-top-right-radius: 4px;
    color: #fff;
}
.bot .msg-content-wrapper {
    display: flex;
    flex-direction: column;
    gap: 8px;
    width: 100%;
}
.bot .msg-text {
    padding-top: 6px;
    color: var(--text-primary);
    white-space: pre-wrap;
}
.bot-avatar {
    width: 34px;
    height: 34px;
    min-width: 34px;
    border-radius: 50%;
    background-image: var(--logo-url);
    background-size: cover;
    box-shadow: 0 2px 6px rgba(0,0,0,0.2);
}
.search-indicator {
    display: inline-flex;
    align-items: center;
    gap: 6px;
    font-size: 0.85rem;
    color: var(--search-active);
    padding: 4px 10px;
    background: rgba(255, 107, 53, 0.1);
    border-radius: 12px;
    margin-bottom: 8px;
}
.search-indicator svg {
    width: 14px;
    height: 14px;
    fill: currentColor;
    animation: spin 1s linear infinite;
}
@keyframes spin {
    from { transform: rotate(0deg); }
    to { transform: rotate(360deg); }
}
.sources-list {
    margin-top: 10px;
    padding: 10px;
    background: rgba(74, 158, 255, 0.05);
    border-left: 3px solid var(--accent-color);
    border-radius: 6px;
    font-size: 0.85rem;
}
.sources-list .source-title {
    font-weight: 600;
    color: var(--accent-color);
    margin-bottom: 6px;
}
.sources-list a {
    color: #8ab4f8;
    text-decoration: none;
    display: block;
    margin: 4px 0;
    transition: color 0.2s;
}
.sources-list a:hover {
    color: var(--accent-color);
    text-decoration: underline;
}
.bot-actions {
    display: flex;
    gap: 10px;
    opacity: 0;
    transition: opacity 0.3s;
    margin-top: 5px;
}
.action-btn {
    background: transparent;
    border: none;
    color: var(--text-secondary);
    cursor: pointer;
    padding: 4px;
    border-radius: 4px;
    display: flex;
    align-items: center;
    transition: color 0.2s, background 0.2s;
    font-size: 0.85rem;
}
.action-btn:hover {
    color: var(--text-primary);
    background: rgba(255,255,255,0.08);
}
.action-btn svg { 
    width: 16px; 
    height: 16px; 
    fill: currentColor; 
    margin-right: 4px;
}
.typing-cursor::after {
    content: '';
    display: inline-block;
    width: 10px;
    height: 10px;
    background: var(--accent-color);
    border-radius: 50%;
    margin-left: 5px;
    vertical-align: middle;
    animation: blink 1s infinite;
}
.footer-container {
    padding: 0 20px 20px 20px;
    background: linear-gradient(to top, var(--bg-color) 85%, transparent);
    position: relative;
    z-index: 60;
}
.search-toggle-wrapper {
    max-width: 850px;
    margin: 0 auto 10px auto;
    display: flex;
    align-items: center;
    gap: 10px;
    padding: 8px 16px;
    background: var(--surface-color);
    border-radius: 16px;
    border: 1px solid rgba(255,255,255,0.08);
}
.search-toggle {
    display: flex;
    align-items: center;
    gap: 8px;
    cursor: pointer;
    user-select: none;
}
.toggle-switch {
    position: relative;
    width: 40px;
    height: 22px;
    background: rgba(255,255,255,0.1);
    border-radius: 11px;
    transition: background 0.3s;
}
.toggle-switch.active {
    background: var(--search-active);
}
.toggle-switch::after {
    content: '';
    position: absolute;
    width: 18px;
    height: 18px;
    background: white;
    border-radius: 50%;
    top: 2px;
    left: 2px;
    transition: left 0.3s;
}
.toggle-switch.active::after {
    left: 20px;
}
.toggle-label {
    font-size: 0.9rem;
    color: var(--text-secondary);
    display: flex;
    align-items: center;
    gap: 6px;
}
.toggle-label svg {
    width: 16px;
    height: 16px;
    fill: currentColor;
}
.toggle-switch.active + .toggle-label {
    color: var(--search-active);
}
.input-box {
    max-width: 850px;
    margin: 0 auto;
    background: var(--surface-color);
    border-radius: 28px;
    padding: 8px 10px 8px 20px;
    display: flex;
    align-items: center;
    border: 1px solid rgba(255,255,255,0.1);
    transition: border-color 0.2s, box-shadow 0.2s;
}
.input-box:focus-within {
    border-color: rgba(74, 158, 255, 0.5);
    box-shadow: 0 0 0 2px rgba(74, 158, 255, 0.1);
}
#userInput {
    flex: 1;
    background: transparent;
    border: none;
    color: white;
    font-size: 1rem;
    font-family: inherit;
    padding: 10px 0;
    resize: none;
    max-height: 120px;
}
#mainBtn {
    background: white;
    color: black;
    border: none;
    width: 36px;
    height: 36px;
    border-radius: 50%;
    display: flex;
    align-items: center;
    justify-content: center;
    cursor: pointer;
    margin-left: 8px;
    transition: transform 0.2s;
}
#mainBtn:hover { transform: scale(1.05); }
#mainBtn:disabled {
    opacity: 0.5;
    cursor: not-allowed;
}
.disclaimer {
    text-align: center;
    font-size: 0.75rem;
    color: #666;
    margin-top: 12px;
}
.stats-badge {
    font-size: 0.7rem;
    color: var(--text-secondary);
    margin-top: 4px;
    font-family: 'Monaco', monospace;
}
@keyframes slideUpFade {
    from { opacity: 0; transform: translateY(15px); }
    to { opacity: 1; transform: translateY(0); }
}
@keyframes blink { 
    0%, 100% { opacity: 1; } 
    50% { opacity: 0; } 
}
@keyframes pulseAvatar {
    0% { box-shadow: 0 0 0 0 rgba(74, 158, 255, 0.4); }
    70% { box-shadow: 0 0 0 8px rgba(74, 158, 255, 0); }
    100% { box-shadow: 0 0 0 0 rgba(74, 158, 255, 0); }
}
.pulsing { animation: pulseAvatar 1.5s infinite; }
::-webkit-scrollbar { width: 8px; }
::-webkit-scrollbar-track { background: transparent; }
::-webkit-scrollbar-thumb { background: #333; border-radius: 4px; }
</style>
</head>
<body>
<header>
    <div class="brand-wrapper" onclick="location.reload()">
        <div class="brand-logo"></div>
        <div class="brand-text">
            MTP <span class="version-badge">3.5</span> <span class="web-badge">WEB</span>
        </div>
    </div>
    <div class="status-indicator" title="Sistema operativo"></div>
</header>
<div id="chatScroll" class="chat-scroll">
    <div class="msg-row bot" style="animation-delay: 0.1s;">
        <div class="bot-avatar"></div>
        <div class="msg-content-wrapper">
            <div class="msg-text">¡Hola! Soy MTP 3.5 Web, un modelo mejorado con capacidades de búsqueda en internet.

🔍 Nueva función: Búsqueda Web
• Activa el botón "Búsqueda Web" para que pueda buscar información en Wikipedia y DuckDuckGo
• Te proporcionaré respuestas basadas en información actualizada de internet
• Incluiré las fuentes consultadas en cada respuesta

Características del modelo:
• RoPE (Rotary Position Embedding)
• RMSNorm para estabilidad
• Control de repetición inteligente
• Generación coherente y fluida

¿En qué puedo ayudarte hoy?</div>
        </div>
    </div>
</div>
<div class="footer-container">
    <div class="search-toggle-wrapper">
        <div class="search-toggle" onclick="toggleWebSearch()">
            <div class="toggle-switch" id="searchToggle"></div>
            <div class="toggle-label">
                <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                    <circle cx="11" cy="11" r="8"></circle>
                    <path d="m21 21-4.35-4.35"></path>
                </svg>
                <span id="searchLabel">Búsqueda Web: Desactivada</span>
            </div>
        </div>
    </div>
    <div class="input-box">
        <textarea id="userInput" placeholder="Escribe un mensaje..." rows="1" autocomplete="off"></textarea>
        <button id="mainBtn" onclick="handleBtnClick()"></button>
    </div>
    <div class="disclaimer">
        MTP 3.5 Web puede cometer errores. Verifica la información importante en las fuentes originales.
    </div>
</div>
<script>
const chatScroll = document.getElementById('chatScroll');
const userInput = document.getElementById('userInput');
const mainBtn = document.getElementById('mainBtn');
const searchToggle = document.getElementById('searchToggle');
const searchLabel = document.getElementById('searchLabel');

let isGenerating = false;
let webSearchEnabled = false;
let lastUserPrompt = "";

const ICON_SEND = `<svg width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M22 2L11 13M22 2l-7 20-4-9-9-4 20-7z"></path></svg>`;
const ICON_STOP = `<svg width="14" height="14" viewBox="0 0 24 24" fill="currentColor"><rect x="2" y="2" width="20" height="20" rx="4"></rect></svg>`;
mainBtn.innerHTML = ICON_SEND;

userInput.addEventListener('input', function() {
    this.style.height = 'auto';
    this.style.height = Math.min(this.scrollHeight, 120) + 'px';
});

function toggleWebSearch() {
    webSearchEnabled = !webSearchEnabled;
    searchToggle.classList.toggle('active', webSearchEnabled);
    searchLabel.textContent = webSearchEnabled ? 'Búsqueda Web: Activada' : 'Búsqueda Web: Desactivada';
}

function scrollToBottom() {
    chatScroll.scrollTop = chatScroll.scrollHeight;
}

function setBtnState(state) {
    if (state === 'sending') {
        mainBtn.innerHTML = ICON_STOP;
        mainBtn.disabled = false;
        isGenerating = true;
    } else if (state === 'disabled') {
        mainBtn.disabled = true;
        isGenerating = false;
    } else {
        mainBtn.innerHTML = ICON_SEND;
        mainBtn.disabled = false;
        isGenerating = false;
    }
}

function handleBtnClick() {
    if (isGenerating) {
        stopGeneration();
    } else {
        sendMessage();
    }
}

function stopGeneration() {
    const activeCursor = document.querySelector('.typing-cursor');
    if (activeCursor) activeCursor.classList.remove('typing-cursor');
    const activeAvatar = document.querySelector('.pulsing');
    if (activeAvatar) activeAvatar.classList.remove('pulsing');
    setBtnState('idle');
    userInput.focus();
}

async function sendMessage(textOverride = null) {
    const text = textOverride || userInput.value.trim();
    if (!text) return;
    
    lastUserPrompt = text;
    
    if (!textOverride) {
        userInput.value = '';
        userInput.style.height = 'auto';
        addMessage(text, 'user');
    }
    
    setBtnState('sending');
    
    const botRow = document.createElement('div');
    botRow.className = 'msg-row bot';
    
    const avatar = document.createElement('div');
    avatar.className = 'bot-avatar pulsing'; 
    
    const wrapper = document.createElement('div');
    wrapper.className = 'msg-content-wrapper';
    
    if (webSearchEnabled) {
        const searchInd = document.createElement('div');
        searchInd.className = 'search-indicator';
        searchInd.innerHTML = `
            <svg viewBox="0 0 24 24">
                <circle cx="12" cy="12" r="10" fill="none" stroke="currentColor" stroke-width="2"></circle>
                <path d="M12 6v6l4 2" stroke="currentColor" stroke-width="2" fill="none"></path>
            </svg>
            <span>Buscando en la web...</span>
        `;
        wrapper.appendChild(searchInd);
    }
    
    const msgText = document.createElement('div');
    msgText.className = 'msg-text'; 
    
    wrapper.appendChild(msgText);
    botRow.appendChild(avatar);
    botRow.appendChild(wrapper);
    chatScroll.appendChild(botRow);
    scrollToBottom();
    
    try {
        const startTime = performance.now();
        
        const response = await fetch('/generate', {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify({ 
                text: text,
                max_tokens: 200,
                temperature: 0.7,
                top_k: 40,
                top_p: 0.92,
                repetition_penalty: 1.15,
                min_length: 20,
                use_web_search: webSearchEnabled
            })
        });
        
        const data = await response.json();
        
        if (!isGenerating) return; 
        
        avatar.classList.remove('pulsing');
        
        const searchInd = wrapper.querySelector('.search-indicator');
        if (searchInd) searchInd.remove();
        
        if (data.error) {
            msgText.innerHTML = `<span style="color: var(--error-color);">Error: ${data.error}</span>`;
            setBtnState('idle');
            return;
        }
        
        const reply = data.reply || "No entendí eso.";
        const endTime = performance.now();
        const totalTime = ((endTime - startTime) / 1000).toFixed(2);
        
        await typeWriter(msgText, reply);
        
        if (isGenerating) {
            if (data.sources && data.sources.length > 0) {
                const sourcesDiv = document.createElement('div');
                sourcesDiv.className = 'sources-list';
                sourcesDiv.innerHTML = '<div class="source-title">📚 Fuentes consultadas:</div>';
                
                data.sources.forEach(source => {
                    const link = document.createElement('a');
                    link.href = source.url;
                    link.target = '_blank';
                    link.textContent = `• ${source.title}`;
                    sourcesDiv.appendChild(link);
                });
                
                wrapper.appendChild(sourcesDiv);
            }
            
            const stats = document.createElement('div');
            stats.className = 'stats-badge';
            let statsText = `${data.tokens_generated} tokens • ${data.tokens_per_second} t/s • ${totalTime}s • ${data.device}`;
            if (data.web_search_used) {
                statsText += ' • 🔍 Web';
            }
            stats.textContent = statsText;
            wrapper.appendChild(stats);
            
            addActions(wrapper, reply);
            setBtnState('idle');
        }
    } catch (error) {
        console.error('Error:', error);
        avatar.classList.remove('pulsing');
        msgText.innerHTML = `<span style="color: var(--error-color);">Error de conexión. Por favor, intenta de nuevo.</span>`;
        setBtnState('idle');
    }
}

function addMessage(text, sender) {
    const row = document.createElement('div');
    row.className = `msg-row ${sender}`;
    
    const content = document.createElement('div');
    content.className = 'msg-content';
    content.textContent = text;
    
    row.appendChild(content);
    chatScroll.appendChild(row);
    scrollToBottom();
}

function typeWriter(element, text, speed = 12) {
    return new Promise(resolve => {
        let i = 0;
        element.classList.add('typing-cursor');
        
        function type() {
            if (!isGenerating) {
                element.classList.remove('typing-cursor');
                resolve();
                return;
            }
            
            if (i < text.length) {
                element.textContent += text.charAt(i);
                i++;
                scrollToBottom();
                setTimeout(type, speed + Math.random() * 5);
            } else {
                element.classList.remove('typing-cursor');
                resolve();
            }
        }
        
        type();
    });
}

function addActions(wrapperElement, textToCopy) {
    const actionsDiv = document.createElement('div');
    actionsDiv.className = 'bot-actions';
    
    const copyBtn = document.createElement('button');
    copyBtn.className = 'action-btn';
    copyBtn.innerHTML = `<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect><path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path></svg>Copiar`;
    copyBtn.onclick = () => {
        navigator.clipboard.writeText(textToCopy).then(() => {
            copyBtn.innerHTML = `<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><polyline points="20 6 9 17 4 12"></polyline></svg>Copiado`;
            setTimeout(() => {
                copyBtn.innerHTML = `<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect><path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path></svg>Copiar`;
            }, 2000);
        });
    };
    
    const regenBtn = document.createElement('button');
    regenBtn.className = 'action-btn';
    regenBtn.innerHTML = `<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M23 4v6h-6"></path><path d="M1 20v-6h6"></path><path d="M3.51 9a9 9 0 0 1 14.85-3.36L23 10M1 14l4.64 4.36A9 9 0 0 0 20.49 15"></path></svg>Regenerar`;
    regenBtn.onclick = () => {
        sendMessage(lastUserPrompt);
    };
    
    actionsDiv.appendChild(copyBtn);
    actionsDiv.appendChild(regenBtn);
    wrapperElement.appendChild(actionsDiv);
    
    requestAnimationFrame(() => actionsDiv.style.opacity = "1");
    scrollToBottom();
}

userInput.addEventListener('keydown', (e) => {
    if (e.key === 'Enter' && !e.shiftKey) {
        e.preventDefault();
        handleBtnClick();
    }
});

window.onload = () => {
    userInput.focus();
    
    fetch('/info')
        .then(r => r.json())
        .then(data => {
            console.log('Modelo cargado:', data);
        })
        .catch(e => console.error('Error cargando info:', e));
};
</script>
</body>
</html>
"""

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    print(f"\n🚀 Iniciando servidor MTP-3.5 Web...")
    print(f"🌐 Interfaz web: http://0.0.0.0:{port}")
    print(f"📡 API docs: http://0.0.0.0:{port}/docs")
    print(f"📊 Health check: http://0.0.0.0:{port}/health")
    print(f"ℹ️  Model info: http://0.0.0.0:{port}/info")
    print(f"🔍 Búsqueda web: Wikipedia + DuckDuckGo")
    print(f"\n✅ Sistema listo. Presiona Ctrl+C para detener.")
    
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=port,
        log_level="info"
    )