import os
import json
import gc
import asyncio
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from duckduckgo_search import DDGS
from llama_cpp import Llama

app = FastAPI(title="Optimized LLM Server")

AVAILABLE_MODELS = {
    "gemma-4-e2b-it": {
        "repo_id": "google/gemma-4-E2B-it-qat-q4_0-gguf",
        "filename": "*q4_0-it.gguf" 
    },
    "jan-3.5-4b": {
        "repo_id": "janhq/Jan-v3.5-4B-gguf",
        "filename": "*Q4_K_M.gguf"
    }
}
DEFAULT_MODEL = "gemma-4-e2b-it"

active_model_id = None
llm = None
HF_TOKEN = os.environ.get("HF_TOKEN")

# SEMÁFORO: Limita a 1 ejecución simultánea del modelo para evitar que la RAM/CPU exploten.
# Las demás peticiones esperarán su turno ordenadamente sin bloquear el servidor web.
model_semaphore = asyncio.Semaphore(1)

def search_web_sync(query):
    """Función bloqueante original de búsqueda"""
    try:
        results = DDGS().text(query, max_results=3)
        if not results:
            return ""
        context = "CONTEXTO DE INTERNET ACTUALIZADO:\n"
        for r in results:
            context += f"- {r.get('title')}: {r.get('body')}\n"
        return context
    except Exception:
        return ""

async def search_web_async(query):
    """Ejecuta la búsqueda en un hilo separado para no bloquear el servidor"""
    return await asyncio.to_thread(search_web_sync, query)

def load_model_into_memory(model_id):
    global active_model_id, llm
    
    if model_id not in AVAILABLE_MODELS:
        model_id = DEFAULT_MODEL
        
    if active_model_id == model_id:
        return
        
    if llm is not None:
        del llm
        gc.collect() # Liberar RAM explícitamente
        
    model_info = AVAILABLE_MODELS[model_id]
    llm = Llama.from_pretrained(
        repo_id=model_info["repo_id"],
        filename=model_info["filename"],
        n_ctx=4096,
        n_threads=2,
        verbose=False
    )
    active_model_id = model_id

@app.get("/")
async def health_check():
    return {"status": "online", "role": "mirror_worker", "backend": "FastAPI + llama.cpp"}

@app.get("/v1/models")
async def list_models():
    data = [{"id": m, "object": "model"} for m in AVAILABLE_MODELS.keys()]
    return {"object": "list", "data": data}

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    data = await request.json()
    if not data or 'messages' not in data:
        raise HTTPException(status_code=400, detail="Petición inválida")

    messages = data.get('messages', [])
    requested_model = data.get('model', DEFAULT_MODEL)
    temperature = data.get('temperature', 0.6)
    max_tokens = min(data.get('max_tokens', 1024), 2048)
    stream = data.get('stream', False)
    use_web_search = data.get('web_search', False)

    # 1. Búsqueda web asíncrona (si se solicita)
    if use_web_search and messages and messages[-1]['role'] == 'user':
        user_query = messages[-1]['content']
        web_context = await search_web_async(user_query)
        if web_context:
            messages[-1]['content'] = f"Responde usando esta info de internet si es útil:\n{web_context}\n\nPregunta: {user_query}"

    # 2. Cola de inferencia protegida por el semáforo
    async def generate_response():
        async with model_semaphore:
            # Cargamos el modelo (si es necesario) dentro del bloqueo de seguridad
            await asyncio.to_thread(load_model_into_memory, requested_model)
            
            # Ejecutamos la inferencia en un hilo separado
            response = await asyncio.to_thread(
                llm.create_chat_completion,
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                stream=stream
            )
            return response

    try:
        response_data = await generate_response()

        if stream:
            async def event_generator():
                # Iteramos sobre el generador síncrono en un hilo para no bloquear
                def get_next_chunk():
                    try:
                        return next(response_data)
                    except StopIteration:
                        return None
                
                while True:
                    chunk = await asyncio.to_thread(get_next_chunk)
                    if chunk is None:
                        break
                    yield f"data: {json.dumps(chunk)}\n\n"
                yield "data: [DONE]\n\n"

            return StreamingResponse(event_generator(), media_type="text/event-stream")
        else:
            return JSONResponse(content=response_data)
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))