"""Patch ddsp training imports not needed for inference."""
import os
import site

site_packages = site.getsitepackages()[0]
ddsp_dir = os.path.join(site_packages, 'ddsp')

# Modules to make optional
optional_modules = [
    'tensorflow_datasets',
    'tensorflow_addons',
    'google.cloud',
    'apache_beam',
]

for root, dirs, files in os.walk(ddsp_dir):
    for fname in files:
        if not fname.endswith('.py'):
            continue
        fpath = os.path.join(root, fname)
        with open(fpath, 'r', errors='ignore') as f:
            content = f.read()

        modified = False
        for mod in optional_modules:
            # Match both "import X" and "from X import Y" style
            for pattern in [f'import {mod}', f'from {mod}']:
                lines = content.split('\n')
                new_lines = []
                i = 0
                while i < len(lines):
                    line = lines[i]
                    stripped = line.lstrip()
                    if stripped.startswith(pattern) and not stripped.startswith('try:'):
                        indent = line[:len(line) - len(stripped)]
                        new_lines.append(f'{indent}try:')
                        new_lines.append(f'{indent}    {stripped}')
                        new_lines.append(f'{indent}except (ImportError, Exception):')
                        new_lines.append(f'{indent}    pass')
                        modified = True
                    else:
                        new_lines.append(line)
                    i += 1
                content = '\n'.join(new_lines)

        if modified:
            with open(fpath, 'w') as f:
                f.write(content)
            print(f'Patched: {fpath}')

print('Done.')