from voice_changer.RVC.RVCModelSlotGenerator import RVCModelSlotGenerator
from voice_changer.VoiceChangerParamsManager import VoiceChangerParamsManager
from voice_changer.utils.LoadModelParams import LoadModelParamFile, LoadModelParams
from voice_changer.utils.VoiceChangerParams import VoiceChangerParams
model_dir1 = "model_dir/"+model_slot+"/"
is_pth = True  # Set this to True if you want to search for .pth files, or False for .onnx files
file_extension = ".pth" if is_pth else ".onnx"
# pth_files = [f for f in os.listdir(model_dir1) if f.endswith(file_extension)]
pth_files = [f for f in os.listdir(model_dir1) if f.endswith(".pth") or f.endswith(".onnx")]
print(pth_files)
index_files = [f for f in os.listdir(model_dir1) if f.endswith(".index")]
if pth_files:
    model_name = pth_files[0].replace(".pth", "")
else:
    model_name = "Null"
if index_files:
    index_name = index_files[0].replace(".index", "")
else:
    index_name = ""
original_string = str(pth_files)
string_pth_files = original_string[2:-2]
print("IM A STRING"+original_string)
print(model_name)
voiceChangerParams = VoiceChangerParams(
        model_dir="./model_dir/"+model_slot,
        content_vec_500="",
        content_vec_500_onnx="",
        content_vec_500_onnx_on="",
        hubert_base="",
        hubert_base_jp="",
        hubert_soft="",
        nsf_hifigan="",
        crepe_onnx_full="",
        crepe_onnx_tiny="",
        rmvpe="",
        rmvpe_onnx="",
        sample_mode=""
  )
vcparams = VoiceChangerParamsManager.get_instance()
vcparams.setParams(voiceChangerParams)
file = LoadModelParamFile(
     name=string_pth_files,
     kind="rvcModel",
     dir="",
)
loadParam = LoadModelParams(
    voiceChangerType="RVC",
    files=[file],
    slot="",
    isSampleMode=False,
    sampleId="",
    params={},
  )
slotInfo = RVCModelSlotGenerator.loadModel(loadParam)
print(slotInfo.samplingRate)
#---------------Make the Json File-----------
params_content = {
    "slotIndex": -1,
    "voiceChangerType": "RVC",
    "name": model_name,
    "description": "",
    "credit": "",
    "termsOfUseUrl": "",
    "iconFile": iconFile,
    "speakers": {
        "0": "target"
    },
    "modelFile": string_pth_files,
    "indexFile": f"{index_name}.index",
    "defaultTune": Tune,
    "defaultIndexRatio": Index,
    "defaultProtect": 0.5,
    "isONNX": False,
    "modelType": "pyTorchRVCv2",
    "samplingRate": slotInfo.samplingRate,
    "f0": True,
    "embChannels": 768,
    "embOutputLayer": 12,
    "useFinalProj": False,
    "deprecated": False,
    "embedder": "hubert_base",
    "sampleId": ""
}
# Write the content to params.json
with open(f"{model_dir1}/params.json", "w") as param_file:
    json.dump(params_content, param_file)