#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Exports RVC (Retrieval-based Voice Conversion) models from PyTorch (.pth)
to ONNX format, simplifying the model optionally.
Prompts the user for the input model path and optionally verifies the output.
Logs output to console and 'conversion.log'. Keeps window open on exit.
Includes workaround for onnxsim dynamic shape issue and weight_norm FutureWarning.
"""

import logging
import os
import sys
import argparse
import time
import warnings # Added for warning suppression

# Third-party libraries
import onnx
import onnxsim
import torch

# Optional verification library
try:
    import onnxruntime as ort
    ORT_AVAILABLE = True
except ImportError:
    ORT_AVAILABLE = False


# --- Configure Logging (File and Console) ---
log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
root_logger = logging.getLogger()
root_logger.setLevel(logging.INFO)

# Console Handler
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(log_formatter)
root_logger.addHandler(console_handler)

# File Handler (log to conversion.log in script's directory)
try:
    script_dir = os.path.dirname(os.path.abspath(__file__))
    log_file_path = os.path.join(script_dir, "conversion.log")
    # 'a' appends to the log, 'w' overwrites
    file_handler = logging.FileHandler(log_file_path, mode='a', encoding='utf-8')
    file_handler.setFormatter(log_formatter)
    root_logger.addHandler(file_handler)
    # Use print only for the initial notification about the log file location
    print(f"Logging detailed output to: {log_file_path}")
except Exception as e:
    print(f"Warning: Could not configure file logging. Error: {e}", file=sys.stderr)


# --- RVC Model Import ---
# Suppress the specific FutureWarning during import and model loading if needed
# Note: Ideally, the RVC library itself should be updated.
try:
    # This context manager is generally better applied specifically around the
    # code that triggers the warning (model instantiation/loading).
    # Placing it here might suppress warnings during import of the module itself.
    # with warnings.catch_warnings():
    #    warnings.filterwarnings("ignore", category=FutureWarning, message=".*`torch.nn.utils.weight_norm` is deprecated.*")
    from infer.lib.infer_pack.models_onnx import SynthesizerTrnMsNSFsidM
except ImportError as e:
    logging.error(f"Error importing SynthesizerTrnMsNSFsidM: {e}", exc_info=True)
    logging.error("Please ensure the script is run from within the RVC "
                  "project structure or that the RVC library's root "
                  "directory is in your PYTHONPATH.")
    input("Import error occurred. Press Enter to exit...")
    sys.exit(1)
except Exception as e:
    logging.error(f"An unexpected error occurred during import: {e}", exc_info=True)
    input("Import error occurred. Press Enter to exit...")
    sys.exit(1)


class RvcOnnxExporter:
    """Handles the export of an RVC model to ONNX format."""

    INPUT_NAMES = ["phone", "phone_lengths", "pitch", "pitchf", "ds", "rnd"]
    OUTPUT_NAMES = ["audio"]

    def __init__(self, opset_version: int = 18, simplify: bool = True):
        """
        Initializes the exporter.

        Args:
            opset_version (int): The ONNX opset version for export.
            simplify (bool): Whether to simplify the exported model using onnxsim.
        """
        self.opset_version = opset_version
        self.simplify = simplify
        self.config = None
        self.model_version = "v1"
        self.vec_channels = 256
        self.n_spk = 1
        self.posterior_channels = 192 # Default value used in error case before
        self.dummy_input_shapes = {} # Store shapes for simplify step

    def _load_checkpoint(self, model_path: str) -> dict | None:
        """Loads the PyTorch checkpoint file."""
        if not os.path.exists(model_path):
            logging.error(f"Model file not found: {model_path}")
            return None
        try:
            logging.info(f"Loading checkpoint: {model_path}")
            checkpoint = torch.load(model_path, map_location=torch.device("cpu"))

            required_keys = ['config', 'weight']
            if not all(key in checkpoint for key in required_keys):
                missing = [k for k in required_keys if k not in checkpoint]
                logging.error(f"Checkpoint missing required keys: {missing}. "
                              f"Found keys: {list(checkpoint.keys())}")
                return None
            if "emb_g.weight" not in checkpoint["weight"]:
                logging.error("Speaker embedding 'emb_g.weight' not found in "
                              "checkpoint weights.")
                return None
            logging.info("Checkpoint loaded successfully.")
            return checkpoint
        except Exception as e:
            logging.error(f"Failed to load checkpoint '{model_path}': {e}",
                          exc_info=True)
            return None

    def _prepare_config(self, checkpoint: dict) -> bool:
        """Determines model version, channels, and speaker count from config."""
        try:
            logging.info("Preparing model configuration...")
            self.config = list(checkpoint["config"])
            self.model_version = checkpoint.get("version", "v1")
            self.vec_channels = 768 if self.model_version == "v2" else 256
            logging.info(f"Detected model version: {self.model_version}, "
                         f"Feature channels: {self.vec_channels}")

            # Use config[2] for posterior_channels if valid, otherwise log error
            if len(self.config) > 2 and isinstance(self.config[2], int) and self.config[2] > 0:
                 self.posterior_channels = self.config[2]
                 logging.info(f"Using posterior encoder channels (config[2]): "
                              f"{self.posterior_channels} for 'rnd' input.")
            else:
                # Log the error but might proceed if the old default was somehow correct (unlikely)
                # Consider returning False here for stricter validation
                config2_val = self.config[2] if len(self.config) > 2 else 'N/A'
                logging.error("Could not determine valid posterior encoder channels "
                              f"from config[2] (value: {config2_val}). "
                              f"Using fallback: {self.posterior_channels}")
                # return False # Uncomment for stricter validation

            n_spk_inferred = checkpoint["weight"]["emb_g.weight"].shape[0]
            if len(self.config) < 3:
                 logging.error(f"Config list has fewer than 3 elements "
                               f"({len(self.config)}), cannot access speaker "
                               f"count at index -3.")
                 return False

            n_spk_config = self.config[-3]
            if n_spk_config != n_spk_inferred:
                logging.warning(f"Mismatch: Config speaker count ({n_spk_config})"
                                f" != Inferred count ({n_spk_inferred}). "
                                f"Updating config.")
                self.config[-3] = n_spk_inferred

            self.n_spk = n_spk_inferred
            logging.info(f"Using speaker count (n_spk): {self.n_spk}")
            self.config = tuple(self.config)
            logging.info("Model configuration prepared successfully.")
            return True

        except (IndexError, KeyError, TypeError, ValueError) as e:
            logging.error(f"Failed to prepare config from checkpoint: {e}",
                          exc_info=True)
            return False

    def _create_dummy_inputs(self, seq_len: int = 200) -> tuple | None:
        """Creates dummy input tensors for ONNX tracing and stores their shapes."""
        try:
            logging.info(f"Creating dummy inputs (sequence length: {seq_len})...")
            # Create tensors
            dummy_phone = torch.rand(1, seq_len, self.vec_channels, dtype=torch.float32)
            dummy_phone_lengths = torch.tensor([seq_len], dtype=torch.long)
            dummy_pitch = torch.randint(low=5, high=255, size=(1, seq_len), dtype=torch.long)
            dummy_pitchf = torch.rand(1, seq_len, dtype=torch.float32)
            dummy_ds = torch.tensor([0], dtype=torch.long) # Speaker ID
            dummy_rnd = torch.rand(1, self.posterior_channels, seq_len, dtype=torch.float32)

            inputs_tuple = (dummy_phone, dummy_phone_lengths, dummy_pitch,
                            dummy_pitchf, dummy_ds, dummy_rnd)

            # Store shapes mapped to names for use in simplification step
            self.dummy_input_shapes = {
                name: list(tensor.shape)
                for name, tensor in zip(self.INPUT_NAMES, inputs_tuple)
            }
            logging.info(f"Dummy input shapes: {self.dummy_input_shapes}")
            logging.info("Dummy inputs created successfully.")
            return inputs_tuple
        except Exception as e:
            logging.error(f"Failed to create dummy inputs: {e}", exc_info=True)
            return None

    def _instantiate_and_load_model(self, checkpoint: dict) -> SynthesizerTrnMsNSFsidM | None:
        """Instantiates the RVC model and loads weights, suppressing specific warnings."""
        try:
            logging.info("Instantiating model SynthesizerTrnMsNSFsidM...")
            # Suppress the specific weight_norm FutureWarning during model init and loading
            # Note: The ideal solution is to update the RVC library code itself.
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", category=FutureWarning, message=".*`torch.nn.utils.weight_norm` is deprecated.*")
                net_g = SynthesizerTrnMsNSFsidM(
                    *self.config,
                    is_half=False,
                    version=self.model_version
                )
                logging.info("Loading state dictionary into the model...")
                net_g.load_state_dict(checkpoint["weight"], strict=False)

            net_g.eval()
            logging.info("Model instantiated and weights loaded successfully.")
            return net_g
        except Exception as e:
            logging.error(f"Failed to instantiate or load model: {e}",
                          exc_info=True)
            return None

    def _simplify_model(self, onnx_path: str) -> bool:
        """Simplifies the exported ONNX model using onnxsim, providing input shapes."""
        try:
            logging.info("Simplifying ONNX model using onnxsim...")
            if not os.path.exists(onnx_path):
                logging.error(f"ONNX file not found at {onnx_path} before simplification.")
                return False
            if not self.dummy_input_shapes:
                 logging.error("Cannot simplify model: Dummy input shapes were not generated.")
                 return False # Treat as failure if shapes are missing

            logging.info(f"Providing input shapes to onnxsim: {self.dummy_input_shapes}")
            onnx_model = onnx.load(onnx_path)
            # Provide the input shapes derived from dummy data
            model_opt, check_ok = onnxsim.simplify(
                onnx_model,
                input_shapes=self.dummy_input_shapes, # Pass the shapes here
                check_n=3,
                perform_optimization=True
            )

            if check_ok:
                onnx.save(model_opt, onnx_path)
                logging.info("Simplified ONNX model saved successfully.")
                return True
            else:
                logging.error("ONNX simplification check failed. Keeping the "
                              "original (unsimplified) model.")
                return True # Treat simplification check failure as non-critical
        except Exception as e:
            # Log the specific error from onnxsim
            logging.error(f"Error during ONNX simplification: {e}", exc_info=True)
            logging.warning("Proceeding with the unsimplified ONNX model.")
            return True # Simplification process failed, but export is usable

    def export(self, model_path: str, exported_path: str) -> bool:
        """
        Executes the full export process.

        Args:
            model_path (str): Path to the input PyTorch model checkpoint (.pth).
            exported_path (str): Path to save the exported ONNX model.

        Returns:
            bool: True if base export was successful, False otherwise.
        """
        checkpoint = self._load_checkpoint(model_path)
        if not checkpoint: return False

        if not self._prepare_config(checkpoint): return False

        dummy_inputs = self._create_dummy_inputs() # Generates shapes needed later
        if not dummy_inputs: return False

        net_g = self._instantiate_and_load_model(checkpoint)
        if not net_g: return False

        dynamic_axes = {
            name: {1: "sequence_length"} for name in ["phone", "pitch", "pitchf"]
        }
        dynamic_axes["rnd"] = {2: "sequence_length"} # Dim 2 for rnd noise
        dynamic_axes[self.OUTPUT_NAMES[0]] = {1: "audio_length"} # Output audio

        try:
            logging.info(f"Exporting model to ONNX format at: {exported_path} "
                         f"(Opset: {self.opset_version})")
            torch.onnx.export(
                net_g, dummy_inputs, exported_path,
                input_names=self.INPUT_NAMES, output_names=self.OUTPUT_NAMES,
                dynamic_axes=dynamic_axes, do_constant_folding=False, # Be cautious with folding + dynamic axes
                opset_version=self.opset_version, verbose=False,
            )
            logging.info(f"Initial ONNX model exported successfully to: {exported_path}")

        except Exception as e:
            logging.error(f"ONNX export failed: {e}", exc_info=True)
            if os.path.exists(exported_path):
                try: os.remove(exported_path)
                except OSError as oe: logging.error(f"Could not remove intermediate ONNX file {exported_path}: {oe}")
            return False

        if self.simplify:
            # _simplify_model now uses self.dummy_input_shapes generated earlier
            self._simplify_model(exported_path)

        logging.info("ONNX export process finished.")
        return True


def get_model_path_from_user() -> str | None:
    """Prompts the user to provide a valid .pth model file path."""
    while True:
        # Use print for direct user interaction
        print("\nPlease provide the path to the RVC model checkpoint file (.pth).")
        print("You can type the full path or drag and drop the file onto "
              "this window and press Enter.")
        print("Type 'quit' or 'exit' to cancel.")
        raw_path = input("Model Path: ").strip()

        if raw_path.lower() in ['quit', 'exit']:
            logging.info("Operation cancelled by user.")
            return None

        cleaned_path = raw_path.strip("'\"")
        if not cleaned_path:
            print("ERROR: No path provided. Please try again or type 'quit'.")
            continue

        try: absolute_path = os.path.abspath(cleaned_path)
        except Exception as e: print(f"ERROR: Could not resolve path '{cleaned_path}'. Error: {e}"); continue

        if not os.path.exists(absolute_path): print(f"ERROR: File not found at '{absolute_path}'. Check the path."); continue
        if not os.path.isfile(absolute_path): print(f"ERROR: Path '{absolute_path}' is a directory. Need a file."); continue
        if not absolute_path.lower().endswith(".pth"): print(f"ERROR: File '{os.path.basename(absolute_path)}' needs .pth extension."); continue

        logging.info(f"Input model accepted: {absolute_path}")
        return absolute_path


def verify_onnx_model(onnx_path: str):
    """Attempts to load the ONNX model using onnxruntime for verification."""
    # Use print for separators in direct user feedback section
    print("-" * 30)
    logging.info(f"Attempting to verify ONNX model: {onnx_path}")

    if not ORT_AVAILABLE:
        logging.warning("ONNX Runtime (onnxruntime) library not found.")
        logging.warning("Verification skipped. Install with: pip install onnxruntime")
        print("-" * 30)
        return

    if not os.path.exists(onnx_path):
        logging.error(f"Verification failed: ONNX file not found at {onnx_path}")
        print("-" * 30)
        return

    try:
        providers = ['CPUExecutionProvider']
        ort_session = ort.InferenceSession(onnx_path, providers=providers)
        inputs = ort_session.get_inputs()
        outputs = ort_session.get_outputs()
        logging.info(f"ONNX Runtime loaded model successfully.")
        logging.info(f"Detected {len(inputs)} inputs: {[inp.name for inp in inputs]}")
        logging.info(f"Detected {len(outputs)} outputs: {[out.name for out in outputs]}")
        logging.info("Verification successful.")

    except ort.OrtLoadError as load_error:
        logging.error(f"ONNX Runtime failed to load the model: {load_error}", exc_info=True)
        logging.error("Verification failed: Model loading error.")
    except Exception as e:
        logging.error(f"An unexpected error occurred during ONNX verification: {e}", exc_info=True)
        logging.error("Verification failed.")
    finally:
        # Ensure separator is printed even on error during verification
        print("-" * 30)


def main():
    """Main execution function."""
    # Use print for initial user-facing messages
    print("-" * 30)
    print("RVC ONNX Exporter")
    print("-" * 30)
    logging.info("Exporter script started.") # Log that script has begun

    checkpoint_path = get_model_path_from_user()
    if not checkpoint_path:
        return 0 # User quit

    input_dir = os.path.dirname(checkpoint_path)
    base_name = os.path.splitext(os.path.basename(checkpoint_path))[0]
    onnx_output_path = os.path.join(input_dir, f"{base_name}.onnx")

    # Log essential paths
    print("-" * 30)
    logging.info(f"Input PyTorch Model: {checkpoint_path}")
    logging.info(f"Output ONNX Model:   {onnx_output_path}")
    print("-" * 30)

    # --- Initialize and Run Exporter ---
    exporter = RvcOnnxExporter(opset_version=18, simplify=True)
    export_successful = exporter.export(checkpoint_path, onnx_output_path)

    # --- Final Status & Optional Verification ---
    print("-" * 30) # Use print for visual separation for user
    if export_successful:
        logging.info(f"Successfully exported model to: {onnx_output_path}")

        while True:
            # Use print for direct user interaction prompt
            verify_choice = input("Verify the exported ONNX model using ONNX Runtime? (y/n): ").lower().strip()
            if verify_choice in ['y', 'yes']:
                verify_onnx_model(onnx_output_path)
                break
            elif verify_choice in ['n', 'no']:
                logging.info("Skipping ONNX model verification.")
                print("-" * 30) # Print separator after skipping
                break
            else:
                print("Invalid input. Please enter 'y' or 'n'.") # Use print for error feedback
    else:
        logging.error("Model export failed. See log messages above for details.")

    logging.info("Exporter script finished.")
    return 0 if export_successful else 1


if __name__ == "__main__":
    exit_code = 1 # Default to error
    try:
        # Start logging session info
        logging.info(f"--- Starting new conversion run ---")
        logging.info(f"Python version: {sys.version}")
        logging.info(f"Torch version: {torch.__version__}")
        logging.info(f"ONNX version: {onnx.__version__}")
        logging.info(f"ONNXSim version: {onnxsim.__version__}")
        logging.info(f"ONNX Runtime Available: {ORT_AVAILABLE}" + (f" (Version: {ort.__version__})" if ORT_AVAILABLE else ""))

        exit_code = main()

    except Exception as e:
        logging.error(f"An unexpected critical error occurred in main execution: {e}", exc_info=True)
        # Ensure the error is logged before the exit prompt

    finally:
        # This block ensures the prompt appears even if main() raises an exception
        logging.info(f"--- Ending conversion run (Exit Code: {exit_code}) ---")
        print("\nScript execution complete.") # Use print for final user message
        input("Press Enter to exit...") # Keep window open
        sys.exit(exit_code)