#!/usr/bin/env python3
"""
Comprehensive fix for MySQL schema - adds PRIMARY KEY to all CREATE TABLE with AUTO_INCREMENT
"""

import re

def fix_schema(input_file, output_file):
    with open(input_file, 'r') as f:
        lines = f.readlines()
    
    # Step 1: Build PK map from ALTER TABLE statements
    pk_map = {}
    for i, line in enumerate(lines):
        match = re.search(r'ALTER TABLE `?([^\s`]+)`?\s+ADD CONSTRAINT[^P]+PRIMARY KEY\s*\(([^)]+)\);', line, re.IGNORECASE)
        if match:
            table = match.group(1).strip('`').strip().lower()
            cols = [c.strip('`').strip() for c in match.group(2).split(',')]
            pk_map[table] = cols
    
    # Step 2: Process CREATE TABLE statements line by line
    output_lines = []
    i = 0
    
    while i < len(lines):
        line = lines[i]
        output_lines.append(line)
        
        # Check if this is a CREATE TABLE line
        create_match = re.search(r'CREATE TABLE `?([^\s`]+)`?', line, re.IGNORECASE)
        if create_match:
            table_name = create_match.group(1).strip('`').strip().lower()
            table_lines = [line]
            j = i + 1
            found_auto_inc = False
            auto_inc_col = None
            last_col_line_idx = -1
            
            # Collect CREATE TABLE lines until we find );
            while j < len(lines):
                next_line = lines[j]
                table_lines.append(next_line)
                
                # Check for AUTO_INCREMENT
                if re.search(r'AUTO_INCREMENT', next_line, re.IGNORECASE):
                    found_auto_inc = True
                    col_match = re.search(r'`([^`]+)`', next_line)
                    if col_match:
                        auto_inc_col = col_match.group(1).strip()
                
                # Track the last column definition line (before );
                if not re.search(r'\);', next_line) and not re.search(r'PRIMARY KEY', next_line, re.IGNORECASE):
                    if next_line.strip() and not next_line.strip().startswith('--'):
                        last_col_line_idx = len(table_lines) - 1
                
                # Check for closing );
                if re.search(r'\);', next_line):
                    # If we need to add PRIMARY KEY
                    if found_auto_inc and table_name in pk_map:
                        pk_cols = pk_map[table_name]
                        if auto_inc_col and auto_inc_col in pk_cols:
                            table_text = ''.join(table_lines)
                            if 'PRIMARY KEY' not in table_text.upper():
                                # Add comma to last column if needed
                                if last_col_line_idx >= 0:
                                    last_col = table_lines[last_col_line_idx]
                                    if not last_col.rstrip().endswith(',') and not last_col.rstrip().endswith('('):
                                        table_lines[last_col_line_idx] = last_col.rstrip() + ',\n'
                                
                                # Insert PRIMARY KEY before );
                                pk_str = ', '.join([f'`{c}`' for c in pk_cols])
                                # Replace the ); line
                                closing_line = table_lines[-1]
                                if closing_line.strip() == ');':
                                    table_lines[-1] = f'    PRIMARY KEY ({pk_str})\n);\n'
                                else:
                                    # ); is on same line
                                    table_lines[-1] = re.sub(r'\);', f',\n    PRIMARY KEY ({pk_str})\n);', closing_line)
                    
                    # Write the fixed table lines (skip first as already written)
                    output_lines.extend(table_lines[1:])
                    i = j
                    break
                j += 1
        i += 1
    
    # Step 3: Comment out ALTER TABLE PRIMARY KEY for tables that have it in CREATE TABLE
    final_output = []
    in_alter_pk = False
    for line in output_lines:
        if re.search(r'ALTER TABLE `?([^\s`]+)`?\s+ADD CONSTRAINT[^P]+PRIMARY KEY', line, re.IGNORECASE):
            table_match = re.search(r'ALTER TABLE `?([^\s`]+)`?', line, re.IGNORECASE)
            if table_match:
                table_name = table_match.group(1).strip('`').strip().lower()
                # Check if CREATE TABLE has PRIMARY KEY
                content_so_far = ''.join(final_output)
                create_pattern = rf'CREATE TABLE `?{re.escape(table_name)}`?\s*\([^)]+PRIMARY KEY[^)]+\);'
                if re.search(create_pattern, content_so_far, re.IGNORECASE | re.DOTALL):
                    line = '-- ' + line.rstrip() + ' -- Removed: PRIMARY KEY already in CREATE TABLE\n'
            in_alter_pk = True
        elif in_alter_pk and ('PRIMARY KEY' in line or line.strip().endswith(';')):
            if not line.strip().startswith('--'):
                line = '-- ' + line
            if line.strip().endswith(';'):
                in_alter_pk = False
        final_output.append(line)
    
    with open(output_file, 'w') as f:
        f.writelines(final_output)
    
    print(f"✅ Fixed all AUTO_INCREMENT tables in {output_file}")

if __name__ == "__main__":
    fix_schema("mysql/schema_mysql_20251113_185155.sql", "mysql/schema_mysql_20251113_185155.sql")