import os
import uvicorn
import asyncio
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

app = FastAPI(title="J.A.R.V.I.S. Core API")

print("[System] Authenticating and downloading private model weights...")
try:
    model_path = hf_hub_download(
        repo_id="omjhamtani77/Jarvis-3B-Core",
        filename="jarvis.gguf"
    )
    print(f"[System] Secure model download path locked: {model_path}")
except Exception as e:
    print(f"[Fatal System Error] Failed to secure private weights: {e}")
    raise e

print("[System] Initializing model parameters into core engine loops...")
# Re-allocating 2 threads now that the main event loop is isolated and safe
llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2
)
print("[System] J.A.R.V.I.S. Brain Arrays: Fully Active.")

@app.get("/")
def health_check():
    return {"status": "online", "engine": "J.A.R.V.I.S. Core"}

# DOUBLE-STACKED ROUTES: This makes the API immune to n8n's hidden URL formatting
@app.post("/chat/completions")
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    try:
        body = await request.json()
        messages = body.get("messages", [])
        temperature = body.get("temperature", 0.1)
        max_tokens = body.get("max_tokens", 150)
        
        # CRITICAL FIX: Run the heavy CPU-bound task in a separate thread pool.
        # This keeps the FastAPI network loop open so it doesn't freeze or timeout.
        response = await asyncio.to_thread(
            llm.create_chat_completion,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens
        )
        return response
        
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"error": f"Internal Model Execution Failure: {str(e)}"}
        )

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