import spaces  # Debe permanecer en la línea 1 para ZeroGPU

import os
import re
import logging
import gradio as ui
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from functools import lru_cache
from typing import Optional

# ─────────────────────────────────────────────────────────────────────────────
# Logging
# ─────────────────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
log = logging.getLogger(__name__)

# ─────────────────────────────────────────────────────────────────────────────
# Constantes de configuración
# ─────────────────────────────────────────────────────────────────────────────
MODELO = "DarksitoBest/DRK-Coder-V1.1"
HF_TOKEN = os.environ.get("HF_TOKEN")
MAX_HISTORIAL = 20
GPU_TIMEOUT_BUFFER = 10          
FALLBACK_CONTEXT = 32_768
MAX_CONTEXT_SANE  = 1_000_000   
DEFAULT_REPETITION_PENALTY = 1.08

if not HF_TOKEN:
    raise RuntimeError(
        "No se encontró HF_TOKEN. Añádelo en Settings → Variables and secrets "
        "del Space y márcalo como Secret."
    )

# ─────────────────────────────────────────────────────────────────────────────
# Carga del modelo
# ─────────────────────────────────────────────────────────────────────────────
log.info("Cargando tokenizer de %s …", MODELO)
tokenizer = AutoTokenizer.from_pretrained(
    MODELO,
    token=HF_TOKEN,
    trust_remote_code=True,
)

log.info("Cargando modelo …")
model = AutoModelForCausalLM.from_pretrained(
    MODELO,
    token=HF_TOKEN,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    low_cpu_mem_usage=True,
)
model.eval()

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

def _calc_context_limit() -> int:
    for attr in ("max_position_embeddings",):
        val = getattr(model.config, attr, None)
        if val and val <= MAX_CONTEXT_SANE:
            return int(val)
    val = getattr(tokenizer, "model_max_length", None)
    if val and val <= MAX_CONTEXT_SANE:
        return int(val)
    return FALLBACK_CONTEXT

CONTEXT_LIMIT: int = _calc_context_limit()
log.info("Límite de contexto detectado: %d tokens", CONTEXT_LIMIT)

# ─────────────────────────────────────────────────────────────────────────────
# Prompts del sistema
# ─────────────────────────────────────────────────────────────────────────────
SYSTEM_PROMPTS: dict[str, str] = {
    "en": (
        "You are DRK Code V1, a senior programming assistant.\n"
        "Always answer in English. Provide correct, secure, and executable solutions.\n"
        "When writing code, use Markdown code blocks with the language specified. Be clear and direct.\n"
        "When useful, add a brief reasoning summary inside <think>...</think> before the final answer. "
        "Do not expose hidden chain-of-thought.\n"
        "Put the main deliverable in one complete fenced code block so it can be opened in Canvas.\n"
        "If important information is missing, ask before making assumptions. Never invent APIs or results."
    ),
    "es": (
        "Eres DRK Code V1, un asistente senior de programación.\n"
        "Responde siempre en español. Entrega soluciones correctas, seguras y ejecutables.\n"
        "Cuando escribas código, usa bloques Markdown con el lenguaje indicado. Sé claro y directo.\n"
        "Cuando sea útil, añade un resumen breve del razonamiento dentro de <think>...</think> antes "
        "de la respuesta final. No expongas razonamiento interno oculto.\n"
        "Coloca el entregable principal en un único bloque de código completo para abrirlo en Canvas.\n"
        "Si faltan datos importantes, pregunta antes de asumir. No inventes APIs ni resultados."
    ),
}

LANGUAGE_ALIASES: dict[str, Optional[str]] = {
    "py": "python", "python3": "python", "js": "javascript", "jsx": "javascript",
    "ts": "typescript", "tsx": "typescript", "htm": "html", "yml": "yaml",
    "bash": "shell", "sh": "shell", "zsh": "shell", "c++": "cpp",
    "md": "markdown", "text": None, "txt": None,
}

SUPPORTED_CANVAS_LANGUAGES: frozenset[str] = frozenset({
    "python", "c", "cpp", "markdown", "latex", "json", "html", "css",
    "javascript", "jinja2", "typescript", "yaml", "dockerfile", "shell", "r", "sql",
})

