#!/usr/bin/env python3
"""
Export Header Fix Script v7
Fixes wrong export header includes in Qt6 migrated code.

Improvements in v7:
- Better directory-based mapping for headers not explicitly in CMakeLists.txt
- Handles subdirectory structures like Lts/Common/LtsWidgets/
- Handles test executables (remove export headers)
"""

import os
import re
import sys
from pathlib import Path

def parse_cmakelists(cmakelists_path):
    """Parse CMakeLists.txt and extract library info with their source files"""
    libraries = []
    
    with open(cmakelists_path, 'r', encoding='utf-8', errors='ignore') as f:
        content = f.read()
    
    # Remove comments
    content = re.sub(r'#.*', '', content)
    
    # Find all SET() variables
    set_vars = {}
    for match in re.finditer(r'SET\s*\(\s*(\w+)\s+([^)]+)\)', content, re.IGNORECASE | re.DOTALL):
        var_name = match.group(1)
        var_value = match.group(2).strip()
        # Split by whitespace or newlines
        files = [f.strip() for f in re.split(r'\s+', var_value) if f.strip()]
        set_vars[var_name] = files
    
    # Find NGL_ADD_LIBRARY and NGL_ADD_LIBRARY2 calls
    for match in re.finditer(r'NGL_ADD_LIBRARY2?\s*\(([^)]+)\)', content, re.IGNORECASE | re.DOTALL):
        args = match.group(1).strip()
        tokens = args.split()
        
        if len(tokens) < 3:
            continue
        
        # Find TYPE keyword (SHARED, STATIC, MODULE, OBJECT)
        type_idx = -1
        for i, token in enumerate(tokens):
            if token.upper() in ['SHARED', 'STATIC', 'MODULE', 'OBJECT']:
                type_idx = i
                break
        
        if type_idx == -1:
            continue
        
        lib_type = tokens[type_idx].upper()
        
        # Target name is right before TYPE
        if type_idx > 0:
            target_name = tokens[type_idx - 1]
            if target_name.upper() in ['SHARED', 'STATIC', 'MODULE', 'OBJECT']:
                continue
        else:
            continue
        
        # Get source files (everything after TYPE)
        src_tokens = tokens[type_idx + 1:]
        all_files = []
        for token in src_tokens:
            if token.startswith('${'):
                var_name = token[2:-1]
                if var_name in set_vars:
                    all_files.extend(set_vars[var_name])
            else:
                all_files.append(token)
        
        libraries.append({
            'name': target_name,
            'type': lib_type,
            'files': all_files,
            'cmakelists': cmakelists_path
        })
    
    # Also find ADD_EXECUTABLE calls (for test files)
    for match in re.finditer(r'ADD_EXECUTABLE\s*\(\s*(\w+)\s+([^)]+)\)', content, re.IGNORECASE | re.DOTALL):
        exe_name = match.group(1)
        files_str = match.group(2).strip()
        files = [f.strip() for f in re.split(r'\s+', files_str) if f.strip()]
        
        libraries.append({
            'name': exe_name,
            'type': 'EXECUTABLE',
            'files': files,
            'cmakelists': cmakelists_path
        })
    
    return libraries

def find_all_cmakelists(base_dir):
    """Find all CMakeLists.txt files"""
    return list(base_dir.rglob('CMakeLists.txt'))

def map_files_to_targets(base_dir):
    """Map header files to their targets"""
    file_to_target = {}
    cmakelists_files = find_all_cmakelists(base_dir)
    
    print(f"Found {len(cmakelists_files)} CMakeLists.txt files")
    
    # First pass: collect all targets and their directories
    all_targets = []
    for cmakelists in cmakelists_files:
        libraries = parse_cmakelists(cmakelists)
        cmakelists_dir = cmakelists.parent
        
        for lib in libraries:
            lib['cmakelists_dir'] = cmakelists_dir
            all_targets.append(lib)
    
    # Second pass: map files to targets using multiple strategies
    
    # Strategy 1: Direct file reference in CMakeLists.txt
    for lib in all_targets:
        cmakelists_dir = lib['cmakelists_dir']
        
        for src_file in lib['files']:
            src_filename = Path(src_file).name
            
            # Try direct path first
            src_path = cmakelists_dir / src_file
            if src_path.exists():
                file_to_target[str(src_path)] = (lib['name'], lib['type'])
            
            # Search in same directory tree
            for header_file in cmakelists_dir.rglob(src_filename):
                if header_file.suffix in ['.h', '.hpp']:
                    file_to_target[str(header_file)] = (lib['name'], lib['type'])
    
    # Strategy 2: Directory-based mapping
    # If a header is in a directory with a CMakeLists.txt that defines a library,
    # assume it belongs to that library
    for lib in all_targets:
        cmakelists_dir = lib['cmakelists_dir']
        
        # Map all .h/.hpp files in the CMakeLists directory and subdirectories
        for header_file in cmakelists_dir.rglob('*.h'):
            if str(header_file) not in file_to_target:
                file_to_target[str(header_file)] = (lib['name'], lib['type'])
        
        for header_file in cmakelists_dir.rglob('*.hpp'):
            if str(header_file) not in file_to_target:
                file_to_target[str(header_file)] = (lib['name'], lib['type'])
    
    # Strategy 3: Global include/ directory mapping by plugin name
    global_include = base_dir / 'include'
    if global_include.exists():
        for lib in all_targets:
            cmakelists_dir = lib['cmakelists_dir']
            plugin_name = cmakelists_dir.name
            
            target_include_dir = global_include / plugin_name
            if target_include_dir.exists():
                for header_file in target_include_dir.rglob('*.h'):
                    if str(header_file) not in file_to_target:
                        file_to_target[str(header_file)] = (lib['name'], lib['type'])
                for header_file in target_include_dir.rglob('*.hpp'):
                    if str(header_file) not in file_to_target:
                        file_to_target[str(header_file)] = (lib['name'], lib['type'])
    
    # Filter out external library directories
    external_lib_dirs = ['include/rti', 'include/ArcGIS', 'include/GeographicLib']
    files_to_remove = []
    for file_path in file_to_target.keys():
        for ext_dir in external_lib_dirs:
            if ext_dir in file_path:
                files_to_remove.append(file_path)
                break
    
    for file_path in files_to_remove:
        del file_to_target[file_path]
    
    print(f"Mapped {len(file_to_target)} files to targets (after filtering external libs)")
    return file_to_target

