# asr_api.py

import os
import tempfile
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
import logging
from dotenv import load_dotenv
from whisper import transcribe, post_process_transcript, setup_logger

# 加载环境变量
load_dotenv('environment.env')

# 设置日志
logger = setup_logger()

app = FastAPI()

@app.post("/transcribe/")
async def transcribe_audio(file: UploadFile = File(...)):
    try:
        content_type = file.content_type
        logger.info(f"Received file: {file.filename}, Content-Type: {content_type}")
        
        if content_type not in ["audio/mpeg", "audio/mp4", "audio/x-m4a", "audio/wav"]:
            raise HTTPException(status_code=400, detail=f"Unsupported file type: {content_type}")

        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as temp_file:
            content = await file.read()
            temp_file.write(content)
            temp_file.flush()

            logger.info(f"Temporary file created: {temp_file.name}")

            transcript = transcribe(temp_file.name)

        os.unlink(temp_file.name)
        logger.info(f"Temporary file deleted: {temp_file.name}")

        if transcript is None:
            raise ValueError("Transcription failed. Check server logs for details.")
        
        corrected_transcript = post_process_transcript(transcript)
        
        return JSONResponse(content={
            "original_transcript": transcript,
            "corrected_transcript": corrected_transcript
        })
    except ValueError as ve:
        logger.exception("Transcription failed")
        raise HTTPException(status_code=400, detail=str(ve))
    except Exception as e:
        logger.exception("An unexpected error occurred during transcription")
        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("asr_api:app", host="0.0.0.0", port=8000, reload=True)