_RE_CODE_BLOCK = re.compile(r"```([^\n`]*)\n(.*?)```", re.DOTALL)
_RE_USER_TURN  = re.compile(r"\n(?:User|Usuario):\s*", re.IGNORECASE)

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────
def texto_del_contenido(content) -> str:
    if isinstance(content, str): return content
    if isinstance(content, list):
        partes = [
            str(b.get("text", "")) if isinstance(b, dict) and b.get("type") == "text"
            else str(b)
            for b in content if isinstance(b, (str, dict))
        ]
        return "\n".join(partes)
    return str(content or "")

def construir_prompt(mensaje: str, historial: list, idioma: str) -> str:
    system_prompt = SYSTEM_PROMPTS.get(idioma, SYSTEM_PROMPTS["en"])
    mensajes = [{"role": "system", "content": system_prompt}]
    
    # Soporte robusto de historial (maneja formato de diccionarios y tuplas)
    for turno in historial[-MAX_HISTORIAL:]:
        if isinstance(turno, dict):
            rol = turno.get("role")
            if rol in {"user", "assistant"}:
                mensajes.append({"role": rol, "content": texto_del_contenido(turno.get("content", ""))})
        elif isinstance(turno, (list, tuple)) and len(turno) == 2:
            user_msg, bot_msg = turno
            if user_msg:
                mensajes.append({"role": "user", "content": texto_del_contenido(user_msg)})
            if bot_msg:
                mensajes.append({"role": "assistant", "content": texto_del_contenido(bot_msg)})

    mensajes.append({"role": "user", "content": mensaje})

    if getattr(tokenizer, "chat_template", None):
        return tokenizer.apply_chat_template(mensajes, tokenize=False, add_generation_prompt=True)

    lineas = [f"System: {system_prompt}"]
    for msg in mensajes[1:]:
        nombre = "User" if msg["role"] == "user" else "Assistant"
        lineas.append(f"{nombre}: {msg['content']}")
    lineas.append("Assistant:")
    return "\n\n".join(lineas)

def extraer_canvas(respuesta: str) -> tuple[Optional[str], Optional[str]]:
    bloques = _RE_CODE_BLOCK.findall(respuesta)
    if not bloques: return None, None
    lenguaje_raw, codigo = max(bloques, key=lambda b: len(b[1].strip()))
    codigo = codigo.strip()
    if not codigo: return None, None
    lenguaje = lenguaje_raw.strip().lower().split()[0] if lenguaje_raw.strip() else None
    lenguaje = LANGUAGE_ALIASES.get(lenguaje, lenguaje)
    if lenguaje not in SUPPORTED_CANVAS_LANGUAGES: lenguaje = None
    return codigo, lenguaje

def _calcular_tokens_salida(prompt: str, max_tokens: int, longitud_automatica: bool) -> int:
    if not longitud_automatica: return int(max_tokens)
    tokens_entrada = len(tokenizer(prompt, add_special_tokens=False, truncation=False).input_ids)
    return max(1, CONTEXT_LIMIT - tokens_entrada - 16)

def _construir_kwargs_generacion(tokens_salida: int, temperatura: float, top_p: float, longitud_automatica: bool) -> dict:
    kwargs: dict = {
        "max_new_tokens": tokens_salida, "return_full_text": False,
        "repetition_penalty": DEFAULT_REPETITION_PENALTY, "pad_token_id": tokenizer.eos_token_id,
    }
    if longitud_automatica: kwargs["max_time"] = 120 - GPU_TIMEOUT_BUFFER
    if float(temperatura) > 0:
        kwargs.update(do_sample=True, temperature=float(temperatura), top_p=float(top_p))
    else: kwargs["do_sample"] = False
    return kwargs

