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:
        lines = f.readlines()
    
    new_lines = []
    # Make sure future annotations is ON to avoid runtime errors with piping if any remain, 
    # BUT wait, if I remove unions from Mapped, I don't need it as much for the scanner.
    # However, keeping it on is usually better for modern code.
    found_future = False
    
    for line in lines:
        if "from __future__ import annotations" in line:
            new_lines.append("from __future__ import annotations\n")
            found_future = True
            continue
        
        # Clean up existing double hashes if any
        line = line.replace("# # from __future__", "# from __future__")
        
        # Match Mapped[T | None] or Mapped[Optional[T]]
        # We want to transform to Mapped[T]
        # and ensure mapped_column has nullable=True
        
        match = re.search(r'(.*?):\s*Mapped\[(?:Optional\[(.*?)\|?(.*?)\]|([^\|]*?)\s*\|\s*None)\]\s*=\s*mapped_column\((.*?)\)(.*)', line)
        if match:
            # We found one.
            indent = match.group(1)
            # Try to get T from various groups
            t = match.group(2) or match.group(4)
            if not t and match.group(3): # For Optional[T] if it captured wrong
                 t = match.group(2)
            
            column_args = match.group(5)
            rest = match.group(6)
            
            # Ensure nullable=True
            if "nullable=" not in column_args:
                if column_args.strip():
                    column_args = f"{column_args}, nullable=True"
                else:
                    column_args = "nullable=True"
            elif "nullable=False" in column_args:
                # Should not happen as we found a Union with None, but if so, override
                column_args = column_args.replace("nullable=False", "nullable=True")
            
            new_line = f"{indent}: Mapped[{t}] = mapped_column({column_args}){rest}\n"
            new_lines.append(new_line)
        else:
            new_lines.append(line)
            
    if not found_future:
        new_lines.insert(0, "from __future__ import annotations\n")
        
    with open(filepath, 'w', encoding='utf-8') as f:
        f.writelines(new_lines)

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))