import os
import shutil
import torch
import logging
from typing import List, Dict, Any, Optional
import nltk
import wget
import json
from omegaconf import OmegaConf
from nemo.collections.asr.models import EncDecClassificationModel


punct_model_langs = [
    "en",
    "fr",
    "de",
    "es",
    "it",
    "nl",
    "pt",
    "bg",
    "pl",
    "cs",
    "sk",
    "sl",
]

LANGUAGES = {
    "en": "english",
    "zh": "chinese",
    "de": "german",
    "es": "spanish",
    "ru": "russian",
    "ko": "korean",
    "fr": "french",
    "ja": "japanese",
    "pt": "portuguese",
    "tr": "turkish",
    "pl": "polish",
    "ca": "catalan",
    "nl": "dutch",
    "ar": "arabic",
    "sv": "swedish",
    "it": "italian",
    "id": "indonesian",
    "hi": "hindi",
    "fi": "finnish",
    "vi": "vietnamese",
    "he": "hebrew",
    "uk": "ukrainian",
    "el": "greek",
    "ms": "malay",
    "cs": "czech",
    "ro": "romanian",
    "da": "danish",
    "hu": "hungarian",
    "ta": "tamil",
    "no": "norwegian",
    "th": "thai",
    "ur": "urdu",
    "hr": "croatian",
    "bg": "bulgarian",
    "lt": "lithuanian",
    "la": "latin",
    "mi": "maori",
    "ml": "malayalam",
    "cy": "welsh",
    "sk": "slovak",
    "te": "telugu",
    "fa": "persian",
    "lv": "latvian",
    "bn": "bengali",
    "sr": "serbian",
    "az": "azerbaijani",
    "sl": "slovenian",
    "kn": "kannada",
    "et": "estonian",
    "mk": "macedonian",
    "br": "breton",
    "eu": "basque",
    "is": "icelandic",
    "hy": "armenian",
    "ne": "nepali",
    "mn": "mongolian",
    "bs": "bosnian",
    "kk": "kazakh",
    "sq": "albanian",
    "sw": "swahili",
    "gl": "galician",
    "mr": "marathi",
    "pa": "punjabi",
    "si": "sinhala",
    "km": "khmer",
    "sn": "shona",
    "yo": "yoruba",
    "so": "somali",
    "af": "afrikaans",
    "oc": "occitan",
    "ka": "georgian",
    "be": "belarusian",
    "tg": "tajik",
    "sd": "sindhi",
    "gu": "gujarati",
    "am": "amharic",
    "yi": "yiddish",
    "lo": "lao",
    "uz": "uzbek",
    "fo": "faroese",
    "ht": "haitian creole",
    "ps": "pashto",
    "tk": "turkmen",
    "nn": "nynorsk",
    "mt": "maltese",
    "sa": "sanskrit",
    "lb": "luxembourgish",
    "my": "myanmar",
    "bo": "tibetan",
    "tl": "tagalog",
    "mg": "malagasy",
    "as": "assamese",
    "tt": "tatar",
    "haw": "hawaiian",
    "ln": "lingala",
    "ha": "hausa",
    "ba": "bashkir",
    "jw": "javanese",
    "su": "sundanese",
    "yue": "cantonese",
}

TO_LANGUAGE_CODE = {
    **{language: code for code, language in LANGUAGES.items()},
    "burmese": "my",
    "valencian": "ca",
    "flemish": "nl",
    "haitian": "ht",
    "letzeburgesch": "lb",
    "pushto": "ps",
    "panjabi": "pa",
    "moldavian": "ro",
    "moldovan": "ro",
    "sinhalese": "si",
    "castilian": "es",
}


def cleanup(path: str):
    """Clean up temporary files and directories."""
    if os.path.isfile(path) or os.path.islink(path):
        os.remove(path)
    elif os.path.isdir(path):
        shutil.rmtree(path)
    else:
        raise ValueError(f"Path {path} is not a file or dir.")


def process_language_arg(language: str, model_name: str) -> str:
    """Process and validate language argument."""
    if language is not None:
        language = language.lower()
    if language not in LANGUAGES:
        if language in TO_LANGUAGE_CODE:
            language = TO_LANGUAGE_CODE[language]
        else:
            raise ValueError(f"Unsupported language: {language}")

    if model_name.endswith(".en") and language != "en":
        if language is not None:
            logging.warning(
                f"{model_name} is an English-only model but received '{language}'; using English instead."
            )
        language = "en"
    return language