# ─────────────────────────────────────────────────────────────────────────────
# Chatbot Principal
# ─────────────────────────────────────────────────────────────────────────────
@spaces.GPU(duration=120)
def chat_bot(mensaje_usuario: str, historial: list, idioma: str, temperatura: float, top_p: float, max_tokens: int, longitud_automatica: bool):
    mensaje_usuario = (mensaje_usuario or "").strip()
    if not mensaje_usuario:
        msg = "Type a question or paste a code snippet to get started." if idioma == "en" else "Escribe una pregunta o pega un fragmento de código para comenzar."
        return msg, ui.skip()

    try:
        prompt = construir_prompt(mensaje_usuario, historial, idioma)
        tokens_salida = _calcular_tokens_salida(prompt, max_tokens, longitud_automatica)
        kwargs = _construir_kwargs_generacion(tokens_salida, temperatura, top_p, longitud_automatica)
        
        log.info("Generando respuesta | idioma=%s temp=%.2f tokens=%d", idioma, temperatura, tokens_salida)
        resultado = pipe(prompt, **kwargs)
        respuesta = resultado[0]["generated_text"].strip()
        respuesta = _RE_USER_TURN.split(respuesta, maxsplit=1)[0].strip()

        if respuesta:
            codigo, lenguaje = extraer_canvas(respuesta)
            canvas_update = ui.update(value=codigo, language=lenguaje) if codigo else ui.skip()
            return respuesta, canvas_update

    except torch.cuda.OutOfMemoryError:
        log.exception("OOM durante la generación")
        msg = "⚠️ Out of GPU memory. Try reducing token length." if idioma == "en" else "⚠️ Memoria de GPU insuficiente. Reduce la longitud."
        return msg, ui.skip()
    except Exception:
        log.exception("Error inesperado en chat_bot")

    msg = "⚠️ I couldn't generate a response. Please try again." if idioma == "en" else "⚠️ No pude generar una respuesta. Inténtalo de nuevo."
    return msg, ui.skip()

# ─────────────────────────────────────────────────────────────────────────────
# Textos UI
# ─────────────────────────────────────────────────────────────────────────────
UI_TEXT: dict[str, dict[str, str]] = {
    "en": {
        "subtitle": "AI coding workspace · Design, debug, and build faster", "online": "Model connected",
        "settings": "Configuration", "settings_sub": "Adjust the generation style",
        "creativity": "Creativity", "creativity_info": "0 = precise · 1.5 = creative", "top_p_info": "Token diversity",
        "length": "Maximum length", "length_info": "Response tokens", "auto_length": "Auto length",
        "auto_length_info": "Stops at EOS, context limit, or the 120s GPU timeout", "tip_title": "Pro Tip",
        "tip": "Include the language, goal, constraints, and full error message for better code.",
        "welcome": "What are we building?", "welcome_sub": "Ask a question, paste code, or describe a bug.",
        "placeholder": "Describe your task or paste code…", "send": "Send", "stop": "Stop",
        "canvas": "Canvas", "canvas_sub": "The main generated code appears here automatically",
    },
    "es": {
        "subtitle": "Espacio de programación con IA · Diseña, depura y construye más rápido", "online": "Modelo conectado",
        "settings": "Configuración", "settings_sub": "Ajusta el estilo de generación",
        "creativity": "Creatividad", "creativity_info": "0 = preciso · 1.5 = creativo", "top_p_info": "Diversidad de tokens",
        "length": "Longitud máxima", "length_info": "Tokens de respuesta", "auto_length": "Longitud automática",
        "auto_length_info": "Se detiene por EOS, límite de contexto o timeout de 120s", "tip_title": "Consejo Pro",
        "tip": "Incluye lenguaje, objetivo, restricciones y el error completo para obtener mejor código.",
        "welcome": "¿Qué vamos a construir?", "welcome_sub": "Pregunta, pega código o describe un bug.",
        "placeholder": "Describe tu tarea o pega código…", "send": "Enviar", "stop": "Detener",
        "canvas": "Canvas", "canvas_sub": "El código principal generado aparece aquí automáticamente",
    },
}

