import os
import re

models_dir = r"c:\dev\git\python\homelab-automation-api-v2\app\models"

def fix_file(filepath):
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Re-enable future annotations correctly
    content = content.replace("# from __future__ import annotations", "from __future__ import annotations")
    content = content.replace("# # from __future__ import annotations", "from __future__ import annotations")
    
    # Match lines with Mapped[... | None] = mapped_column(...)
    # We want to change to Mapped[...] and add nullable=True
    
    def replacer(match):
        indent_and_name = match.group(1)
        type_name = match.group(2)
        column_args = match.group(3)
        rest = match.group(4)
        
        # Strip trailing comma in args if it was at the end
        args = column_args.strip()
        if args and "nullable=" not in args:
             if args.endswith(")"):
                  # It was likely something(args)
                  # We append , nullable=True at the end but INSIDE the mapped_column? 
                  # Wait, mapped_column(String, nullable=True)
                  pass # Handle below
        
        if "nullable=" not in args:
             if args:
                  args += ", nullable=True"
             else:
                  args = "nullable=True"
        
        return f"{indent_and_name}: Mapped[{type_name}] = mapped_column({args}){rest}"

    # Regex to find: some_field: Mapped[SOME_TYPE | None] = mapped_column(SOME_ARGS)
    pattern = r'(\s+[\w_]+)\s*:\s*Mapped\[\s*([^\|\[\]]+)\s*\|\s*None\s*\]\s*=\s*mapped_column\((.*?)\)(.*)'
    content = re.sub(pattern, replacer, content)
    
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(content)

for filename in os.listdir(models_dir):
    if filename.endswith(".py") and filename != "__init__.py":
        print(f"Fixing {filename}...")
        fix_file(os.path.join(models_dir, filename))