from fastapi import APIRouter, FastAPI, Form, UploadFile
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

from modules.asr_frontend import ASRModel, audio_from_file
from modules.tts_frontend import TextInput, TTSModel

router = APIRouter()
tts_model = TTSModel()
asr_model = ASRModel()


@router.get("/")
async def index() -> FileResponse:
    return FileResponse(path="template/index.html", media_type="text/html")


@router.get("/api/tts/load_model")
async def load_tts_model(model_name: str, fs: int) -> JSONResponse:
    try:
        tts_model.load_model(model_name, fs)
        return JSONResponse(status_code=200, content="Successfully loaded model")
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})


@router.post("/api/tts/generate")
async def generate_tts(text: TextInput):
    try:
        tts_model.generate(text)
        return JSONResponse(status_code=200)
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})


@router.get("/api/asr/load_model")
async def load_asr_model(model_name: str) -> JSONResponse:
    try:
        asr_model.load_model(model_name)
        return JSONResponse(status_code=200, content="Successfully loaded model")
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})


@router.post("/api/asr/recognize")
async def recognize(audio: UploadFile = Form(...)):
    audio = await audio_from_file(audio.file)
    try:
        text = asr_model.generate(audio)
        return JSONResponse(status_code=200, content={"text": text})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})


app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")

app.include_router(router)