from diarization_processing import diarization_instance
from asr_processing import asr_instance
from audio_processing import process_instance
from fastapi import UploadFile, HTTPException
# from langdetection import lang_detect
import datetime
import aiofiles
from logger import logger
import numpy as np
import json
import os


class SpeechRecognitionService:
    def __init__(self, output_path="output"): 
        self.output_path = output_path
        self.enhanced_path_dir = os.path.join(output_path, "enhanced_chunks")
        os.makedirs(self.output_path, exist_ok=True)
        os.makedirs(self.enhanced_path_dir, exist_ok=True)

    async def save_temp_file(self, file: UploadFile, original_filename : str):
        temp_file_path = os.path.join(self.output_path,  "temp_" + original_filename)

        # Check if the file is an mp3 by examining the content type 
        if file.content_type == "audio/mpeg":
            logger.error("Uploaded file is not an WAV file.")
            raise HTTPException(status_code=422, detail="Invalid file format. Please upload an WAV file.")
        
        async with aiofiles.open(temp_file_path, "wb") as temp_file:
            content = await file.read()
            await temp_file.write(content)

        return temp_file_path
    

    def process_and_enhance_file(self, temp_file_path, noise_remover, speech_enhancement, vet_id, diarize_flag, original_filename):

        # Language detecion
        predicted_lang = process_instance.lang_detection(temp_file_path)
        default_langs = {"nl": "nl-NL", "en": "en-US"}
        predicted_lang = default_langs.get(predicted_lang)
        
        # Reassigning the audio file path variable
        enhanced_audio_path = temp_file_path
        enhanced_output_path = os.path.join(self.output_path, "enhanced_" + original_filename)
        
        # Using noise reduction tool to enhance audio quality 
        if noise_remover:
            logger.info(" Noise removing is being used to improve poor quality of the input audio.")
            enhanced_audio_path = process_instance.audoai_noise_remover(temp_file_path, enhanced_output_path)
            asr_instance.blob_storage(enhanced_audio_path)
        
        # Using speech enhancement tool to enhance audio quality 
        if speech_enhancement:
            logger.info(" Speech enhancing is being used to improve poor quality of the input audio.")
            enhanced_audio_path = process_instance.elevenlabs_speech_enhancer(temp_file_path, enhanced_output_path)
            asr_instance.blob_storage(enhanced_audio_path)

        def split_audio_and_save_chunks(audio_path, process_instance, enhanced_dir, original_filename, noise_remover=False, speech_enhancement=False):
            audio_chunks = process_instance.split_audio_on_silence(audio_path, 1000, None, 250)
            logger.info(f"Number of chunks: {len(audio_chunks)}")
            chunk_paths = process_instance.save_chunks(audio_chunks, enhanced_dir)
            
            os.remove(audio_path)
            return_filename = f"enhanced_{original_filename}" if (noise_remover or speech_enhancement) else original_filename
            return chunk_paths, return_filename
    
        try:
            if diarize_flag:
                file_in_blob = vet_id + '.json'
                local_file = diarization_instance.search_blob(file_in_blob)
                logger.info(f"VetId: {vet_id}, local_file: {local_file}")

                if local_file is None:  # if vet-id doesn't exist
                    logger.info("Embeddings file does not exist or input name is invalid. Proceeding without vet's config file.")
                    
                    chunk_paths, filename = split_audio_and_save_chunks(
                        enhanced_audio_path, 
                        process_instance, 
                        self.enhanced_path_dir, 
                        original_filename, 
                        noise_remover, 
                        speech_enhancement
                    )
                    return chunk_paths, filename, None, predicted_lang
                
                elif os.path.isfile(local_file): # if vet-id exist
                    with open(local_file, 'r') as f:
                        embeddings = json.load(f)
                        embeddings_array = np.array(embeddings)
                        os.remove(local_file)

                    # Diarization processing (distinguishing between vet and patient voices)
                    try:
                        logger.info(f"Enhanced_audio_path: {enhanced_audio_path}")
                        chunk_paths, speakers = diarization_instance.diarization(enhanced_audio_path, embeddings_array)  
                        
                        filename = f"enhanced_{original_filename}" if (noise_remover or speech_enhancement) else original_filename
                        return chunk_paths, filename, speakers, predicted_lang
                    
                    except Exception as e:
                        logger.error(f"Diarization couldn't be completed: {e}")
                        raise HTTPException(status_code=400, detail=f"Diarization couldn't be completed: {str(e)}")
            else:
                # If diarize_flag is False, process audio normally without diarization
                chunk_paths, filename = split_audio_and_save_chunks(
                    enhanced_audio_path, 
                    process_instance, 
                    self.enhanced_path_dir, 
                    original_filename, 
                    noise_remover, 
                    speech_enhancement
                )
                return chunk_paths, filename, None, predicted_lang
        
        except Exception as e:
            logger.error(f"An error occurred during processing: {e}")
            raise HTTPException(status_code=500, detail=f"Error processing audio: {str(e)}")

    def transcribe_chunks(self, chunk_paths, engine, diarize_flag, speakers, predicted_lang):
        
        transcripts = []
        for i, chunk_path in enumerate(chunk_paths):
            try:
                # Detect language of the audio sample

                # Google ASR for English and Dutch / Resample to 24kHz for Google Telephony
                if engine == "google":
                    process_instance.resample_audio(chunk_path, 24000)
                    phrase_set_id = "en-vet-words-domain" if predicted_lang == "en-US" else "nl-vet-words-domain"
                    transcript = asr_instance.asr_google_word_boosting(audio_uri=chunk_path, phrase_set_id=phrase_set_id, lang=predicted_lang)
                
                # Azure ASR / Resample to 16kHz for Azure
                elif engine == "azure":
                    process_instance.resample_audio(chunk_path, 16000)
                    transcript = asr_instance.asr_azure_continuous_recognition(chunk_path, predicted_lang)
                else:
                    raise HTTPException(status_code=403, detail="Invalid ASR engine or language. Supported: 'google', 'azure' with 'en-US', 'nl-NL'.")

                if diarize_flag and speakers:
                    speaker = speakers[i] if i < len(speakers) else "Unknown"
                    transcripts.append(f"({speaker}): {transcript}")
                else:
                    transcripts.append(transcript)

            except HTTPException as e:
                raise e
            except Exception as e:
                logger.error(f"Error transcribing {chunk_path}: {e}")
                raise HTTPException(status_code=400, detail=f"Error transcribing {chunk_path}: {str(e)}")
            finally:
                os.remove(chunk_path)  

        return " ".join(transcripts)



