import json
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

app = FastAPI()

# Guarantee seamless web streaming connections from your Vercel frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

print("Downloading highly optimized VibeThinker-3B Q4_K_M GGUF model...")
model_path = hf_hub_download(
    repo_id="prithivMLmods/VibeThinker-3B-GGUF", 
    filename="VibeThinker-3B.Q4_K_M.gguf"
)

print("Initializing memory-optimized llama.cpp execution runtime...")
llm = Llama(
    model_path=model_path, 
    n_ctx=6144,   # Expanded context window from 4096 to 6144 for longer historical tracks
    n_batch=32,   # Kept at 32 to guarantee flat peak memory profiles
    n_threads=2
)

@app.get("/")
def read_root():
    return {"status": "online", "engine": "llama.cpp Memory-Hardened Core"}

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    body = await request.json()
    messages = body.get("messages", [])
    response = llm.create_chat_completion(
        messages=messages,
        temperature=0.3,
        top_p=0.95,
        stream=True,
        max_tokens=3072  # Doubled token headroom from 1536 to 3072
    )
    
    def stream_generator():
        try:
            for chunk in response:
                delta = chunk.get("choices", [{}])[0].get("delta", {})
                if "content" in delta:
                    yield f"data: {json.dumps(chunk)}\n\n"
            yield "data: [DONE]\n\n"
        except Exception as e:
            # Catch silent disconnects or timeouts cleanly without crashing the Uvicorn thread
            print(f"Streaming trace intercepted safely: {str(e)}")
            return

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

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