from pydantic import BaseModel, Field
import os
from pathlib import Path
from enum import Enum
from typing import Any, Tuple


# Constants
EXT_MODELS_DIRT = f"data{os.sep}ckpt{os.sep}ppg_extractor"
ENC_MODELS_DIRT = f"data{os.sep}ckpt{os.sep}encoder"


if os.path.isdir(EXT_MODELS_DIRT):    
    extractors =  Enum('extractors', list((file.name, file) for file in Path(EXT_MODELS_DIRT).glob("**/*.pt")))
    print("Loaded extractor models: " + str(len(extractors)))
else:
    raise Exception(f"Model folder {EXT_MODELS_DIRT} doesn't exist.")

if os.path.isdir(ENC_MODELS_DIRT):    
    encoders = Enum('encoders', list((file.name, file) for file in Path(ENC_MODELS_DIRT).glob("**/*.pt")))
    print("Loaded encoders models: " + str(len(encoders)))
else:
    raise Exception(f"Model folder {ENC_MODELS_DIRT} doesn't exist.")

class Model(str, Enum):
    VC_PPG2MEL = "ppg2mel"

class Dataset(str, Enum):
    AIDATATANG_200ZH = "aidatatang_200zh"
    AIDATATANG_200ZH_S = "aidatatang_200zh_s"

class Input(BaseModel):
    # def render_input_ui(st, input) -> Dict: 
    #     input["selected_dataset"] = st.selectbox(
    #         '选择数据集', 
    #         ("aidatatang_200zh", "aidatatang_200zh_s")
    #     )
    # return input
    model: Model = Field(
        Model.VC_PPG2MEL, title="目标模型",
    )
    dataset: Dataset = Field(
        Dataset.AIDATATANG_200ZH, title="数据集选择",
    )
    datasets_root: str = Field(
        ..., alias="数据集根目录", description="输入数据集根目录（相对/绝对）",
        format=True,
        example="..\\trainning_data\\"
    )
    output_root: str = Field(
        ..., alias="输出根目录", description="输出结果根目录（相对/绝对）",
        format=True,
        example="..\\trainning_data\\"
    )
    n_processes: int = Field(   
        2, alias="处理线程数", description="根据CPU线程数来设置",
        le=32, ge=1
    )
    extractor: extractors = Field(
        ..., alias="特征提取模型", 
        description="选择PPG特征提取模型文件."
    )
    encoder: encoders = Field(
        ..., alias="语音编码模型", 
        description="选择语音编码模型文件."
    )

class AudioEntity(BaseModel):
    content: bytes
    mel: Any

class Output(BaseModel):
    __root__: Tuple[str, int]

    def render_output_ui(self, streamlit_app, input) -> None:  # type: ignore
        """Custom output UI.
        If this method is implmeneted, it will be used instead of the default Output UI renderer.
        """
        sr, count = self.__root__
        streamlit_app.subheader(f"Dataset {sr} done processed total of {count}")

def preprocess(input: Input) -> Output:
    """Preprocess(预处理)"""
    finished = 0
    if input.model == Model.VC_PPG2MEL:
        from models.ppg2mel.preprocess import preprocess_dataset
        finished = preprocess_dataset(
            datasets_root=Path(input.datasets_root),
            dataset=input.dataset,
            out_dir=Path(input.output_root),
            n_processes=input.n_processes,
            ppg_encoder_model_fpath=Path(input.extractor.value),
            speaker_encoder_model=Path(input.encoder.value)
        )
    # TODO: pass useful return code
    return Output(__root__=(input.dataset, finished))