import os
import re
from pathlib import Path
from PIL import Image
from tqdm import tqdm
import xml.etree.ElementTree as ET
from sqlalchemy import create_engine, Column, String, MetaData, Table
from sqlalchemy.orm import sessionmaker
from sqlalchemy.dialects.mysql import insert

#chatgpt kekw 

# Define SQLAlchemy model
metadata = MetaData()

def define_table(data_type):
    existing_table = metadata.tables.get(data_type)
    if existing_table is not None:
        return existing_table

    columns = [
        Column('id', String(length=255), primary_key=True),
        Column('str', String(length=255, collation='utf8mb4_unicode_ci')),
        Column('imagePath', String(length=255, collation='utf8mb4_unicode_ci'), nullable=True),
        Column('sortName', String(length=255, collation='utf8mb4_unicode_ci'), nullable=True),
        Column('category', String(length=255, collation='utf8mb4_unicode_ci'), nullable=True),
        Column('netOpenName', String(length=255, collation='utf8mb4_unicode_ci'), nullable=True),
        Column('rareType', String(length=255, collation='utf8mb4_unicode_ci'), nullable=True)
    ]

    return Table(
        data_type,
        metadata,
        *columns,
        mysql_charset='utf8mb4',
        mysql_collate='utf8mb4_unicode_ci'
    )

def parse_xml(file_path):
    tree = ET.parse(file_path)
    return tree.getroot()

def find_data(root, data_type):
    id_element = root.find('.//name/id')
    str_element = root.find('.//name/str')

    data = {'id': id_element.text, 'str': str_element.text}

    if data_type == "cozynet_chuni_static_accessory":
        data.update({
            'imagePath': root.find('.//image/path').text or "",
            'sortName': root.find('.//sortName').text or "",
            'category': root.find('.//category').text or "",
            'netOpenName': root.find('.//netOpenName/str').text or ""
        })

    elif data_type == "cozynet_chuni_static_nameplate":
        data.update({
            'imagePath': root.find('.//image/path').text or "",
            'sortName': root.find('.//sortName').text or "",
            'netOpenName': root.find('.//netOpenName/str').text or "",
        })

    elif data_type == "cozynet_chuni_static_systemvoice":
        data.update({
            'imagePath': root.find('.//image/path').text or "",
            'sortName': root.find('.//sortName').text or ""
        })

    elif data_type == "cozynet_chuni_static_trophies":
        data.update({
            'rareType': root.find('.//rareType').text or "",
            'netOpenName': root.find('.//netOpenName/str').text or ""
        })

    elif data_type == "cozynet_chuni_static_mapicon":
        data.update({
            'imagePath': root.find('.//image/path').text or "",
            'sortName': root.find('.//sortName').text or ""
        })

    return data

def insert_data_into_database(session, data_type, data):
    table = define_table(data_type)
    stmt = insert(table).values(data).on_duplicate_key_update(data)
    session.execute(stmt)

def process_directories(base_path, session, data_types, xml_file_names):
    for data_type, xml_file_name in zip(data_types, xml_file_names):
        base_directories = [base_path]
        # Find the XML files with a progress bar
        xml_files = list(find_xml_files(base_path, xml_file_name))
        for xml_path in tqdm(xml_files, desc=f'Processing {data_type}'):
            root = parse_xml(xml_path)
            data = find_data(root, data_type)
            if data:
                insert_data_into_database(session, data_type, data)

def find_xml_files(directory, xml_file_name):
    for dirpath, _, filenames in os.walk(directory):
        if xml_file_name in filenames:
            yield Path(dirpath) / xml_file_name

def find_dds_files(source_folders, file_pattern=None):
    dds_files = []
    for source_folder in source_folders:
        for root, _, files in os.walk(source_folder):
            for file in files:
                if file.endswith('.dds') and (file_pattern is None or file_pattern.match(file)):
                    dds_files.append(os.path.join(root, file))
    return dds_files

def convert_dds_to_png(dds_file_path, png_file_path, scale_percent=50):
    with Image.open(dds_file_path) as img:
        new_width = int(img.width * scale_percent / 100)
        new_height = int(img.height * scale_percent / 100)
        img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
        img.save(png_file_path, 'PNG')

def process_files(files, destination_folder, progress_prefix, scale_percent=100):
    files_to_convert = []

    for file_path in files:
        file_name = os.path.splitext(os.path.basename(file_path))[0] + '.png'
        png_file_path = os.path.join(destination_folder, file_name)
        if not os.path.exists(png_file_path):
            files_to_convert.append(file_path)

    if not files_to_convert:
        print(f"{progress_prefix}: Already converted!")
        return

    for file_path in tqdm(files_to_convert, desc=progress_prefix, unit='file'):
        file_name = os.path.splitext(os.path.basename(file_path))[0] + '.png'
        png_file_path = os.path.join(destination_folder, file_name)
        convert_dds_to_png(file_path, png_file_path, scale_percent)

def process_avatarAccessories(source_folders, destination_folder, progress_prefix):
    avatar_accessory_files = find_dds_files(source_folders)

    if not os.path.exists(destination_folder):
        os.makedirs(destination_folder)

    process_files(avatar_accessory_files, destination_folder,
                  progress_prefix, scale_percent=50)

