import uvicorn
import json
import asyncio
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse
from llama_cpp import Llama

app = FastAPI(
    title="Dimercia AI",
    version="0.1.4",
    description="API OpenAI compatible souveraine - Anti-Timeout pour Cline"
)

# --- Initialisation globale ---
MODEL_PATH = "/app/models/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf"

print("Chargement de Dimercia AI v0.1.4 en mémoire RAM...")
llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=16384,     # 16k conserve un parfait équilibre sur CPU
    n_threads=4,     # Exploite à fond le calcul parallèle
    verbose=False
)
print("Dimercia AI v0.1.4 est prêt.")

@app.get("/")
def home():
    return {
        "name": "Dimercia AI",
        "version": "0.1.4",
        "status": "running"
    }

@app.get("/v1/models")
def models():
    return {
        "object": "list",
        "data": [
            {"id": "dimercia-coder", "object": "model", "owned_by": "dimercia"}
        ]
    }

@app.post("/v1/chat/completions")
async def chat(request: Request):
    try:
        body = await request.json()
    except Exception:
        return JSONResponse(status_code=400, content={"detail": "JSON invalide"})

    # --- Nettoyage adaptatif du Payload ---
    raw_messages = body.get("messages", [])
    cleaned_messages = []
    
    for msg in raw_messages:
        role = msg.get("role", "user")
        raw_content = msg.get("content", "")
        
        if isinstance(raw_content, list):
            text_pieces = [piece.get("text", "") if isinstance(piece, dict) else str(piece) for piece in raw_content]
            content = "".join(text_pieces)
        else:
            content = str(raw_content)
            
        cleaned_messages.append({"role": role, "content": content})

    temperature = body.get("temperature", 0.2)
    max_tokens = body.get("max_tokens", 512)
    stream = body.get("stream", False)

    temp_val = float(temperature) if temperature is not None else 0.2
    tokens_val = int(max_tokens) if max_tokens is not None else 512

    # --- Gestion du mode STREAMING (Avec système anti-timeout) ---
    if stream:
        async def chunk_generator():
            # Étape 1 : Lancer la création de l'itérateur dans un thread séparé (non-bloquant)
            task = asyncio.create_task(asyncio.to_thread(
                llm.create_chat_completion,
                messages=cleaned_messages,
                temperature=temp_val,
                max_tokens=tokens_val,
                stream=True
            ))

            # Étape 2 : Tant que llama.cpp calcule le prefill, on envoie des pings invisibles à Cline
            while not task.done():
                # Envoi d'un chunk de commentaire SSE pour garder la connexion ouverte
                yield ": heartbeat\n\n"
                await asyncio.sleep(1.0) # Attendre 1 seconde avant le prochain ping

            try:
                iterator = task.result()
            except Exception as e:
                yield f"data: {json.dumps({'error': str(e)})}\n\n"
                yield "data: [DONE]\n\n"
                return

            # Étape 3 : Consommer les tokens normalement dès qu'ils sont prêts
            def get_next_chunk(it):
                try:
                    return next(it)
                except StopIteration:
                    return None
                except Exception as ex:
                    return ex

            while True:
                chunk = await asyncio.to_thread(get_next_chunk, iterator)
                if chunk is None:
                    break
                if isinstance(chunk, Exception):
                    yield f"data: {json.dumps({'error': str(chunk)})}\n\n"
                    break
                
                if "model" in chunk:
                    chunk["model"] = "dimercia-coder"
                yield f"data: {json.dumps(chunk)}\n\n"
                
            yield "data: [DONE]\n\n"

        return StreamingResponse(chunk_generator(), media_type="text/event-stream")

    # --- Gestion du mode STANDARD ---
    else:
        try:
            response = await asyncio.to_thread(
                llm.create_chat_completion,
                messages=cleaned_messages,
                temperature=temp_val,
                max_tokens=tokens_val,
                stream=False
            )
            if isinstance(response, dict):
                response["model"] = "dimercia-coder"
            return response
        except Exception as e:
            return JSONResponse(status_code=500, content={"detail": str(e)})

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)