def format_timestamp(milliseconds, always_include_hours=False, decimal_marker='.'):
    seconds = milliseconds / 1000
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    seconds = seconds % 60
    hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
    return f"{hours_marker}{minutes:02d}:{seconds:06.3f}".replace('.', decimal_marker)


def setup_diarizer_config(output_dir: str) -> OmegaConf:
    """Set up configuration for the diarizer."""
    DOMAIN_TYPE = "telephonic"
    CONFIG_FILE_NAME = f"diar_infer_{DOMAIN_TYPE}.yaml"
    CONFIG_URL = f"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/inference/{CONFIG_FILE_NAME}"
    MODEL_CONFIG = os.path.join(output_dir, CONFIG_FILE_NAME)

    if not os.path.exists(MODEL_CONFIG):
        MODEL_CONFIG = wget.download(CONFIG_URL, output_dir)

    config = OmegaConf.load(MODEL_CONFIG)
    data_dir = os.path.join(output_dir, "data")
    os.makedirs(data_dir, exist_ok=True)

    meta = {
        "audio_filepath": os.path.join(output_dir, "mono_file.wav"),
        "offset": 0,
        "duration": None,
        "label": "infer",
        "text": "-",
        "rttm_filepath": None,
        "uem_filepath": None,
    }

    with open(os.path.join(data_dir, "input_manifest.json"), "w") as fp:
        json.dump(meta, fp)
        fp.write("\n")

    # Configure the diarizer
    config.num_workers = 0
    config.diarizer.manifest_filepath = os.path.join(
        data_dir, "input_manifest.json")
    config.diarizer.out_dir = output_dir
    config.diarizer.speaker_embeddings.model_path = "titanet_large"
    config.diarizer.oracle_vad = False

    config.diarizer.vad.model_path = "vad_multilingual_marblenet"
    config.diarizer.vad.parameters.onset = 0.8
    config.diarizer.vad.parameters.offset = 0.6
    config.diarizer.vad.parameters.pad_offset = -0.05
    config.diarizer.msdd_model.model_path = (
        "diar_msdd_telephonic"  # Telephonic speaker diarization model
    )
    # Set a default batch size for processing
    config.diarizer.msdd_model.parameters.batch_size = 16
    config.diarizer.clustering.parameters = {
        "oracle_num_speakers": False,
        "max_num_speakers": 8,
        "enhanced_count_thres": 80,
        "max_rp_threshold": 0.25,
        "sparse_search_volume": 30,
        "maj_vote_spk_count": False,
        "chunk_cluster_count": 50,
        "embeddings_per_chunk": 10000
    }
    return config


def check_audio_for_speech(audio_file_path):
    """
    Check if the given audio file contains speech using NeMo's VAD model.

    :param audio_file_path: Path to the audio file to check.
    :return: True if speech is detected, False otherwise.
    """
    # Load the pre-trained VAD model
    vad_model = EncDecClassificationModel.from_pretrained(
        model_name="vad_multilingual_marblenet")

    # Run VAD on the audio file
    vad_output = vad_model.transcribe([audio_file_path])

    # Convert vad_output to a list if it's a tensor
    if isinstance(vad_output, torch.Tensor):
        vad_output = vad_output.tolist()

    # Debugging: Print the structure of vad_output after conversion
    print("Converted VAD Output:", vad_output)

    # Assuming vad_output is now a list of lists or similar structure
    speech_detected = any('speech' in segment for segment in vad_output)

    print(f"Speech detected: {speech_detected}")
    return speech_detected


def check_audio_for_human_voice(audio_file_path):
    """
    Check if the given audio file contains human voice using NeMo's VAD model.

    :param audio_file_path: Path to the audio file to check.
    :return: True if human voice is detected, False otherwise.
    """
    try:
        # Load the pre-trained VAD model
        vad_model = EncDecClassificationModel.from_pretrained(model_name="vad_multilingual_marblenet")

        # Run VAD on the audio file
        vad_output = vad_model.transcribe([audio_file_path])

        # Convert vad_output to a list if it's a tensor
        if isinstance(vad_output, torch.Tensor):
            vad_output = vad_output.tolist()

        # Check if any human voice segments are detected
        human_voice_detected = any('speech' in segment for segment in vad_output)

        return human_voice_detected
    except ValueError as e:
        print(f"ValueError encountered: {e}")
        return False
    except Exception as e:
        print(f"An error occurred: {e}")
        return False
