import os
import time
from fastapi import FastAPI, HTTPException, Header, Depends, Request
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

app = FastAPI()
SECRET_API_KEY = os.getenv("API_PASSWORD")

model_id = "Qwen/Qwen2.5-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)

@app.get("/health")
def health_check():
    return {"status": "active"}

@app.get("/")
def home_endpoint():
    return {"status": "Server is running perfectly!"}

# n8n OpenAI মডেলে ডিফল্টভাবে "Bearer " যুক্ত করে পাসওয়ার্ড পাঠায়
def verify_api_key(authorization: str = Header(None)):
    if not SECRET_API_KEY:
        raise HTTPException(status_code=500, detail="API Key not configured in Space Secrets")
    
    expected_auth = f"Bearer {SECRET_API_KEY}"
    if authorization != expected_auth:
        raise HTTPException(status_code=401, detail="Unauthorized: Invalid API Key")

# n8n-এর AI Agent ঠিক এই লিংকেই রিকোয়েস্ট পাঠাবে
@app.post("/chat/completions", dependencies=[Depends(verify_api_key)])
@app.post("/v1/chat/completions", dependencies=[Depends(verify_api_key)])
async def chat_completions(request: Request):
    data = await request.json()
    
    # n8n নিজে থেকেই সিস্টেম প্রম্পট, পোস্টগ্রেসের মেমোরি এবং নতুন মেসেজ এই 'messages'-এর ভেতর পাঠিয়ে দেবে
    messages = data.get("messages", [])
    
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt")
    
    outputs = model.generate(
        **inputs, 
        max_new_tokens=150,
        do_sample=True,
        temperature=0.6,
        top_p=0.85,
        num_beams=1,
        pad_token_id=tokenizer.eos_token_id
    )
    response_text = tokenizer.decode(outputs[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True)
    
    # n8n-কে ঠিক OpenAI-এর ফরম্যাটেই উত্তর ফেরত দিতে হবে
    return {
        "id": f"chatcmpl-{int(time.time())}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model_id,
        "choices": [{
            "index": 0,
            "message": {
                "role": "assistant",
                "content": response_text,
            },
            "finish_reason": "stop"
        }]
    }