# ─────────────────────────────────────────────────────────────────────────────
# HTML Generators
# ─────────────────────────────────────────────────────────────────────────────
@lru_cache(maxsize=8)
def hero_html(lang: str = "en") -> str:
    t = UI_TEXT.get(lang, UI_TEXT["en"])
    return f"""
    <div id="top-bar">
      <div class="brand">
        <div class="logo">DRK</div>
        <div class="brand-text">
          <h1>DRK Code V1</h1>
          <p>{t['subtitle']}</p>
        </div>
      </div>
      <div class="online"><span class="dot"></span> {t['online']}</div>
    </div>
    """

@lru_cache(maxsize=8)
def panel_html(lang: str = "en") -> str:
    t = UI_TEXT.get(lang, UI_TEXT["en"])
    return f'<div class="section-header">{t["settings"]}</div><div class="section-sub">{t["settings_sub"]}</div>'

@lru_cache(maxsize=8)
def tip_html(lang: str = "en") -> str:
    t = UI_TEXT.get(lang, UI_TEXT["en"])
    return f'<div class="quick-tip"><b>{t["tip_title"]}</b>{t["tip"]}</div>'

@lru_cache(maxsize=8)
def canvas_html(lang: str = "en") -> str:
    t = UI_TEXT.get(lang, UI_TEXT["en"])
    return f'<div class="canvas-header"><div class="canvas-title"><span>✦</span> {t["canvas"]}</div><div class="section-sub">{t["canvas_sub"]}</div></div>'

@lru_cache(maxsize=8)
def placeholder_md(lang: str = "en") -> str:
    """Devuelve Markdown para que Gradio lo renderice correctamente en el placeholder."""
    t = UI_TEXT.get(lang, UI_TEXT["en"])
    return (
        "### ⌁\n\n"
        f"**{t['welcome']}**\n\n"
        f"{t['welcome_sub']}"
    )

def cambiar_idioma(lang: str):
    t = UI_TEXT.get(lang, UI_TEXT["en"])
    return (
        hero_html(lang), panel_html(lang),
        ui.update(label=t["creativity"], info=t["creativity_info"]),
        ui.update(info=t["top_p_info"]),
        ui.update(label=t["length"], info=t["length_info"]),
        ui.update(label=t["auto_length"], info=t["auto_length_info"]),
        tip_html(lang), canvas_html(lang),
        ui.update(placeholder=placeholder_md(lang)),
        ui.update(placeholder=t["placeholder"], submit_btn=t["send"], stop_btn=t["stop"]),
    )

# ─────────────────────────────────────────────────────────────────────────────
# CSS & Theme
# ─────────────────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');

:root {
  --bg-base: #09090b;
  --bg-panel: #0f0f12;
  --bg-input: #18181b;
  --border-subtle: #27272a;
  --border-default: #3f3f46;
  --text-primary: #fafafa;
  --text-secondary: #a1a1aa;
  --text-tertiary: #71717a;
  --accent: #6366f1;
  --accent-hover: #4f46e5;
  --accent-soft: rgba(99, 102, 241, 0.1);
}

html, body, .gradio-container {
  background-color: var(--bg-base) !important;
  font-family: 'Inter', system-ui, -apple-system, sans-serif !important;
  color: var(--text-primary) !important;
}

.gradio-container {
  max-width: 1600px !important;
  padding: 0 24px 24px !important;
}

/* Dot Matrix Background */
.gradio-container::before {
  content: "";
  position: fixed;
  inset: 0;
  pointer-events: none;
  background-image: radial-gradient(rgba(255, 255, 255, 0.025) 1px, transparent 1px);
  background-size: 24px 24px;
  z-index: 0;
}
.gradio-container > * { position: relative; z-index: 1; }
footer { display: none !important; }

