import asyncio
import datetime
import logging
import os
import time
import traceback
import tempfile
from concurrent.futures import ThreadPoolExecutor

import edge_tts
import librosa
import torch
from fairseq import checkpoint_utils
import uuid

from config import Config
from lib.infer_pack.models import (
    SynthesizerTrnMs256NSFsid,
    SynthesizerTrnMs256NSFsid_nono,
    SynthesizerTrnMs768NSFsid,
    SynthesizerTrnMs768NSFsid_nono,
)
from rmvpe import RMVPE
from vc_infer_pipeline import VC

# Set logging levels
logging.getLogger("fairseq").setLevel(logging.WARNING)
logging.getLogger("numba").setLevel(logging.WARNING)
logging.getLogger("markdown_it").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)
logging.getLogger("matplotlib").setLevel(logging.WARNING)

limitation = os.getenv("SYSTEM") == "spaces"

config = Config()

# Edge TTS
tts_voice_list = asyncio.get_event_loop().run_until_complete(edge_tts.list_voices())
tts_voices = ["mn-MN-BataaNeural", "mn-MN-YesuiNeural"]  # Specific voices

# RVC models
model_root = "weights"
models = [d for d in os.listdir(model_root) if os.path.isdir(f"{model_root}/{d}")]
models.sort()

def get_unique_filename(extension):
    return f"{uuid.uuid4()}.{extension}"

def model_data(model_name):
    pth_path = [
        f"{model_root}/{model_name}/{f}"
        for f in os.listdir(f"{model_root}/{model_name}")
        if f.endswith(".pth")
    ][0]
    print(f"Loading {pth_path}")
    cpt = torch.load(pth_path, map_location="cpu", weights_only=True)
    tgt_sr = cpt["config"][-1]
    cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]  # n_spk
    if_f0 = cpt.get("f0", 1)
    version = cpt.get("version", "v1")
    if version == "v1":
        if if_f0 == 1:
            net_g = SynthesizerTrnMs256NSFsid(*cpt["config"], is_half=config.is_half)
        else:
            net_g = SynthesizerTrnMs256NSFsid_nono(*cpt["config"])
    elif version == "v2":
        if if_f0 == 1:
            net_g = SynthesizerTrnMs768NSFsid(*cpt["config"], is_half=config.is_half)
        else:
            net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"])
    else:
        raise ValueError("Unknown version")
    del net_g.enc_q
    net_g.load_state_dict(cpt["weight"], strict=False)
    print("Model loaded")
    net_g.eval().to(config.device)
    if config.is_half:
        net_g = net_g.half()
    else:
        net_g = net_g.float()
    vc = VC(tgt_sr, config)

    index_files = [
        f"{model_root}/{model_name}/{f}"
        for f in os.listdir(f"{model_root}/{model_name}")
        if f.endswith(".index")
    ]
    if len(index_files) == 0:
        print("No index file found")
        index_file = ""
    else:
        index_file = index_files[0]
        print(f"Index file found: {index_file}")

    return tgt_sr, net_g, vc, version, index_file, if_f0

def load_hubert():
    models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
        ["hubert_base.pt"],
        suffix="",
    )
    hubert_model = models[0]
    hubert_model = hubert_model.to(config.device)
    if config.is_half:
        hubert_model = hubert_model.half()
    else:
        hubert_model = hubert_model.float()
    return hubert_model.eval()

def get_model_names():
    return [d for d in os.listdir(model_root) if os.path.isdir(f"{model_root}/{d}")]

import nest_asyncio
nest_asyncio.apply()

def get_unique_filename(extension):
    return f"{uuid.uuid4().hex[:8]}.{extension}"

class TTSProcessor:
    def __init__(self, config):
        self.config = config
        self.executor = ThreadPoolExecutor(max_workers=config.n_cpu)
        self.semaphore = asyncio.Semaphore(config.max_concurrent_tts)
        self.last_request_time = time.time()
        self.rate_limit = config.tts_rate_limit
        self.temp_dir = tempfile.mkdtemp()

    async def tts(self, model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice):
        async with self.semaphore:
            current_time = time.time()
            time_since_last_request = current_time - self.last_request_time
            if time_since_last_request < 1 / self.rate_limit:
                await asyncio.sleep(1 / self.rate_limit - time_since_last_request)
            
            self.last_request_time = time.time()

            loop = asyncio.get_running_loop()
            return await loop.run_in_executor(
                self.executor,
                self._tts_process,
                model_name,
                tts_text,
                tts_voice,
                index_rate,
                use_uploaded_voice,
                uploaded_voice
            )

    def _tts_process(self, model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice):
        try:
            edge_output_filename = os.path.join(self.temp_dir, get_unique_filename("mp3"))

            if use_uploaded_voice:
                if uploaded_voice is None:
                    return "No voice file uploaded.", None, None
                
                uploaded_file_path = os.path.join(self.temp_dir, get_unique_filename("wav"))
                with open(uploaded_file_path, "wb") as f:
                    f.write(uploaded_voice)

                audio, sr = librosa.load(uploaded_file_path, sr=16000, mono=True)
            else:
                if limitation and len(tts_text) > 12000:
                    return (
                        f"Text characters should be at most 12000 in this huggingface space, but got {len(tts_text)} characters.",
                        None,
                        None,
                    )
                
                speed = 0  # Default speech speed
                speed_str = f"+{speed}%" if speed >= 0 else f"{speed}%"
                
                # Use synchronous approach for Edge TTS
                communicate = edge_tts.Communicate(tts_text, tts_voice, rate=speed_str)
                asyncio.get_event_loop().run_until_complete(communicate.save(edge_output_filename))

                audio, sr = librosa.load(edge_output_filename, sr=16000, mono=True)

            duration = len(audio) / sr
            if limitation and duration >= 20000:
                return (
                    f"Audio should be less than 20 seconds in this huggingface space, but got {duration}s.",
                    None,
                    None,
                )

            f0_up_key = 0
            tgt_sr, net_g, vc, version, index_file, if_f0 = model_data(model_name)

            if hasattr(self, 'model_rmvpe'):
                vc.model_rmvpe = self.model_rmvpe

            times = [0, 0, 0]
            audio_opt = vc.pipeline(
                hubert_model,
                net_g,
                0,
                audio,
                edge_output_filename if not use_uploaded_voice else uploaded_file_path,
                times,
                f0_up_key,
                "rmvpe",
                index_file,
                index_rate,
                if_f0,
                3,  # filter_radius
                tgt_sr,
                0,  # resample_sr
                0.25,  # rms_mix_rate
                version,
                0.33,  # protect
                None,
            )

            info = f"Success. Time: tts: {times[0]}s, npy: {times[1]}s, f0: {times[2]}s"
            print(info)
            return (
                info,
                edge_output_filename if not use_uploaded_voice else None,   
                (tgt_sr, audio_opt),
            )

        except Exception as e:
            logging.error(f"Error in TTS processing: {str(e)}")
            logging.error(traceback.format_exc())
            return str(e), None, None

    def __del__(self):
        # Clean up temporary directory
        import shutil
        shutil.rmtree(self.temp_dir, ignore_errors=True)

# Initialize global variables
tts_processor = TTSProcessor(config)
hubert_model = load_hubert()
rmvpe_model = RMVPE("rmvpe.pt", config.is_half, config.device)
tts_processor.model_rmvpe = rmvpe_model

voice_mapping = {
    "Mongolian Male": "mn-MN-BataaNeural",
    "Mongolian Female": "mn-MN-YesuiNeural"
}