import os
from pathlib import Path

import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer


MODEL_PATH = Path(os.getenv("MODEL_PATH", "/app/speechCleaner_t5_model")).resolve()
MAX_LENGTH = int(os.getenv("MAX_LENGTH", "256"))
NUM_BEAMS = int(os.getenv("NUM_BEAMS", "4"))

app = FastAPI(title="SignApp Disfluency Remover")

tokenizer = None
model = None
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class TextInput(BaseModel):
    text: str


def load_model():
    global tokenizer, model
    if tokenizer is None or model is None:
        if not MODEL_PATH.exists():
            raise RuntimeError(f"Model directory not found: {MODEL_PATH}")
        tokenizer = AutoTokenizer.from_pretrained(str(MODEL_PATH), local_files_only=True)
        model = AutoModelForSeq2SeqLM.from_pretrained(str(MODEL_PATH), local_files_only=True)
        model.to(device)
        model.eval()


@app.on_event("startup")
def startup():
    load_model()


@app.get("/health")
def health():
    return {"status": "ok", "device": str(device), "model_path": str(MODEL_PATH)}


@app.post("/clean/")
def clean(body: TextInput):
    text = body.text.strip()
    if not text:
        raise HTTPException(status_code=400, detail="Text is empty")

    load_model()
    inputs = tokenizer(
        "clean speech: " + text,
        return_tensors="pt",
        truncation=True,
        padding=True,
    ).to(device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=MAX_LENGTH,
            num_beams=NUM_BEAMS,
            early_stopping=True,
        )

    cleaned_text = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
    return {"cleaned_text": cleaned_text}