#!/usr/bin/env python3
"""
Fix all CREATE TABLE statements - add PRIMARY KEY for AUTO_INCREMENT columns
Simple, robust line-by-line approach
"""

import re

def fix_all_tables(input_file, output_file):
    with open(input_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    
    # Step 1: Build PK map from ALTER TABLE
    print("Building PRIMARY KEY map...")
    pk_map = {}
    i = 0
    while i < len(lines):
        if 'ALTER TABLE' in lines[i]:
            # Extract table name
            table_match = re.search(r'ALTER TABLE `?([^\s`]+)`?', lines[i], re.IGNORECASE)
            if table_match:
                table_name = table_match.group(1).strip('`').strip().lower()
                # Look ahead for PRIMARY KEY (may be on next line)
                j = i
                while j < len(lines) and j < i + 3:  # Check up to 3 lines ahead
                    if 'PRIMARY KEY' in lines[j]:
                        pk_match = re.search(r'PRIMARY KEY\s*\(([^)]+)\)', lines[j], re.IGNORECASE)
                        if pk_match:
                            pk_cols = [col.strip('`').strip() for col in pk_match.group(1).split(',')]
                            pk_map[table_name] = pk_cols
                            print(f"  Found PK for {table_name}: {pk_cols}")
                            break
                    if lines[j].strip().endswith(';'):
                        break
                    j += 1
        i += 1
    
    print(f"Found {len(pk_map)} PRIMARY KEY definitions\n")
    
    # Step 2: Fix CREATE TABLE statements
    print("Fixing CREATE TABLE statements...")
    output_lines = []
    i = 0
    tables_fixed = 0
    
    while i < len(lines):
        line = lines[i]
        output_lines.append(line)
        
        # Check for CREATE TABLE
        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_idx = -1
            
            # Collect table definition
            while j < len(lines):
                next_line = lines[j]
                table_lines.append(next_line)
                
                # Check for AUTO_INCREMENT
                if 'AUTO_INCREMENT' in next_line:
                    found_auto_inc = True
                    col_match = re.search(r'`([^`]+)`', next_line)
                    if col_match:
                        auto_inc_col = col_match.group(1).strip()
                
                # Track last column
                if ('`' in next_line and 
                    ');' not in next_line and 
                    'PRIMARY KEY' not in next_line.upper() and
                    'ENGINE' not in next_line.upper()):
                    last_col_idx = len(table_lines) - 1
                
                # Check for closing );
                if ');' in next_line:
                    # Fix if needed
                    if (found_auto_inc and 
                        'PRIMARY KEY' not in ''.join(table_lines).upper() and
                        table_name in pk_map):
                        pk_cols = pk_map[table_name]
                        if auto_inc_col and auto_inc_col in pk_cols:
                            # Add comma to last column
                            if last_col_idx >= 0:
                                last_col = table_lines[last_col_idx]
                                if not last_col.rstrip().endswith(',') and not last_col.rstrip().endswith('('):
                                    table_lines[last_col_idx] = last_col.rstrip() + ',\n'
                            
                            # Replace ); with PRIMARY KEY + );
                            pk_str = ', '.join([f'`{c}`' for c in pk_cols])
                            table_lines[-1] = f'    PRIMARY KEY ({pk_str})\n);\n'
                            tables_fixed += 1
                            print(f"  ✓ Fixed {table_name}")
                    
                    # Write table lines (skip first, already written)
                    output_lines.extend(table_lines[1:])
                    i = j
                    break
                j += 1
        i += 1
    
    # Step 3: Comment out duplicate ALTER TABLE PRIMARY KEY
    print(f"\nCommenting out duplicate ALTER TABLE statements...")
    final_output = []
    content_str = ''.join(output_lines)
    commented_count = 0
    
    i = 0
    while i < len(output_lines):
        line = output_lines[i]
        
        # Check for ALTER TABLE PRIMARY KEY
        if 'ALTER TABLE' in line:
            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
                create_pattern = rf'CREATE TABLE `?{re.escape(table_name)}`?\s*\([^)]+PRIMARY KEY[^)]+\);'
                if re.search(create_pattern, content_str, re.IGNORECASE | re.DOTALL):
                    # Comment out this line and the next line if it has PRIMARY KEY
                    line = '-- ' + line.rstrip() + ' -- Removed: PRIMARY KEY already in CREATE TABLE\n'
                    if i+1 < len(output_lines) and 'PRIMARY KEY' in output_lines[i+1]:
                        output_lines[i+1] = '-- ' + output_lines[i+1].rstrip() + '\n'
                    commented_count += 1
        
        final_output.append(line)
        i += 1
    
    print(f"  Commented out {commented_count} ALTER TABLE PRIMARY KEY statements")
    
    with open(output_file, 'w', encoding='utf-8') as f:
        f.writelines(final_output)
    
    print(f"\n✅ Fixed {tables_fixed} CREATE TABLE statements")

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