import ast
import json
import os

import torch
from mosestokenizer import *
from tts_infer.transliterate import XlitEngine
from tts_infer.tts import TextToMel, MelToWav

from src import log_setup, utilities
from src.config import settings

LOGGER = log_setup.get_logger(__name__)


def get_gpu_info(gpu):
    LOGGER.info(f"*** GPU is enabled: {gpu} ***")
    if gpu:
        no_gpus = torch.cuda.device_count()
        LOGGER.info(f"*** Total number of gpus allocated are {no_gpus} ***")
        LOGGER.info(f"*** Cuda Version {torch.version.cuda} ***")
        LOGGER.info(f"*** Python process id {os.getpid()} ***")
        LOGGER.info("*** The gpu device info : ***")
        for gpu in range(0, no_gpus):
            LOGGER.info(f"GPU {str(gpu)} - {str(torch.cuda.get_device_name(gpu))}")


class ModelService:

    def __init__(self):
        LOGGER.info(f'Loading with settings {settings}')
        gpu_present = torch.cuda.is_available()
        LOGGER.info("Gpu present : %s", gpu_present)
        get_gpu_info(settings.gpu)

        self.device = "cuda" if gpu_present & settings.gpu else "cpu"
        LOGGER.info("Using device : %s", self.device)

        model_config_file_path = settings.models_base_path + settings.model_config_file_path
        if os.path.exists(model_config_file_path):
            with open(model_config_file_path, 'r') as f:
                model_config = json.load(f)
        else:
            raise Exception(f'Model configuration file is missing at {model_config_file_path}')
        languages = ast.literal_eval(utilities.get_env_var('languages', '["all"]'))
        self.supported_languages = list(model_config.keys())
        LOGGER.info(f'supported languages {self.supported_languages}')
        self.available_choice = {}
        LOGGER.info(f'requested languages are {languages}')
        for language_code, lang_config in model_config.items():
            if language_code in languages or 'all' in languages:
                if lang_config.get("male_glow") is not None:
                    self.available_choice[f"{language_code}_male"] = [
                        TextToMel(glow_model_dir=settings.models_base_path + lang_config.get("male_glow"),
                                  device=self.device),
                        MelToWav(hifi_model_dir=settings.models_base_path + lang_config.get("male_hifi"),
                                 device=self.device)]
                if lang_config.get("female_glow") is not None:
                    self.available_choice[f"{language_code}_female"] = [
                        TextToMel(glow_model_dir=settings.models_base_path + lang_config.get("female_glow"),
                                  device=self.device),
                        MelToWav(hifi_model_dir=settings.models_base_path + lang_config.get("female_hifi"),
                                 device=self.device)]
                LOGGER.info(f'{language_code} Models initialized successfully')
        LOGGER.info(f'Model service available_choices are {self.available_choice}')
        if len(self.available_choice) == 0:
            LOGGER.error(
                f'Invalid languages requested in {languages} ,only supported languages are {self.supported_languages}')
            raise Exception(f'Invalid languages requested')
        self.transliterate_obj = XlitEngine()