import os
import wave
import asyncio
import torch
from fastapi import FastAPI, Depends, HTTPException, status, UploadFile, File
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from fastapi.responses import FileResponse
import nemo.collections.asr as nemo_asr
from piper.voice import PiperVoice
from pydantic import BaseModel

app = FastAPI(title="ASR & TTS API")
security = HTTPBasic()

# Basic Authentication Configuration
USERNAME = os.environ.get("API_USERNAME", "admin")
PASSWORD = os.environ.get("API_PASSWORD", "secret")

def verify_credentials(credentials: HTTPBasicCredentials = Depends(security)):
    if not (credentials.username == USERNAME and credentials.password == PASSWORD):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Basic"},
        )
    return credentials

# Global references for the models
asr_model = None
tts_voice = None

@app.on_event("startup")
async def load_models():
    global asr_model, tts_voice
    
    # 1. Load and Quantize NeMo ASR
    # Ensure you have uploaded your downloaded NeMo model to the Space with this filename
    nemo_path = "model.nemo" 
    if os.path.exists(nemo_path):
        device = torch.device('cpu')
        model = nemo_asr.models.EncDecCTCModel.restore_from(restore_path=nemo_path, map_location=device)
        model.freeze()
        
        # Apply CPU Dynamic Quantization for memory reduction and speed
        # model = torch.quantization.quantize_dynamic(
        #     model, {torch.nn.Linear}, dtype=torch.qint8
        # )
        torch.quantization.quantize_dynamic(
            model, {torch.nn.Linear}, dtype=torch.qint8, inplace=True
        )
        
        model.cur_decoder = 'ctc'
        asr_model = model
        print("ASR Model loaded and dynamically quantized.")
    else:
        print("WARNING: model.nemo not found. Please upload it.")

    # 2. Load Piper TTS (Nepali Chitwan)
    tts_model_path = "chitwan.onnx"
    if os.path.exists(tts_model_path):
        tts_voice = PiperVoice.load(tts_model_path)
        print("TTS Model loaded.")

@app.post("/asr")
async def transcribe(file: UploadFile = File(...), _: str = Depends(verify_credentials)):
    if not asr_model:
        raise HTTPException(status_code=503, detail="ASR model not loaded")
    
    # Save the uploaded audio temporarily
    audio_path = f"/tmp/{file.filename}"
    with open(audio_path, "wb") as f:
        f.write(await file.read())
    
    # Run the CPU-bound ASR transcription in a thread pool
    loop = asyncio.get_event_loop()
    transcription = await loop.run_in_executor(
        None, 
        # Pass the list positionally, and explicitly declare Nepali ('ne')
        lambda: asr_model.transcribe([audio_path], batch_size=1, language_id='ne')[0] 
    )
    
    os.remove(audio_path)
    return {"text": transcription}

class TTSRequest(BaseModel):
    text: str
    
@app.post("/tts")
async def synthesize(req: TTSRequest, _: str = Depends(verify_credentials)):
    if not tts_voice:
        raise HTTPException(status_code=503, detail="TTS model not loaded")
    
    output_path = "/tmp/output.wav"
    
    # Run the CPU-bound TTS synthesis in a thread pool
    loop = asyncio.get_event_loop()
    
    def generate_audio():
        with wave.open(output_path, "wb") as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(tts_voice.config.sample_rate)
            # Make sure to call req.text here!
            tts_voice.synthesize(req.text, wav_file)

    await loop.run_in_executor(None, generate_audio)
        
    return FileResponse(
        output_path, 
        media_type="audio/wav",
        headers={"Content-Disposition": "attachment; filename=tts_output.wav"}
    )