from pydantic import BaseModel
import os
import time
import asyncio
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoTokenizer, AutoModelForCausalLM
from fastapi import FastAPI
import torch
import os
from dotenv import load_dotenv
load_dotenv()  


# HF_TOKEN = os.getenv("HF_TOKEN")
MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")


# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="cpu",
    torch_dtype=torch.float32,
    # use_auth_token=token
)
 
model.eval()

if hasattr(model, "config"):
    model.config.use_cache = True

app = FastAPI()

# CORS: env orqali
origins_env = os.getenv("ALLOWED_ORIGINS", "http://localhost:3000")
allowed_origins = [o.strip() for o in origins_env.split(",") if o.strip()]

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

# app.add_middleware(
#     CORSMiddleware,
#     allow_origins=["http://localhost:3000"],  # frontend
#     allow_credentials=True,
#     allow_methods=["*"],
#     allow_headers=["*"],
# )

# CPU model uchun parallel requestlarni navbatga qo'yish
generation_lock = asyncio.Lock()


# Request model
class GenerateRequest(BaseModel):
    prompt: str

# Response model (ixtiyoriy)
class GenerateResponse(BaseModel):
    response: str


@app.post("/generate", response_model=GenerateResponse)
async def generate(req: GenerateRequest):
    prompt = f"You are an helpful assistant. Provide a concise answer to the question below.\nQuestion: {req.prompt}\nAnswer:"
    start = time.time()
    async with generation_lock:
        with torch.inference_mode():
            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
            outputs = model.generate(**inputs, max_new_tokens=100, temperature=0.7, top_p=0.9, top_k=50, do_sample=True)
    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    answer = decoded.split("Answer:")[-1].strip()
    print(f"/generate took {time.time() - start:.2f}s")
    return {"response": answer}