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")
    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():
    model_root = "weights"
    return [d for d in os.listdir(model_root) if os.path.isdir(f"{model_root}/{d}")]

def run_async_in_thread(fn, *args):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    result = loop.run_until_complete(fn(*args))
    loop.close()
    return result

async def tts(
    model_name,
    tts_text,
    tts_voice,
    index_rate,
    use_uploaded_voice,
    uploaded_voice,
):
    # ... (keep the existing implementation)

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

hubert_model = load_hubert()

rmvpe_model = RMVPE("rmvpe.pt", config.is_half, config.device)

class TTSProcessor:
    def __init__(self, config):
        self.config = config
        self.executor = ThreadPoolExecutor(max_workers=config.n_cpu)
        self.semaphore = asyncio.Semaphore(10)  # Limit to 10 concurrent tasks
        self.queue = asyncio.Queue()
        self.is_processing = False

    async def tts(self, model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice):
        task = asyncio.create_task(self._tts_task(model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice))
        await self.queue.put(task)
        
        if not self.is_processing:
            asyncio.create_task(self._process_queue())

        return await task

    async def _tts_task(self, model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice):
        async with self.semaphore:
            return await tts(model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice)

    async def _process_queue(self):
        self.is_processing = True
        while not self.queue.empty():
            task = await self.queue.get()
            try:
                await task
            except Exception as e:
                print(f"Error processing TTS task: {str(e)}")
            finally:
                self.queue.task_done()
        self.is_processing = False

# Initialize the TTSProcessor
tts_processor = TTSProcessor(config)

async def parallel_tts(tasks):
    return await asyncio.gather(*(tts_processor.tts(*task) for task in tasks))

def parallel_tts_wrapper(tasks):
    loop = asyncio.get_event_loop()
    return loop.run_until_complete(parallel_tts(tasks))