"""
@Desc: 全局配置文件读取
"""

import os
import shutil
from typing import Dict, List

import torch
import yaml

from common.log import logger

# If not cuda available, set possible devices to cpu
cuda_available = torch.cuda.is_available()


class Resample_config:
    """重采样配置"""

    def __init__(self, in_dir: str, out_dir: str, sampling_rate: int = 44100):
        self.sampling_rate: int = sampling_rate  # 目标采样率
        self.in_dir: str = in_dir  # 待处理音频目录路径
        self.out_dir: str = out_dir  # 重采样输出路径

    @classmethod
    def from_dict(cls, dataset_path: str, data: Dict[str, any]):
        """从字典中生成实例"""

        # 不检查路径是否有效，此逻辑在resample.py中处理
        data["in_dir"] = os.path.join(dataset_path, data["in_dir"])
        data["out_dir"] = os.path.join(dataset_path, data["out_dir"])

        return cls(**data)


class Preprocess_text_config:
    """数据预处理配置"""

    def __init__(
        self,
        transcription_path: str,
        cleaned_path: str,
        train_path: str,
        val_path: str,
        config_path: str,
        val_per_lang: int = 5,
        max_val_total: int = 10000,
        clean: bool = True,
    ):
        self.transcription_path: str = (
            transcription_path  # 原始文本文件路径，文本格式应为{wav_path}|{speaker_name}|{language}|{text}。
        )
        self.cleaned_path: str = (
            cleaned_path  # 数据清洗后文本路径，可以不填。不填则将在原始文本目录生成
        )
        self.train_path: str = (
            train_path  # 训练集路径，可以不填。不填则将在原始文本目录生成
        )
        self.val_path: str = (
            val_path  # 验证集路径，可以不填。不填则将在原始文本目录生成
        )
        self.config_path: str = config_path  # 配置文件路径
        self.val_per_lang: int = val_per_lang  # 每个speaker的验证集条数
        self.max_val_total: int = (
            max_val_total  # 验证集最大条数，多于的会被截断并放到训练集中
        )
        self.clean: bool = clean  # 是否进行数据清洗

    @classmethod
    def from_dict(cls, dataset_path: str, data: Dict[str, any]):
        """从字典中生成实例"""

        data["transcription_path"] = os.path.join(
            dataset_path, data["transcription_path"]
        )
        if data["cleaned_path"] == "" or data["cleaned_path"] is None:
            data["cleaned_path"] = None
        else:
            data["cleaned_path"] = os.path.join(dataset_path, data["cleaned_path"])
        data["train_path"] = os.path.join(dataset_path, data["train_path"])
        data["val_path"] = os.path.join(dataset_path, data["val_path"])
        data["config_path"] = os.path.join(dataset_path, data["config_path"])

        return cls(**data)


class Bert_gen_config:
    """bert_gen 配置"""

    def __init__(
        self,
        config_path: str,
        num_processes: int = 2,
        device: str = "cuda",
        use_multi_device: bool = False,
    ):
        self.config_path = config_path
        self.num_processes = num_processes
        if not cuda_available:
            device = "cpu"
        self.device = device
        self.use_multi_device = use_multi_device

    @classmethod
    def from_dict(cls, dataset_path: str, data: Dict[str, any]):
        data["config_path"] = os.path.join(dataset_path, data["config_path"])

        return cls(**data)


class Style_gen_config:
    """style_gen 配置"""

    def __init__(
        self,
        config_path: str,
        num_processes: int = 4,
        device: str = "cuda",
    ):
        self.config_path = config_path
        self.num_processes = num_processes
        if not cuda_available:
            device = "cpu"
        self.device = device

    @classmethod
    def from_dict(cls, dataset_path: str, data: Dict[str, any]):
        data["config_path"] = os.path.join(dataset_path, data["config_path"])

        return cls(**data)


class Train_ms_config:
    """训练配置"""

    def __init__(
        self,
        config_path: str,
        env: Dict[str, any],
        # base: Dict[str, any],
        model_dir: str,
        num_workers: int,
        spec_cache: bool,
        keep_ckpts: int,
    ):
        self.env = env  # 需要加载的环境变量
        # self.base = base  # 底模配置
        self.model_dir = model_dir  # 训练模型存储目录，该路径为相对于dataset_path的路径，而非项目根目录
        self.config_path = config_path  # 配置文件路径
        self.num_workers = num_workers  # worker数量
        self.spec_cache = spec_cache  # 是否启用spec缓存
        self.keep_ckpts = keep_ckpts  # ckpt数量

    @classmethod
    def from_dict(cls, dataset_path: str, data: Dict[str, any]):
        # data["model"] = os.path.join(dataset_path, data["model"])
        data["config_path"] = os.path.join(dataset_path, data["config_path"])

        return cls(**data)


