import os
import re
from contextlib import asynccontextmanager
from pathlib import Path
from typing import List, Optional, Dict, Any, Tuple

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field

from llama_cpp import Llama
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, BaseMessage
from langgraph.graph import StateGraph, MessagesState, START, END

# Importador para asegurar descarga automática si falta el modelo
from download_model import ensure_model


# ==========================================
# 1. Configuración del Modelo
# ==========================================
BASE_DIR = Path(__file__).resolve().parent
MODEL_DIR = BASE_DIR / "model"
CLIENTE_DIR = BASE_DIR / "cliente"

def get_model_path() -> Path:
    """Obtiene la ruta del modelo .gguf; si no existe, lo descarga automáticamente."""
    gguf_files = list(MODEL_DIR.glob("*.gguf"))
    if gguf_files:
        return gguf_files[0]
    
    print("[INFO] No se encontró modelo local. Iniciando descarga automática...")
    return ensure_model()

MODEL_FILE = get_model_path()

# Configuración del LLM optimizada (soporta tanto Render como HF Spaces de 16 GB)
MAX_CONTEXT_TOKENS = int(os.environ.get("MAX_CONTEXT_TOKENS", 2048))
SAFE_PROMPT_LIMIT = int(os.environ.get("SAFE_PROMPT_LIMIT", 1600))
MODEL_THREADS = int(os.environ.get("MODEL_THREADS", 2))

llm: Optional[Llama] = None

def get_llm() -> Llama:
    """Retorna la instancia del modelo en memoria, cargándola si aún no está lista."""
    global llm
    if llm is None:
        print(f"[INFO] Cargando modelo en memoria: {MODEL_FILE.name} (Contexto: {MAX_CONTEXT_TOKENS}, Threads: {MODEL_THREADS})...")
        llm = Llama(
            model_path=str(MODEL_FILE),
            n_ctx=MAX_CONTEXT_TOKENS,
            n_threads=MODEL_THREADS,
            n_batch=256,
            verbose=False
        )
        print(f"[INFO] Modelo {MODEL_FILE.name} cargado en memoria exitosamente.")
    return llm

@asynccontextmanager
async def lifespan(app: FastAPI):
    get_llm()
    yield
    print("[INFO] Apagando servicio...")


# ==========================================
# 2. Funciones de Optimización de Contexto
# ==========================================
DEFAULT_SYSTEM_PROMPT = (
    "Eres una mente sintética superinteligente, hiperprecisa y de pocas palabras. "
    "Respondes con exactitud milimétrica en un máximo de 1 a 2 oraciones en español. "
    "Cero relleno, sin introducciones vacías, directo a la verdad esencial."
)

def estimate_tokens(text: str) -> int:
    """Estima la cantidad de tokens de un texto."""
    model = get_llm()
    return len(model.tokenize(text.encode("utf-8")))

def prune_history_if_needed(messages: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], bool]:
    """
    Poda el historial para evitar saturación de memoria y degradación del modelo.
    Mantiene siempre el system_prompt y recorta las conversaciones más viejas.
    """
    total_tokens = sum(estimate_tokens(m.get("content", "")) for m in messages)
    pruned = False

    while total_tokens > SAFE_PROMPT_LIMIT and len(messages) > 3:
        messages.pop(1)
        pruned = True
        total_tokens = sum(estimate_tokens(m.get("content", "")) for m in messages)

    return messages, pruned

def clean_output(text: str) -> str:
    """Limpia posibles artefactos residuales de plantillas."""
    text = re.sub(r"<\|im_\w+\|>", "", text)
    return text.strip()


# ==========================================
# 3. Workflow con LangGraph
# ==========================================
def llm_node(state: MessagesState) -> Dict[str, Any]:
    """Nodo LangGraph que ejecuta inferencia precisa y concisa en español."""
    model = get_llm()

    raw_messages = state["messages"]
    formatted = []

    for msg in raw_messages:
        if isinstance(msg, dict):
            formatted.append({
                "role": msg.get("role", "user"),
                "content": str(msg.get("content", ""))
            })
        elif isinstance(msg, AIMessage):
            formatted.append({"role": "assistant", "content": str(msg.content)})
        elif isinstance(msg, SystemMessage):
            formatted.append({"role": "system", "content": str(msg.content)})
        elif isinstance(msg, HumanMessage):
            formatted.append({"role": "user", "content": str(msg.content)})
        elif isinstance(msg, BaseMessage):
            formatted.append({"role": getattr(msg, "type", "user"), "content": str(msg.content)})

    # Poda preventiva de contexto para cuidar la memoria
    optimized_messages, was_pruned = prune_history_if_needed(formatted)

    completion = model.create_chat_completion(
        messages=optimized_messages,
        max_tokens=80,
        temperature=0.3,
        repeat_penalty=1.18,
        stop=["<|im_end|>", "<|im_start|>", "\nUser:", "User:"]
    )

    raw_reply = completion["choices"][0]["message"]["content"]
    bot_reply = clean_output(raw_reply)
    if not bot_reply:
        bot_reply = "Respuesta procesada con éxito."

    usage = completion.get("usage", {})

    ai_msg = AIMessage(
        content=bot_reply,
        additional_kwargs={
            "usage": usage,
            "was_pruned": was_pruned,
            "optimized_messages": optimized_messages
        }
    )
    return {"messages": [ai_msg]}


graph_builder = StateGraph(MessagesState)
graph_builder.add_node("agent", llm_node)
graph_builder.add_edge(START, "agent")
graph_builder.add_edge("agent", END)
graph_app = graph_builder.compile()


