#!/usr/bin/env python3
import os
import sys
import subprocess
import shutil
from pathlib import Path

# Configuration
REPO_URL = "https://github.com/hpbyte/myanmar-tts.git"
REPO_DIR = "myanmar-tts"
MODEL_DIR = "trained_model"

def setup_repository():
    """Clone and set up the Myanmar TTS repository."""
    print(f"Setting up Myanmar TTS repository...")
    
    # Clone the repository if it doesn't exist
    if not os.path.exists(REPO_DIR):
        print(f"Cloning repository from {REPO_URL}...")
        subprocess.run(["git", "clone", REPO_URL], check=True)
        print("Repository cloned successfully!")
    else:
        print(f"Repository already exists at {REPO_DIR}")
    
    # Create the model directory if it doesn't exist
    if not os.path.exists(MODEL_DIR):
        os.makedirs(MODEL_DIR)
        print(f"Created model directory at {MODEL_DIR}")
    else:
        print(f"Model directory already exists at {MODEL_DIR}")
    
    # Create placeholder files with instructions if they don't exist
    checkpoint_path = os.path.join(MODEL_DIR, "checkpoint_latest.pth.tar")
    config_path = os.path.join(MODEL_DIR, "hparams.yml")
    
    if not os.path.exists(checkpoint_path):
        with open(checkpoint_path, 'w') as f:
            f.write("This is a placeholder file. Replace with the actual model checkpoint.")
        print(f"Created placeholder file at {checkpoint_path}")
    
    if not os.path.exists(config_path):
        with open(config_path, 'w') as f:
            f.write("# This is a placeholder file. Replace with the actual hparams.yml file.")
        print(f"Created placeholder file at {config_path}")
    
    # Make sure the repository packages are installed
    try:
        # Create a minimal setup.py in the repository if it doesn't exist
        setup_path = os.path.join(REPO_DIR, "setup.py")
        if not os.path.exists(setup_path):
            with open(setup_path, 'w') as f:
                f.write("""
from setuptools import setup, find_packages

setup(
    name="myanmar_tts",
    version="0.1.0",
    packages=find_packages(),
    install_requires=[
        "inflect",
        "unidecode",
        "torch",
        "numpy",
        "scipy",
    ],
)
""")
            print("Created setup.py file in repository")
        
        # Install the package in development mode
        print("Installing Myanmar TTS package...")
        subprocess.run([sys.executable, "-m", "pip", "install", "-e", REPO_DIR], check=True)
        print("Myanmar TTS package installed successfully!")
        
        # Create __init__.py files to make imports work
        for subdir in ["text", "utils"]:
            init_path = os.path.join(REPO_DIR, subdir, "__init__.py")
            if not os.path.exists(init_path):
                os.makedirs(os.path.dirname(init_path), exist_ok=True)
                with open(init_path, 'w') as f:
                    f.write("# Auto-generated __init__.py file")
                print(f"Created {init_path}")
    
    except Exception as e:
        print(f"Error during package installation: {str(e)}")
        print("Continuing with setup...")
    
    print("Repository setup complete!")
    print("\nIMPORTANT: You need to upload the following files to the 'trained_model' directory:")
    print("1. checkpoint_latest.pth.tar - The model checkpoint file")
    print("2. hparams.yml - The hyperparameters configuration file")
    print("\nThese files can be obtained from the original repository or by training the model.")

if __name__ == "__main__":
    setup_repository()