class Webui_config:
    """webui 配置 (for webui.py, not supported now)"""

    def __init__(
        self,
        device: str,
        model: str,
        config_path: str,
        language_identification_library: str,
        port: int = 7860,
        share: bool = False,
        debug: bool = False,
    ):
        if not cuda_available:
            device = "cpu"
        self.device: str = device
        self.model: str = model  # 端口号
        self.config_path: str = config_path  # 是否公开部署，对外网开放
        self.port: int = port  # 是否开启debug模式
        self.share: bool = share  # 模型路径
        self.debug: bool = debug  # 配置文件路径
        self.language_identification_library: str = (
            language_identification_library  # 语种识别库
        )

    @classmethod
    def from_dict(cls, dataset_path: str, data: Dict[str, any]):
        data["config_path"] = os.path.join(dataset_path, data["config_path"])
        data["model"] = os.path.join(dataset_path, data["model"])
        return cls(**data)


class Server_config:
    def __init__(
        self,
        port: int = 5000,
        device: str = "cuda",
        limit: int = 100,
        language: str = "JP",
        origins: List[str] = None,
    ):
        self.port: int = port
        if not cuda_available:
            device = "cpu"
        self.device: str = device
        self.language: str = language
        self.limit: int = limit
        self.origins: List[str] = origins

    @classmethod
    def from_dict(cls, data: Dict[str, any]):
        return cls(**data)


class Translate_config:
    """翻译api配置"""

    def __init__(self, app_key: str, secret_key: str):
        self.app_key = app_key
        self.secret_key = secret_key

    @classmethod
    def from_dict(cls, data: Dict[str, any]):
        return cls(**data)


class Config:
    def __init__(self, config_path: str, path_config: dict[str, str]):
        if not os.path.isfile(config_path) and os.path.isfile("default_config.yml"):
            shutil.copy(src="default_config.yml", dst=config_path)
            logger.info(
                f"A configuration file {config_path} has been generated based on the default configuration file default_config.yml."
            )
            logger.info(
                "If you have no special needs, please do not modify default_config.yml."
            )
            # sys.exit(0)
        with open(file=config_path, mode="r", encoding="utf-8") as file:
            yaml_config: Dict[str, any] = yaml.safe_load(file.read())
            model_name: str = yaml_config["model_name"]
            self.model_name: str = model_name
            if "dataset_path" in yaml_config:
                dataset_path = yaml_config["dataset_path"]
            else:
                dataset_path = os.path.join(path_config["dataset_root"], model_name)
            self.dataset_path: str = dataset_path
            self.assets_root: str = path_config["assets_root"]
            self.out_dir = os.path.join(self.assets_root, model_name)
            self.resample_config: Resample_config = Resample_config.from_dict(
                dataset_path, yaml_config["resample"]
            )
            self.preprocess_text_config: Preprocess_text_config = (
                Preprocess_text_config.from_dict(
                    dataset_path, yaml_config["preprocess_text"]
                )
            )
            self.bert_gen_config: Bert_gen_config = Bert_gen_config.from_dict(
                dataset_path, yaml_config["bert_gen"]
            )
            self.style_gen_config: Style_gen_config = Style_gen_config.from_dict(
                dataset_path, yaml_config["style_gen"]
            )
            self.train_ms_config: Train_ms_config = Train_ms_config.from_dict(
                dataset_path, yaml_config["train_ms"]
            )
            self.webui_config: Webui_config = Webui_config.from_dict(
                dataset_path, yaml_config["webui"]
            )
            self.server_config: Server_config = Server_config.from_dict(
                yaml_config["server"]
            )
            # self.translate_config: Translate_config = Translate_config.from_dict(
            #     yaml_config["translate"]
            # )


with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f:
    path_config: dict[str, str] = yaml.safe_load(f.read())
    # Should contain the following keys:
    # - dataset_root: the root directory of the dataset, default to "Data"
    # - assets_root: the root directory of the assets, default to "model_assets"


try:
    config = Config("config.yml", path_config)
except (TypeError, KeyError):
    logger.warning("Old config.yml found. Replace it with default_config.yml.")
    shutil.copy(src="default_config.yml", dst="config.yml")
    config = Config("config.yml", path_config)