# ==========================================
# 4. Esquemas Pydantic y API FastAPI
# ==========================================
app = FastAPI(
    title="Qwen 2.5 Hyper-Concise Chat API",
    description="API de chat superinteligente y concisa con Qwen2.5-0.5B-Instruct y LangGraph en un solo proceso. Optimizada para Render Free (512 MB RAM).",
    version="3.0.0",
    lifespan=lifespan
)


class ChatRequest(BaseModel):
    message: str = Field(..., description="Mensaje del usuario", json_schema_extra={"example": "¿Qué es la entropía?"})
    history: Optional[List[Dict[str, str]]] = Field(
        default=None,
        description="Historial previo [{'role': 'user'|'assistant', 'content': '...'}]"
    )
    system_prompt: Optional[str] = Field(
        default=DEFAULT_SYSTEM_PROMPT,
        description="Prompt de sistema para guiar la precisión y concisión"
    )
    temperature: Optional[float] = Field(default=0.3, ge=0.0, le=1.5, description="Temperatura baja para máxima precisión")
    max_tokens: Optional[int] = Field(default=80, ge=10, le=256, description="Límite para respuestas cortas")


class ContextTelemetry(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    max_context: int
    usage_percent: float
    health_status: str  # "optimo", "alerta", "poda_aplicada"


class ChatResponse(BaseModel):
    reply: str
    model: str
    updated_history: List[Dict[str, str]]
    context_telemetry: ContextTelemetry


# ==========================================
# 5. Endpoints de la Aplicación
# ==========================================
@app.get("/", include_in_schema=False)
@app.head("/", include_in_schema=False)
@app.get("/cliente", include_in_schema=False)
async def serve_chat_ui():
    """Sirve la interfaz web del chat cliente/index.html."""
    html_file = CLIENTE_DIR / "index.html"
    if not html_file.exists():
        raise HTTPException(status_code=404, detail="No se encontró cliente/index.html")
    return FileResponse(html_file)


@app.get(
    "/healthz",
    tags=["Diagnóstico"],
    summary="Verificar el estado del servicio y la memoria"
)
async def health_check():
    return {
        "status": "healthy",
        "model_loaded": llm is not None,
        "model_file": MODEL_FILE.name,
        "max_context": MAX_CONTEXT_TOKENS,
        "persona": "Superhumano Hiperpreciso & Conciso (Español)",
        "platform": "Render Free Ready"
    }


@app.post(
    "/chat",
    response_model=ChatResponse,
    tags=["Chat & Inferencia"],
    summary="Enviar mensaje al chat",
    description="Procesa la conversación a través de LangGraph con Qwen2.5, monitorea el contexto y aplica poda si es necesario."
)
async def chat_endpoint(request: ChatRequest):
    if llm is None:
        raise HTTPException(status_code=503, detail="El modelo aún se está cargando en memoria.")

    # Reconstrucción de mensajes con historial
    messages = [SystemMessage(content=request.system_prompt or DEFAULT_SYSTEM_PROMPT)]
    if request.history:
        for item in request.history:
            role = item.get("role", "user")
            content = item.get("content", "")
            if role == "assistant":
                messages.append(AIMessage(content=content))
            elif role == "user":
                messages.append(HumanMessage(content=content))

    # Agregar el mensaje actual del usuario
    messages.append(HumanMessage(content=request.message))

    # Invocar el grafo de LangGraph
    output = graph_app.invoke({"messages": messages})
    last_msg = output["messages"][-1]

    reply_content = last_msg.content
    kwargs = getattr(last_msg, "additional_kwargs", {})
    usage = kwargs.get("usage", {})
    was_pruned = kwargs.get("was_pruned", False)

    prompt_toks = usage.get("prompt_tokens", estimate_tokens(request.message))
    comp_toks = usage.get("completion_tokens", estimate_tokens(reply_content))
    tot_toks = prompt_toks + comp_toks
    usage_pct = (tot_toks / MAX_CONTEXT_TOKENS) * 100

    # Determinar estado de salud del contexto
    if was_pruned:
        health_status = "poda_aplicada"
    elif usage_pct >= 65.0:
        health_status = "alerta"
    else:
        health_status = "optimo"

    # Actualizar historial para el cliente
    current_history = request.history or []
    current_history.append({"role": "user", "content": request.message})
    current_history.append({"role": "assistant", "content": reply_content})

    # Si se aplicó poda en el servidor, sincronizar historial
    if was_pruned and "optimized_messages" in kwargs:
        current_history = [
            m for m in kwargs["optimized_messages"] if m.get("role") in ("user", "assistant")
        ]

    telemetry = ContextTelemetry(
        prompt_tokens=prompt_toks,
        completion_tokens=comp_toks,
        total_tokens=tot_toks,
        max_context=MAX_CONTEXT_TOKENS,
        usage_percent=round(usage_pct, 1),
        health_status=health_status
    )

    return ChatResponse(
        reply=reply_content,
        model=MODEL_FILE.name,
        updated_history=current_history,
        context_telemetry=telemetry
    )


# ==========================================
# 6. Arranque Local y en Render
# ==========================================
if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    print(f"\n=======================================================")
    print(f"🚀 Qwen2.5 Chat API Activa en: http://localhost:{port}")
    print(f"💬 Interfaz de Chat Web:      http://localhost:{port}/cliente")
    print(f"📖 Swagger Docs:              http://localhost:{port}/docs")
    print(f"=======================================================\n")
    uvicorn.run("main:app", host="0.0.0.0", port=port, reload=False)