"""
Generate clean requirements.txt from conda environment for Hugging Face deployment
"""
import subprocess
import sys
import re

def get_conda_packages():
    """Get packages from conda environment"""
    print("Getting packages from conda environment 'momask'...")
    try:
        result = subprocess.run(
            ['conda', 'list', '-n', 'momask', '--export'],
            capture_output=True,
            text=True,
            check=True
        )
        return result.stdout
    except subprocess.CalledProcessError as e:
        print(f"Error getting conda packages: {e}")
        return None

def convert_to_pip_format(conda_output):
    """Convert conda export to pip requirements format"""
    requirements = []

    for line in conda_output.split('\n'):
        line = line.strip()
        # Skip comments and empty lines
        if not line or line.startswith('#'):
            continue

        # Parse conda format: package=version=build
        parts = line.split('=')
        if len(parts) >= 2:
            package = parts[0]
            version = parts[1]

            # Skip pip itself and conda-specific packages
            if package in ['pip', 'conda', 'python']:
                continue

            # Handle special cases
            if package == 'pytorch':
                package = 'torch'
            elif package == 'pytorch-cuda':
                continue  # Skip CUDA packages

            requirements.append(f"{package}=={version}")

    return requirements

def create_minimal_requirements():
    """Create minimal requirements.txt with essential packages"""
    print("\nCreating minimal requirements.txt with essential packages...")

    essential_packages = [
        "torch>=2.0.0",
        "torchvision>=0.15.0",
        "torchaudio>=2.0.0",
        "numpy==1.23.5",
        "gradio>=4.0.0",
        "einops>=0.6.1",
        "scipy>=1.13.0",
        "matplotlib>=3.5.0",
        "trimesh>=4.0.0",
        "chumpy>=0.70",
        "scikit-learn>=1.6.0",
        "tqdm>=4.67.0",
        "pandas>=2.0.0",
        "Pillow>=9.0.0",
        "ffmpeg-python>=0.2.0",
        "vector-quantize-pytorch>=1.6.0",
        "huggingface-hub>=1.0.0",
        "git+https://github.com/openai/CLIP.git",
    ]

    return essential_packages

def main():
    print("=" * 70)
    print(" " * 20 + "Requirements Generator")
    print("=" * 70)

    # Try to get packages from conda
    conda_output = get_conda_packages()

    if conda_output:
        print("✓ Successfully retrieved conda packages")
        requirements = convert_to_pip_format(conda_output)
        print(f"✓ Found {len(requirements)} packages")
    else:
        print("⚠ Could not retrieve conda packages, using minimal requirements")
        requirements = create_minimal_requirements()

    # Always ensure essential packages are included
    essential_packages = create_minimal_requirements()

    # Write requirements.txt
    output_file = 'requirements_clean.txt'
    print(f"\nWriting to {output_file}...")

    with open(output_file, 'w', encoding='utf-8') as f:
        # Write header
        f.write("# Requirements for Hugging Face Spaces deployment\n")
        f.write("# Generated for MoMask Text-to-Motion Generator\n\n")

        # Write essential packages
        f.write("# Essential packages\n")
        for pkg in essential_packages:
            f.write(f"{pkg}\n")

    print(f"✓ Requirements written to {output_file}")
    print("\nNext steps:")
    print("1. Review requirements_clean.txt")
    print("2. Rename it to requirements.txt")
    print("3. Test locally: pip install -r requirements.txt")
    print("4. Run pre_deploy_check.py before deploying")

    print("\n" + "=" * 70)

if __name__ == "__main__":
    main()