!pip install onnx

import json
import os
import onnx
import urllib.request
from urllib.error import URLError, HTTPError
from google.colab import files

def download_file(url, save_path):
    """
    Helper function to download a file from a URL.
    Returns True if successful, False otherwise.
    """
    if not url:
        return False
        
    try:
        print(f"Downloading from {url} ...")
        urllib.request.urlretrieve(url, save_path)
        print(f"Successfully downloaded to Colab: {save_path}")
        return True
    except HTTPError as e:
        print(f"HTTP Error failed to download {url}: {e.code}")
        return False
    except URLError as e:
        print(f"URL Error failed to download {url}: {e.reason}")
        return False
    except Exception as e:
        print(f"An unexpected error occurred during download: {e}")
        return False

def process_piper_model_from_url(onnx_url, json_url, output_dir):
    # Guard Clauses: Verify if required inputs are provided
    if not onnx_url or not json_url:
        print("Error: Both ONNX URL and JSON URL must be provided.")
        return

    # Create output directory if it does not exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Extract filenames from URLs or assign default names
    onnx_filename = onnx_url.split('/')[-1] if '/' in onnx_url else "model.onnx"
    json_filename = json_url.split('/')[-1] if '/' in json_url else "config.json"
    
    # Define local file paths
    onnx_path = os.path.join(output_dir, onnx_filename)
    json_path = os.path.join(output_dir, json_filename)
    tokens_path = os.path.join(output_dir, "tokens.txt")

    # Download required files to Colab environment
    if not download_file(onnx_url, onnx_path):
        print("Failed to download ONNX file. Aborting process.")
        return
        
    if not download_file(json_url, json_path):
        print("Failed to download JSON config. Aborting process.")
        return

    print("Files fetched successfully. Starting processing...")

    # Read JSON Configuration File
    try:
        with open(json_path, "r", encoding="utf-8") as f:
            config = json.load(f)
    except Exception as e:
        print(f"Error reading JSON file: {e}")
        return

    # Guard Clause: Verify required keys exist in JSON
    if "language" not in config or "espeak" not in config or "audio" not in config:
        print("Error: JSON config is missing required Piper metadata keys.")
        return

    # Step 1: Embed/Add Metadata into the ONNX File
    print("Step 1: Adding metadata to the ONNX file...")
    try:
        model = onnx.load(onnx_path)
    except Exception as e:
        print(f"Error loading ONNX model: {e}")
        return

    meta_data = {
        "model_type": "vits",
        "comment": "piper",
        "language": config["language"]["code"],
        "voice": config["espeak"]["voice"],
        "has_espeak": 1,
        "n_speakers": config["num_speakers"],
        "sample_rate": config["audio"]["sample_rate"],
    }

    for key, value in meta_data.items():
        meta = model.metadata_props.add()
        meta.key = key
        meta.value = str(value)

    try:
        onnx.save(model, onnx_path)
        print(f"-> Metadata successfully saved into '{onnx_filename}'!")
    except Exception as e:
        print(f"Error saving modified ONNX model: {e}")
        return

    # Step 2: Generate tokens.txt (Phoneme Map) file for Sherpa-ONNX
    print("Step 2: Generating tokens.txt file for Sherpa-ONNX...")
    if "phoneme_id_map" not in config:
        print("Error: 'phoneme_id_map' not found in JSON config.")
        return
        
    id_map = config["phoneme_id_map"]

    try:
        with open(tokens_path, "w", encoding="utf-8") as f_tokens:
            for s, i in id_map.items():
                f_tokens.write(f"{s} {i[0]}\n")
        
        print(f"-> tokens.txt successfully generated! Total tokens: {len(id_map)}")
    except Exception as e:
        print(f"Error writing tokens.txt: {e}")
        return

    # Step 3: Trigger Automatic Browser Downloads
    print("\n[SUCCESS] Conversion complete! Starting automatic browser downloads...")
    try:
        print(f"Downloading {onnx_filename} to your device...")
        files.download(onnx_path)
        
        print("Downloading tokens.txt to your device...")
        files.download(tokens_path)
    except Exception as e:
        print(f"Browser download failed or not running in Colab environment: {e}")

# ==========================================
# MAIN LOGIC - Users only need to edit below
# ==========================================

ONNX_FILE_URL = "YOUR_ONNX_FILE_URL_HERE"
JSON_FILE_URL = "YOUR_JSON_FILE_URL_HERE"
OUTPUT_DIRECTORY = "./sherpa_piper_model"

# Execute the function directly
process_piper_model_from_url(ONNX_FILE_URL, JSON_FILE_URL, OUTPUT_DIRECTORY)