from fastapi import FastAPI, Depends, HTTPException
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
import os

app = FastAPI()

# Root kontrol (HF health check için)
@app.get("/")
async def root():
    return {"message": "Joke API is running!"}

security = HTTPBearer()
AUTH_TOKEN = os.getenv("AUTH_TOKEN")

# Model yükle (joke modeli)
base_model = AutoModelForCausalLM.from_pretrained("gpt2")
model = PeftModel.from_pretrained(base_model, "skaterx/joke-lora-v1")
model.eval()

tokenizer = AutoTokenizer.from_pretrained("skaterx/joke-lora-v1")
tokenizer.pad_token = tokenizer.eos_token

generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")

class InferRequest(BaseModel):
    prompt: str
    max_length: int = 100
    temperature: float = 0.9

@app.post("/infer")
async def infer(request: InferRequest, credentials: HTTPAuthorizationCredentials = Depends(security)):
    if credentials.credentials != AUTH_TOKEN:
        raise HTTPException(status_code=401, detail="Invalid token")
    
    full_prompt = f"Şaka anlat:\n{request.prompt}\n\nŞaka:"
    result = generator(full_prompt, max_length=request.max_length, temperature=request.temperature, do_sample=True)
    return {"generated_text": result[0]["generated_text"]}