def process_nameplates(source_folders, destination_folder, progress_prefix):
    nameplate_files = find_dds_files(source_folders)

    if not os.path.exists(destination_folder):
        os.makedirs(destination_folder)

    process_files(nameplate_files, destination_folder, progress_prefix)

def process_systemVoiceImages(source_folders, destination_folder, progress_prefix):
    nameplate_files = find_dds_files(source_folders)

    if not os.path.exists(destination_folder):
        os.makedirs(destination_folder)

    process_files(nameplate_files, destination_folder, progress_prefix)

def process_jacketArts(source_folders, destination_folder, progress_prefix):
    nameplate_files = find_dds_files(source_folders)

    if not os.path.exists(destination_folder):
        os.makedirs(destination_folder)

    process_files(nameplate_files, destination_folder, progress_prefix)

def process_mapIcons(source_folders, destination_folder, progress_prefix):
    # Find DDS files matching the pattern for map icons
    map_icon_files = find_dds_files(source_folders)

    # Create destination folder if it doesn't exist
    if not os.path.exists(destination_folder):
        os.makedirs(destination_folder)

    # Process and convert DDS files to PNGs
    process_files(map_icon_files, destination_folder, progress_prefix)
    
def process_partners(source_folders, destination_chunithm, destination_chusan):
    pattern_chunithm = re.compile(
        r'CHU_UI_Character_0([0-9]{3})_(00|01)_[0-9]{2}\.dds$')

    pattern_chusan = re.compile(
        r'CHU_UI_Character_([1-9]\d{3,})_(00|01)_[0-9]{2}\.dds$')

    chunithm_files = find_dds_files(source_folders, pattern_chunithm)
    chusan_files = find_dds_files(source_folders, pattern_chusan)

    if not os.path.exists(destination_chunithm):
        os.makedirs(destination_chunithm)
    process_files(chunithm_files, destination_chunithm,
                  "Chunithm Partners")

    if not os.path.exists(destination_chusan):
        os.makedirs(destination_chusan)
    process_files(chusan_files, destination_chusan,
                  "Chusan Partners")

def find_subdirectories(base_directories, subdirectory_names):
    subdirectory_paths = {name: [] for name in subdirectory_names}
    for base_directory in base_directories:
        for subdir in Path(base_directory).rglob('*'):
            if subdir.is_dir() and subdir.name in subdirectory_names:
                subdirectory_paths[subdir.name].append(str(subdir))
    return subdirectory_paths

def main():
    # Define paths and database connection details
    data_folder = "C:\\Users\\Polaris\\Documents\\Chunithm Luminous Plus\\App\\data"
    option_folder = "C:\\Users\\Polaris\\Documents\\Chunithm Luminous Plus\\option"
    output_folder = "Assets"

    mysql_host = "localhost"
    mysql_port = 3306
    mysql_user = "root"
    mysql_password = "password"
    mysql_db = "aime"

    data_types = ['cozynet_chuni_static_accessory', 'cozynet_chuni_static_nameplate',
                  'cozynet_chuni_static_systemvoice', 'cozynet_chuni_static_trophies', 'cozynet_chuni_static_mapicon']
    xml_file_names = ['AvatarAccessory.xml', 'NamePlate.xml',
                      'SystemVoice.xml', 'Trophy.xml', 'MapIcon.xml']

 # Create MySQL database connection URL
    database_url = f"mysql://{mysql_user}:{mysql_password}@{mysql_host}:{mysql_port}/{mysql_db}?charset=utf8mb4"

    # Create an engine and session for the MySQL database
    engine = create_engine(database_url)
    Session = sessionmaker(bind=engine)
    session = Session()

    # Drop existing tables
    for data_type in data_types:
        table = define_table(data_type)
        table.drop(engine, checkfirst=True)

    # Create tables in the MySQL database
    for data_type in data_types:
        define_table(data_type).create(engine, checkfirst=True)

    # Process XML files and insert data into the database
    process_directories(data_folder, session, data_types, xml_file_names)

    # Commit changes and close the session
    session.commit()
    session.close()

    # Find relevant subdirectories for file conversion
    relevant_subdirs = ['music', 'avatarAccessory', 'mapIcon', 'namePlate', 'ddsImage', 'systemVoice']

    subdirectories = find_subdirectories(
        [data_folder, option_folder], relevant_subdirs)

    # Define output path
    output_path = os.path.join(
        os.path.expanduser('~'), 'Desktop', output_folder)
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    # Process and convert DDS files to PNGs
    process_partners(subdirectories['ddsImage'], os.path.join(
        output_path, 'public', 'chunithm_partners'), os.path.join(output_path, 'public', 'chusan_partners'))
    process_nameplates(subdirectories['namePlate'], os.path.join(
        output_path, 'public', 'namePlates'), 'namePlate')
    process_systemVoiceImages(subdirectories['systemVoice'], os.path.join(
        output_path, 'public', 'systemVoiceThumbnails'), 'systemVoice thumbnail')
    process_jacketArts(subdirectories['music'], os.path.join(
        output_path, 'public', 'jacketArts'), 'jacket art')
    process_mapIcons(subdirectories['mapIcon'], os.path.join(
        output_path, 'public', 'mapIcons'), 'map icon')
    process_avatarAccessories(subdirectories['avatarAccessories'], os.path.join(
        output_path, 'public', 'avatarAccessories'), 'avatarAccessories Progress')

    print("All conversions and transfers complete.")

if __name__ == "__main__":
    main()