/* Top Bar */
#top-bar {
  display: flex; align-items: center; justify-content: space-between;
  padding: 20px 0; border-bottom: 1px solid var(--border-subtle);
  margin-bottom: 24px;
}
.brand { display: flex; align-items: center; gap: 16px; }
.logo {
  width: 40px; height: 40px; border-radius: 8px;
  display: grid; place-items: center;
  font: 700 15px/1 'Inter', sans-serif; color: white;
  background: var(--bg-input); border: 1px solid var(--border-default);
  box-shadow: 0 1px 2px rgba(0,0,0,0.3);
}
.brand-text h1 { margin: 0; font-size: 18px; font-weight: 600; letter-spacing: -0.01em; color: var(--text-primary); }
.brand-text p { margin: 2px 0 0; font-size: 13px; color: var(--text-tertiary); }
.online {
  display: flex; align-items: center; gap: 8px;
  padding: 6px 12px; border-radius: 6px;
  background: var(--bg-panel); border: 1px solid var(--border-subtle);
  color: var(--text-secondary); font-size: 12px; font-weight: 500;
}
.dot { width: 8px; height: 8px; border-radius: 50%; background: #10b981; box-shadow: 0 0 8px rgba(16, 185, 129, 0.6); }

/* Section Headers */
.section-header { font-size: 11px; font-weight: 600; text-transform: uppercase; letter-spacing: 0.08em; color: var(--text-secondary); margin-bottom: 4px; }
.section-sub { font-size: 13px; color: var(--text-tertiary); margin-bottom: 20px; }

/* Panels */
#settings-panel, #main-chat, #canvas-panel {
  background: var(--bg-panel) !important;
  border: 1px solid var(--border-subtle) !important;
  border-radius: 8px !important;
  box-shadow: 0 1px 3px rgba(0,0,0,0.2) !important;
}
#settings-panel { padding: 24px !important; }
#canvas-panel { padding: 20px !important; display: flex; flex-direction: column; }
.canvas-header { display: flex; flex-direction: column; margin-bottom: 12px; }
.canvas-title { font-size: 11px; font-weight: 600; text-transform: uppercase; letter-spacing: 0.08em; color: var(--text-secondary); display: flex; align-items: center; gap: 6px; }

/* Chat styles */
#main-chat { padding: 8px !important; overflow: hidden; }
#chatbot { background: transparent !important; border: 0 !important; height: 100% !important; }
#chatbot .message {
  border-radius: 8px !important;
  border: 1px solid var(--border-subtle) !important;
  padding: 16px !important;
  margin-bottom: 16px !important;
  font-size: 14px !important;
  line-height: 1.6 !important;
}
#chatbot .message.user {
  background: var(--accent-soft) !important;
  border-color: rgba(99, 102, 241, 0.3) !important;
}
#chatbot .message.bot {
  background: var(--bg-base) !important;
}
#chatbot pre {
  background: var(--bg-base) !important;
  border: 1px solid var(--border-subtle) !important;
  border-radius: 6px !important;
  font-family: 'JetBrains Mono', monospace !important;
  font-size: 13px !important;
}

/* Code Editor (Canvas) */
#canvas-code {
  min-height: 540px !important;
  flex-grow: 1;
  background: var(--bg-base) !important;
  border-radius: 6px !important;
  border: 1px solid var(--border-subtle) !important;
  font-family: 'JetBrains Mono', monospace !important;
}

/* Form Elements */
input[type="text"], textarea, select {
  background-color: var(--bg-input) !important;
  border: 1px solid var(--border-subtle) !important;
  border-radius: 6px !important;
  color: var(--text-primary) !important;
  font-size: 14px !important;
  transition: border-color 0.2s ease, box-shadow 0.2s ease !important;
}
input[type="text"]:focus, textarea:focus {
  border-color: var(--accent) !important;
  box-shadow: 0 0 0 3px var(--accent-soft) !important;
  outline: none !important;
}

/* Gradio Form Overrides */
.gradio-container .form { background: transparent !important; border: 0 !important; }
label { color: var(--text-secondary) !important; font-size: 13px !important; font-weight: 500 !important; }
input[type="range"] { accent-color: var(--accent); }

/* Buttons */
button {
  border-radius: 6px !important;
  transition: all 0.2s ease !important;
  font-weight: 500 !important;
  border: 1px solid var(--border-subtle) !important;
  background: var(--bg-input) !important;
  color: var(--text-primary) !important;
  font-size: 14px !important;
}
button:hover { background: var(--border-subtle) !important; }
button.primary {
  background: var(--accent) !important;
  border-color: var(--accent) !important;
  color: white !important;
}
button.primary:hover {
  background: var(--accent-hover) !important;
  border-color: var(--accent-hover) !important;
}

