import os
import sys

with open("api.py", "r") as f:
    content = f.read()

# Re-apply ATK_DATA_ROOT
import re
replacements = [
    (r'INPUTS_DIR = PyPath\("\./inputs"\)\.resolve\(\)', r'DATA_ROOT = PyPath(os.environ.get("ATK_DATA_ROOT", "/tmp/augmentoolkit")).resolve()\nINPUTS_DIR = (DATA_ROOT / "inputs").resolve()'),
    (r'OUTPUTS_DIR = PyPath\("\./outputs"\)\.resolve\(\)', r'OUTPUTS_DIR = (DATA_ROOT / "outputs").resolve()'),
    (r'GENERATION_DIR = PyPath\("\./generation"\)\.resolve\(\)', r'GENERATION_DIR = (DATA_ROOT / "generation").resolve()'),
    (r'CONFIGS_DIR = PyPath\("\./external_configs"\)\.resolve\(\)', r'CONFIGS_DIR = (DATA_ROOT / "external_configs").resolve()'),
    (r'LOGS_DIR = PyPath\("\./logs"\)\.resolve\(\)', r'LOGS_DIR = (DATA_ROOT / "logs").resolve()')
]

for old, new in replacements:
    content = re.sub(old, new, content)

# Replace .mkdir(exist_ok=True) with parents=True
content = content.replace("INPUTS_DIR.mkdir(exist_ok=True)", "INPUTS_DIR.mkdir(parents=True, exist_ok=True)")
content = content.replace("OUTPUTS_DIR.mkdir(exist_ok=True)", "OUTPUTS_DIR.mkdir(parents=True, exist_ok=True)")
content = content.replace("GENERATION_DIR.mkdir(exist_ok=True)", "GENERATION_DIR.mkdir(parents=True, exist_ok=True)")
content = content.replace("CONFIGS_DIR.mkdir(exist_ok=True)", "CONFIGS_DIR.mkdir(parents=True, exist_ok=True)")
content = content.replace("LOGS_DIR.mkdir(exist_ok=True)", "LOGS_DIR.mkdir(parents=True, exist_ok=True)")

with open("api.py", "w") as f:
    f.write(content)