import numpy as np
import torch
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from faster_whisper import WhisperModel

app = FastAPI()

# Load ASR model using faster-whisper from local directory
print("Loading Whisper model from local whisper-small-yor-ct2...")
# Select device and compute type dynamically
device = "cuda" if torch.cuda.is_available() else "cpu"
compute_type = "float16" if torch.cuda.is_available() else "int8"

asr_model = WhisperModel(
    "whisper-small-yor-ct2",
    device=device,
    compute_type=compute_type
)
print("ASR model loaded successfully.")

from silero_vad import VADIterator, load_silero_vad

# Load Silero VAD model
print("Loading Silero VAD model...")
vad_model = load_silero_vad()
print("VAD model loaded successfully.")

@app.get("/")
def read_root():
    return {"message": "Yoruba Real-time ASR Server is running."}

@app.websocket("/stream")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    print("WebSocket client connected.")
    
    # Initialize VAD iterator with a sample rate of 16000 and 0.6 threshold for noise robustness
    vad_iterator = VADIterator(vad_model, threshold=0.6, sampling_rate=16000)
    
    # We will accumulate audio samples for transcription
    speech_buffer = []
    accumulator = np.array([], dtype=np.float32)
    
    try:
        while True:
            # Receive audio bytes (16 kHz PCM 16-bit mono)
            data = await websocket.receive_bytes()
            if not data:
                break
                
            # Convert bytes to numpy int16, then float32 scaled to [-1.0, 1.0]
            audio_chunk = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
            
            # Append new chunk to the accumulator
            accumulator = np.concatenate((accumulator, audio_chunk))
            
            # Process in blocks of exactly 512 samples (Silero VAD requirement)
            chunk_size = 512
            while len(accumulator) >= chunk_size:
                sub_chunk = accumulator[:chunk_size]
                accumulator = accumulator[chunk_size:]
                
                tensor_chunk = torch.from_numpy(sub_chunk)
                speech_dict = vad_iterator(tensor_chunk, return_seconds=True)
                
                # If speech is detected, accumulate the frames
                if vad_iterator.triggered:
                    speech_buffer.append(sub_chunk)
                
                # Force transcription if a single speech segment exceeds ~5.0 seconds
                # 5 seconds @ 16kHz is 80000 samples. 156 chunks * 512 = 79872 samples (~5.0s).
                force_transcribe = len(speech_buffer) >= 156
                
                # When speech ends, or we force it due to length
                if (speech_dict and "end" in speech_dict) or force_transcribe:
                    if speech_buffer:
                        full_audio = np.concatenate(speech_buffer)
                        # Clear buffer
                        speech_buffer = []
                        
                        print(f"Transcribing {len(full_audio)/16000:.2f}s of speech (forced={force_transcribe})...")
                        # Transcribe using faster-whisper
                        segments, info = asr_model.transcribe(
                            full_audio, 
                            language="yo", 
                            beam_size=2,
                            condition_on_previous_text=False,
                            no_repeat_ngram_size=3
                        )
                        text = " ".join([segment.text for segment in segments]).strip()
                        
                        if text:
                            print(f"Result: {text}")
                            await websocket.send_json({
                                "status": "final",
                                "text": text
                            })
                            
                    vad_iterator.reset_states()
                    
    except WebSocketDisconnect:
        print("WebSocket client disconnected.")
    except Exception as e:
        print(f"Error in stream processing: {e}")
    finally:
        # Transcribe any remaining audio in the buffer at disconnection
        # If we are currently in a speech segment, add remaining accumulator samples
        if len(accumulator) > 0 and vad_iterator.triggered:
            speech_buffer.append(accumulator)
            
        if speech_buffer:
            try:
                full_audio = np.concatenate(speech_buffer)
                print(f"Transcribing remaining {len(full_audio)/16000:.2f}s of speech...")
                segments, info = asr_model.transcribe(
                    full_audio, 
                    language="yo", 
                    beam_size=2,
                    condition_on_previous_text=False,
                    no_repeat_ngram_size=3
                )
                text = " ".join([segment.text for segment in segments]).strip()
                if text:
                    print(f"Final Result: {text}")
                    await websocket.send_json({
                        "status": "final",
                        "text": text
                    })
            except Exception as e:
                print(f"Error sending final transcript: {e}")
        try:
            await websocket.close()
        except Exception:
            pass
        print("WebSocket connection cleaned up.")