/* Quick Tip */
.quick-tip {
  margin-top: 24px; padding: 12px 16px;
  border-radius: 6px; font-size: 13px; line-height: 1.5;
  color: var(--text-secondary);
  border: 1px solid var(--border-subtle);
  background: var(--bg-input);
}
.quick-tip b { color: var(--text-primary); display: block; margin-bottom: 4px; font-size: 12px; text-transform: uppercase; letter-spacing: 0.05em; }

/* Responsive */
@media (max-width: 1024px) {
  .gradio-container { padding: 0 16px 16px !important; }
  #top-bar { flex-direction: column; align-items: flex-start; gap: 16px; }
}
"""

THEME = ui.themes.Base(
    primary_hue="indigo",
    secondary_hue="slate",
    neutral_hue="zinc",
).set(
    body_background_fill="#09090b",
    body_text_color="#fafafa",
    block_background_fill="#0f0f12",
    block_border_color="#27272a",
    input_background_fill="#18181b",
    input_border_color="#27272a",
    button_primary_background_fill="#6366f1",
    button_primary_background_fill_hover="#4f46e5",
    button_primary_text_color="#ffffff",
)

# ─────────────────────────────────────────────────────────────────────────────
# Interfaz Gradio
# ─────────────────────────────────────────────────────────────────────────────
with ui.Blocks(fill_height=True, title="DRK Code V1") as demo:

    hero = ui.HTML(hero_html("en"))

    canvas = ui.Code(
        value="", language=None, lines=28, max_lines=40,
        show_label=False, interactive=True, wrap_lines=True,
        show_line_numbers=True, buttons=["copy", "download"],
        elem_id="canvas-code", render=False,
    )

    with ui.Row(equal_height=True):

        # Panel de Configuración
        with ui.Column(scale=1, min_width=260, elem_id="settings-panel"):
            idioma = ui.Dropdown(
                choices=[("English", "en"), ("Español", "es")], value="en",
                label="Language / Idioma", interactive=True,
            )
            panel_header = ui.HTML(panel_html("en"))
            temperatura = ui.Slider(0, 1.5, value=0.35, step=0.05, label="Creativity", info="0 = precise · 1.5 = creative")
            top_p = ui.Slider(0.1, 1, value=0.9, step=0.05, label="Top P", info="Token diversity")
            max_tokens = ui.Slider(128, 4096, value=1024, step=128, label="Maximum length", info="Response tokens")
            longitud_automatica = ui.Checkbox(value=False, label="Auto length", info="Stops at EOS, context limit, or 120s timeout")
            tip = ui.HTML(tip_html("en"))

        # Chat Principal
        with ui.Column(scale=4, min_width=400, elem_id="main-chat"):
            chatbot = ui.Chatbot(
                elem_id="chatbot", height=620, show_label=False,
                buttons=["copy", "copy_all"], reasoning_tags=[("<think>", "</think>")],
                placeholder=placeholder_md("en"), render=False,
            )
            caja = ui.Textbox(
                placeholder="Describe your task or paste code…", lines=2, max_lines=10,
                show_label=False, container=False, submit_btn="Send", stop_btn="Stop", render=False,
            )
            ui.ChatInterface(
                fn=chat_bot, chatbot=chatbot, textbox=caja,
                additional_inputs=[idioma, temperatura, top_p, max_tokens, longitud_automatica],
                additional_outputs=[canvas], editable=True, fill_height=True,
            )

        # Panel Canvas
        with ui.Column(scale=2, min_width=360, elem_id="canvas-panel"):
            canvas_header = ui.HTML(canvas_html("en"))
            canvas.render()

    # Eventos
    longitud_automatica.change(
        fn=lambda activada: ui.update(interactive=not activada),
        inputs=longitud_automatica, outputs=max_tokens, queue=False,
    )

    idioma.change(
        fn=cambiar_idioma, inputs=idioma,
        outputs=[hero, panel_header, temperatura, top_p, max_tokens, longitud_automatica, tip, canvas_header, chatbot, caja],
        queue=False,
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=20).launch(theme=THEME, css=CSS)