def get_export_header_name(target_name):
    """Generate export header name from target name"""
    return f"{target_name.lower()}_export.h"

def fix_export_header(file_path, target_name, lib_type):
    """Fix export header include in a file"""
    export_header = get_export_header_name(target_name)
    
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Find existing export header includes
    patterns = [
        r'#include\s*<\w+_export\.h>',
        r'#include\s*"\w+_export\.h"',
        r'#include\s*<\w+/\w+_export\.h>',
        r'#include\s*"\w+/\w+_export\.h"'
    ]
    
    found_export = False
    for pattern in patterns:
        if re.search(pattern, content):
            found_export = True
            break
    
    if not found_export:
        return False, "No export header found"
    
    # Replace with correct export header
    new_content = re.sub(
        r'#include\s*[<"]\w+/?\w*_export\.h[>"]',
        f'#include "{export_header}"',
        content
    )
    
    if new_content == content:
        return False, "No changes made"
    
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(new_content)
    
    return True, f"Fixed to {export_header}"

def remove_export_header(file_path):
    """Remove export header include from STATIC library or executable file"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    patterns = [
        r'#include\s*[<"]\w+/?\w*_export\.h[>"]\n'
    ]
    
    original = content
    for pattern in patterns:
        content = re.sub(pattern, '', content)
    
    if content == original:
        return False, "No export header found to remove"
    
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(content)
    
    return True, "Removed export header"

def main():
    if len(sys.argv) < 2:
        print("Usage: fix_export_headers_v7.py <base_dir>")
        sys.exit(1)
    
    base_dir = Path(sys.argv[1]).resolve()
    
    # Map all files to targets
    file_to_target = map_files_to_targets(base_dir)
    
    # Find all header files
    header_files = list(base_dir.rglob('*.h')) + list(base_dir.rglob('*.hpp'))
    
    fix_count = 0
    remove_count = 0
    skip_count = 0
    error_count = 0
    error_files = []
    
    for header_path in header_files:
        header_str = str(header_path)
        
        # Skip build directories
        if '/build' in header_str or 'build2/' in header_str or 'build-' in header_str:
            skip_count += 1
            continue
        
        # Skip external code
        if 'grantlee/' in header_str.lower():
            skip_count += 1
            continue
        
        # Skip external library directories
        if '/include/rti/' in header_str or '/include/ArcGIS/' in header_str or '/include/GeographicLib/' in header_str:
            skip_count += 1
            continue
        
        # Check if file has export header include FIRST
        with open(header_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        has_export = bool(re.search(r'#include\s*[<"].*_export\.h[>"]', content))
        
        if not has_export:
            skip_count += 1
            continue
        
        # Now check if file is mapped to a target
        if header_str not in file_to_target:
            error_count += 1
            error_files.append((header_path, "Could not determine target"))
            continue
        
        target_name, lib_type = file_to_target[header_str]
        
        # Handle based on library type
        if lib_type == 'STATIC' or lib_type == 'EXECUTABLE':
            success, msg = remove_export_header(header_path)
            if success:
                remove_count += 1
                print(f"REMOVED ({lib_type}): {header_path} - {msg}")
            else:
                skip_count += 1
        else:
            # SHARED or MODULE
            success, msg = fix_export_header(header_path, target_name, lib_type)
            if success:
                fix_count += 1
                print(f"FIXED: {header_path} -> {msg}")
            else:
                skip_count += 1
    
    print("\n" + "="*80)
    print("Summary:")
    print(f"  Fixed:   {fix_count}")
    print(f"  Removed (STATIC/EXEC): {remove_count}")
    print(f"  Skipped: {skip_count}")
    print(f"  Errors:  {error_count}")
    print("="*80)
    
    if error_files and error_count > 0:
        print("\nErrors:")
        for path, reason in error_files[:30]:
            print(f"  ERROR: {path} - {reason}")
        if error_count > 30:
            print(f"  ... and {error_count - 30} more")

if __name__ == '__main__':
    main()