# refers to ("/recognize/v1/")
    async def recognize_speech(self, file: UploadFile, engine: str, vet_id : str, diarize_flag : bool, noise_remover: bool, speech_enhancement: bool):

        start_time = datetime.datetime.now()

        original_filename = file.filename
        temp_file_path = await self.save_temp_file(file, original_filename)
        chunk_paths, processed_filename, speakers, predicted_lang = self.process_and_enhance_file(temp_file_path, noise_remover, speech_enhancement, vet_id, diarize_flag, original_filename)
        full_transcript = self.transcribe_chunks(chunk_paths, engine, diarize_flag, speakers, predicted_lang)

        end_time = datetime.datetime.now()
        duration = end_time - start_time
    
        return {"processed-time" : duration, "processed-file": processed_filename, "transcription": full_transcript, "predicted-lang": predicted_lang, "engine": engine,"vet_id" : vet_id, "diarization": diarize_flag,  "noise-reduction": noise_remover, "speech-enhancement": speech_enhancement}

# refers to ("/upload_voice")
    async def voice_embeddings(self, doctor_speech_file: UploadFile):

        original_filename = doctor_speech_file.filename
        config_name = await diarization_instance.create_embeddings(self.output_path, doctor_speech_file, original_filename)
    
        return {"message" : "Successfully created and saved vet voice embeddings config file!", "processed-audio" : original_filename, "config-file" : config_name}

How to create and prepare an unit tests for this rest api applicatio?