import os
import time
from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

app = FastAPI(title="Qwen 2.5 Coder Public Secure Server")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 🔐 YOUR PREFERRED CUSTOM API KEY
CUSTOM_API_KEY = "varsshancoder"

security = HTTPBearer()

def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
    if credentials.credentials != CUSTOM_API_KEY:
        raise HTTPException(
            status_code=401,
            detail="Invalid or missing custom API Key."
        )
    return credentials.credentials

MODEL_REPO = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF"
MODEL_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf"

print("Checking and downloading Qwen 2.5 Coder 7B GGUF model...")
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)

print(f"Loading model into memory from: {model_path}")
llm = Llama(
    model_path=model_path,
    n_ctx=4096,      
    n_threads=2,     
    n_batch=512,
    verbose=False
)
print("Model loaded successfully and secured!")

@app.post("/v1/chat/completions", dependencies=[Depends(verify_api_key)])
async def chat_completions(request: Request):
    body = await request.json()
    messages = body.get("messages", [])
    temperature = body.get("temperature", 0.2)
    max_tokens = body.get("max_tokens", 1024)

    formatted_prompt = ""
    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        formatted_prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
    formatted_prompt += "<|im_start|>assistant\n"

    response = llm(
        formatted_prompt,
        max_tokens=max_tokens,
        stop=["<|im_end|>", "<|im_start|>"],
        temperature=temperature
    )

    text_output = response["choices"]["text"]
    
    return JSONResponse(content={
        "id": f"chatcmpl-{int(time.time())}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": "qwen2.5-coder-7b",
        "choices": [{
            "index": 0,
            "message": {
                "role": "assistant",
                "content": text_output
            },
            "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens": len(formatted_prompt) // 4, 
            "completion_tokens": len(text_output) // 4,
            "total_tokens": (len(formatted_prompt) + len(text_output)) // 4
        }
    })

@app.get("/")
def index():
    return {
        "status": "healthy", 
        "security": "enabled",
        "model": "qwen2.5-coder-7b-instruct-q4_k_m", 
        "endpoint": "/v1/chat/completions"
    }