import torch
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

BASE = "Qwen/Qwen2.5-0.5B"
ADAPTER = "BobCodeur/qwen2.5-0.5b-wolof"

# Chargement une seule fois au démarrage : base + adaptateur LoRA wolof
tokenizer = AutoTokenizer.from_pretrained(ADAPTER)
base = AutoModelForCausalLM.from_pretrained(BASE)
model = PeftModel.from_pretrained(base, ADAPTER)
model.eval()

app = FastAPI(title="Wolof Text Generation API")

# CORS ouvert : l'API est appelable depuis n'importe quel site web
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class GenIn(BaseModel):
    text: str
    max_length: int = 80
    temperature: float = 0.6


def generer(text: str, max_length: int, temperature: float) -> str:
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        sortie = model.generate(
            **inputs,
            max_new_tokens=int(max_length),
            do_sample=True,
            temperature=float(temperature),
            top_p=0.9,
            repetition_penalty=1.15,
        )
    return tokenizer.decode(sortie[0], skip_special_tokens=True)


@app.get("/")
def root():
    return {
        "message": "Wolof Text Generation API",
        "model": ADAPTER,
        "endpoints": {
            "GET /health": "vérifie que le service est réveillé",
            "POST /generate": "corps JSON {text, max_length?, temperature?}",
            "GET /generate?text=...": "test rapide au navigateur",
        },
    }


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


@app.post("/generate")
def generate_post(inp: GenIn):
    generated = generer(inp.text, inp.max_length, inp.temperature)
    return {"prompt": inp.text, "generated": generated}


@app.get("/generate")
def generate_get(text: str, max_length: int = 80, temperature: float = 0.6):
    generated = generer(text, max_length, temperature)
    return {"prompt": text, "generated": generated}