import os
import zipfile
import shutil
import re

# Source and Destination
SOURCE_DIR = 'sql_tools_zip'
DEST_DIR = 'sql_tools'

# Mapping structure: (Zip Pattern, Target Folder Structure, Executables to Extract)
# Target Folder Structure: e.g. "mysql/{version}" where {version} is extracted from filename
RULES = [
    {
        'pattern': r'mysql-(\d+\.\d+)\.\d+-winx64\.zip',
        'type': 'mysql',
        'bin_path': 'bin/', # usually inside root folder of zip
        'executables': ['mysqldump.exe']
    },
    {
        'pattern': r'mariadb-(\d+\.\d+)\.\d+-winx64.*\.zip',
        'type': 'mariadb',
        'bin_path': 'bin/',
        'executables': ['mariadb-dump.exe', 'mysqldump.exe'] # Try both
    },
    {
        'pattern': r'postgresql-(\d+\.\d+).*windows-x64-binaries\.zip',
        'type': 'postgres',
        'bin_path': 'pgsql/bin/', # pg binaries zip usually has pgsql/bin
        'executables': ['pg_dump.exe', 'pg_dumpall.exe', 'libpq.dll', 'libiconv-2.dll', 'libintl-8.dll', 'libssl-1_1-x64.dll', 'libcrypto-1_1-x64.dll'] # DLLs might be needed!
        # Modern PG zip structure might differ, checking pgsql/bin is standard for binaries zip
    },
    {
         # Special handler for PG 18 if pattern differs or standard regex fails
         'pattern': r'postgresql-(\d+)\.\d+.*windows-x64-binaries\.zip',
         'type': 'postgres',
         'bin_path': 'pgsql/bin/',
         'executables': ['pg_dump.exe', 'pg_dumpall.exe']
    }
]

def ensure_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)

def extract_tools():
    if not os.path.exists(SOURCE_DIR):
        print("Source directory '{}' not found.".format(SOURCE_DIR))
        return

    ensure_dir(DEST_DIR)

    files = [f for f in os.listdir(SOURCE_DIR) if f.lower().endswith('.zip')]
    
    print("Found {} zip files.".format(len(files)))

    for filename in files:
        # Match rule
        matched_rule = None
        version = None
        
        for rule in RULES:
            match = re.search(rule['pattern'], filename)
            if match:
                matched_rule = rule
                version = match.group(1)
                break
        
        if not matched_rule:
            print("Skipping unknown file: {}".format(filename))
            continue

        db_type = matched_rule['type']
        target_dir = os.path.join(DEST_DIR, db_type, version)
        
        print("Processing {} -> {}...".format(filename, target_dir))
        
        try:
            with zipfile.ZipFile(os.path.join(SOURCE_DIR, filename), 'r') as zf:
                # Find the internal path prefix (usually top folder)
                # We assume binaries are in {top_folder}/{bin_path} or just {bin_path} 
                # Let's search for the executable to determine path
                
                for exe in matched_rule['executables']:
                    # find file in zip
                    target_member = None
                    for name in zf.namelist():
                        if name.endswith("/{}".format(exe)) or name == exe:
                            # Verify if it looks like a bin folder
                            if matched_rule['bin_path'].strip('/') in name:
                                target_member = name
                                break
                    
                    if target_member:
                        ensure_dir(target_dir)
                        source = zf.open(target_member)
                        target_file_path = os.path.join(target_dir, exe)
                        with open(target_file_path, "wb") as target:
                            shutil.copyfileobj(source, target)
                        print("  Extracted: {}".format(exe))
                        
                        # For Postgres, we might need DLLs. 
                        # Simple logic: if 'postgres' in db_type, try to grab standard DLLs in same dir as exe
                        if db_type == 'postgres':
                            parent = os.path.dirname(target_member)
                            for item in zf.namelist():
                                if os.path.dirname(item) == parent and item.lower().endswith('.dll'):
                                    dll_name = os.path.basename(item)
                                    with open(os.path.join(target_dir, dll_name), "wb") as dll_out:
                                        shutil.copyfileobj(zf.open(item), dll_out)

                    else:
                        print("  Warning: {} not found in {}".format(exe, filename))

        except Exception as e:
            print("  Error extracting {}: {}".format(filename, e))

if __name__ == "__main__":
    extract_tools()