#!/usr/bin/env python3
"""
Fix all CREATE TABLE statements to include PRIMARY KEY for AUTO_INCREMENT columns
"""

import re

def fix_all_tables(input_file, output_file):
    with open(input_file, 'r') as f:
        content = f.read()
    
    # Step 1: Build map of table -> primary key columns from ALTER TABLE statements
    pk_map = {}
    for match in re.finditer(r'ALTER TABLE `?([^\s`]+)`?\s+ADD CONSTRAINT[^P]+PRIMARY KEY\s*\(([^)]+)\);', content, re.IGNORECASE):
        table = match.group(1).strip('`').strip().lower()
        cols = [c.strip('`').strip() for c in match.group(2).split(',')]
        pk_map[table] = cols
    
    # Step 2: Fix each CREATE TABLE statement
    def fix_create_table(match):
        create_stmt = match.group(0)
        table_match = re.search(r'CREATE TABLE `?([^\s`]+)`?', create_stmt, re.IGNORECASE)
        if not table_match:
            return create_stmt
        
        table_name = table_match.group(1).strip('`').strip().lower()
        
        # Check if table has AUTO_INCREMENT
        if 'AUTO_INCREMENT' not in create_stmt:
            return create_stmt
        
        # Check if PRIMARY KEY already in CREATE TABLE
        if 'PRIMARY KEY' in create_stmt.upper():
            return create_stmt
        
        # Find AUTO_INCREMENT column
        auto_match = re.search(r'`([^`]+)`\s+[^,\n)]+AUTO_INCREMENT', create_stmt, re.IGNORECASE)
        if not auto_match:
            return create_stmt
        
        auto_col = auto_match.group(1).strip()
        
        # Check if this table has a PRIMARY KEY definition
        if table_name not in pk_map:
            return create_stmt
        
        pk_cols = pk_map[table_name]
        if auto_col not in pk_cols:
            return create_stmt
        
        # Add PRIMARY KEY before closing );
        pk_str = ', '.join([f'`{c}`' for c in pk_cols])
        create_stmt = re.sub(r'(\s+\));', f',\n    PRIMARY KEY ({pk_str})\n\\1', create_stmt, count=1)
        
        return create_stmt
    
    # Apply fix to all CREATE TABLE statements
    content = re.sub(
        r'CREATE TABLE `?[^\s`]+`?\s*\([^)]+\)\s*;',
        fix_create_table,
        content,
        flags=re.IGNORECASE | re.DOTALL
    )
    
    # Step 3: Comment out ALTER TABLE PRIMARY KEY for tables that now have it in CREATE TABLE
    def comment_alter_pk(match):
        table_name = match.group(1).strip('`').strip().lower()
        # Check if CREATE TABLE has PRIMARY KEY
        create_pattern = rf'CREATE TABLE `?{re.escape(table_name)}`?\s*\([^)]+PRIMARY KEY[^)]+\);'
        if re.search(create_pattern, content, re.IGNORECASE | re.DOTALL):
            return '-- ' + match.group(0) + ' -- Removed: PRIMARY KEY already in CREATE TABLE'
        return match.group(0)
    
    content = re.sub(
        r'ALTER TABLE `?([^\s`]+)`?\s+ADD CONSTRAINT[^P]+PRIMARY KEY[^;]+;',
        comment_alter_pk,
        content,
        flags=re.IGNORECASE | re.MULTILINE
    )
    
    with open(output_file, 'w') as f:
        f.write(content)
    
    print(f"✅ Fixed all AUTO_INCREMENT tables in {output_file}")

if __name__ == "__main__":
    input_file = "mysql/schema_mysql_20251113_185155.sql"
    output_file = "mysql/schema_mysql_20251113_185155.sql"
    
    fix_all_tables(input_file, output_file)