import os
import sys
import time
import signal
import io

from fastapi import FastAPI, Request, status, Form, UploadFile
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse

import fn
import gradio as gr
from app import demo

app = FastAPI()

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

gr.mount_gradio_app(app, demo, path="/gradio")

fn.load_model()

@app.post("/transcribe")
async def transcribe_audio(file: UploadFile = Form(...)):
    try:
        file_content = await file.read()
        file_stream = io.BytesIO(file_content)

        text_only, text_with_timestamps = fn.speech_to_text(file_stream)

        return {"transcription": text_only, "text_with_timestamps": text_with_timestamps}
    except Exception as e:
        return {"error": str(e)}

@app.post("/set_prompt")
async def set_prompt(prompt: str):
    try:
        fn.set_prompt(prompt)

        return {"status": 0}
    except Exception as e:
        return {"error": str(e)}