# -*- coding: utf-8 -*-
import os, sys
import tkinter as tk
from tkinter import filedialog
from pathlib import Path
import soundfile as sf

# =========================
#        CONFIG
# =========================
AUTO_MODE = True  # <<<< Set False to ask the user on launch

# Root directories (relative to this script’s folder)
RVC_DIR     = Path(__file__).resolve().parent
WEIGHTS_DIR = (RVC_DIR / "assets" / "weights").resolve()
INDEX_DIR   = (RVC_DIR / "logs").resolve()
UVR5_DIR    = (RVC_DIR / "assets" / "uvr5_weights").resolve()
HUBERT_PATH = (RVC_DIR / "assets" / "hubert" / "hubert_base.pt").resolve()

# Default params (only model/index filenames, not full paths)
AUTO_PARAMS = {
    "model_name": "Marceline EN.pth",
    "index_name": "Marceline EN.index",
    "voices_dir": r"SET YOUR VOICE PATH IN THE \"myInferFast.py\" SCRIPT FIRST <3",
    "pitch": 0,
    "f0method": "crepe",   # pm | harvest | crepe | rmvpe
    "index_rate": 0.75,
    "protect": 0.33,
}

# Build absolute paths automatically
AUTO_PARAMS["model_path"] = str((WEIGHTS_DIR / AUTO_PARAMS["model_name"]).resolve())
AUTO_PARAMS["index_path"] = str((INDEX_DIR / AUTO_PARAMS["index_name"]).resolve())

# =========================
#   RVC ENVIRONMENT SETUP
# =========================
sys.path.append(str(RVC_DIR))
os.chdir(str(RVC_DIR))

os.environ["weight_root"]      = str(WEIGHTS_DIR)
os.environ["index_root"]       = str(INDEX_DIR)
os.environ["weight_uvr5_root"] = str(UVR5_DIR)
if HUBERT_PATH.exists():
    os.environ["hubert_path"] = str(HUBERT_PATH)

print("[ENV] RVC_DIR     =", RVC_DIR)
print("[ENV] weight_root =", os.environ.get("weight_root"))
print("[ENV] index_root  =", os.environ.get("index_root"))
print("[ENV] hubert_path =", os.environ.get("hubert_path"))

# Shim for omegaconf import issue
try:
    import importlib, omegaconf
    sys.modules['_utils'] = importlib.import_module('omegaconf._utils')
except Exception:
    pass

# Imports RVC
from configs.config import Config
from infer.modules.vc.modules import VC

config = Config()
vc = VC(config)

# =========================
#   INTERACTIVE UTILITIES
# =========================
def choose_file(title, patterns):
    root = tk.Tk()
    root.withdraw()
    return filedialog.askopenfilename(title=title, filetypes=patterns)

def choose_dir(title):
    root = tk.Tk()
    root.withdraw()
    return filedialog.askdirectory(title=title)

def ask_params():
    print("=== INTERACTIVE MODE ===")
    model_path = choose_file("Choose a model (.pth)", [("RVC Model", "*.pth")])
    index_path = choose_file("Choose an index (.index)", [("RVC Index", "*.index")])
    voices_dir = choose_dir("Choose the VOICES folder")

    use_auto = (input("Load default params? (y/n) [y]: ").strip().lower() or "y") == "y"
    if use_auto:
        pitch     = AUTO_PARAMS["pitch"]
        f0method  = AUTO_PARAMS["f0method"]
        index_rate= AUTO_PARAMS["index_rate"]
        protect   = AUTO_PARAMS["protect"]
    else:
        pitch     = int(input("Pitch shift [0]: ") or "0")
        f0method  = (input("F0 method (pm/harvest/crepe/rmvpe) [crepe]: ").strip() or "crepe")
        index_rate= float(input("Index rate (0..1) [0.75]: ") or "0.75")
        protect   = float(input("Protect (0..0.5) [0.33]: ") or "0.33")

    return {
        "model_path": model_path,
        "index_path": index_path,
        "voices_dir": voices_dir,
        "pitch": pitch,
        "f0method": f0method,
        "index_rate": index_rate,
        "protect": protect,
    }

# =========================
#     FILE CONVERSION
# =========================
def convert_file(model_path, index_path, input_path, output_path, pitch, f0method, index_rate, protect):
    print(f"[+] Converting {os.path.basename(input_path)}")

    model_name = os.path.basename(model_path)
    try:
        vc.get_vc(model_name, protect, protect)
    except Exception as e:
        print(f"[!] Model load error {model_name}: {e}")
        return

    use_index = True
    if not index_path or not os.path.exists(index_path):
        print(f"[!] Index not found → disabled (index_rate=0). Expected: {index_path}")
        use_index = False

    info, (sr, wav_opt) = vc.vc_single(
        0,                                   # sid
        input_path,                          # audio source
        pitch,                               # semitones
        None,                                # f0_file
        f0method,                            # pm / harvest / crepe / rmvpe
        index_path if use_index else "",     # file_index
        None,                                # file_index2
        index_rate if use_index else 0.0,    # index_rate
        3,                                   # filter_radius
        0,                                   # resample_sr
        1.0,                                 # rms_mix_rate
        protect,                             # protect
    )

    if wav_opt is None:
        print(f"[!] Error with {input_path}:\n{info}")
        return

    try:
        sf.write(output_path, wav_opt, sr)
        print(f"[OK] Saved → {output_path}")
    except Exception as e:
        print(f"[!] Write failed {output_path}: {e}")

# =========================
#          MAIN
# =========================
def main():
    if AUTO_MODE:
        params = AUTO_PARAMS.copy()
        print("=== AUTO MODE ===")
    else:
        use_auto = (input("Load auto config? (y/n) [y]: ").strip().lower() or "y") == "y"
        if use_auto:
            params = AUTO_PARAMS.copy()
            print("=== AUTO MODE (forced) ===")
        else:
            params = ask_params()

    model_path = params["model_path"]
    index_path = params["index_path"]
    voices_dir = params["voices_dir"]
    pitch      = params["pitch"]
    f0method   = params["f0method"]
    index_rate = params["index_rate"]
    protect    = params["protect"]

    if not voices_dir or not os.path.isdir(voices_dir):
        print(f"[!] Invalid voices directory: {voices_dir}")
        input("\nPress any key to close...")
        return

    output_dir = os.path.join(voices_dir, "clone")
    os.makedirs(output_dir, exist_ok=True)

    exts = (".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg")
    files = [f for f in os.listdir(voices_dir) if f.lower().endswith(exts)]

    if not files:
        print(f"[!] No audio files found in: {voices_dir}")
        input("\nPress any key to close...")
        return

    for file in files:
        in_path = os.path.join(voices_dir, file)
        out_path = os.path.join(output_dir, Path(file).stem + "_converted.wav")
        try:
            convert_file(
                model_path,
                index_path,
                in_path,
                out_path,
                pitch,
                f0method,
                index_rate,
                protect,
            )
        except Exception as e:
            print(f"[!] Error with {file}: {e}")

    print(f"\n✅ Conversion complete → {output_dir}")
    input("Press any key to close...")

if __name__ == "__main__":
    main()