import tkinter as tk
from tkinter import filedialog
import os
import sys
import traceback

# --- Configuration ---
TIA_PORTAL_VERSION = "18.0" # Target TIA Portal version (e.g., "18.0")
EXPORT_OPTIONS = None # Use default export options
KEEP_FOLDER_STRUCTURE = True # Replicate TIA project folder structure in export directory

# --- TIA Scripting Import Handling ---
# Check if the TIA_SCRIPTING environment variable is set
if os.getenv('TIA_SCRIPTING'):
    sys.path.append(os.getenv('TIA_SCRIPTING'))
else:
    # Optional: Define a fallback path if the environment variable isn't set
    # fallback_path = "C:\\path\\to\\your\\TIA_Scripting_binaries"
    # if os.path.exists(fallback_path):
    #     sys.path.append(fallback_path)
    pass # Allow import to fail if not found

try:
    import siemens_tia_scripting as ts
    EXPORT_OPTIONS = ts.Enums.ExportOptions.WithDefaults # Set default options now that 'ts' is imported
except ImportError:
    print("ERROR: Failed to import 'siemens_tia_scripting'.")
    print("Ensure:")
    print(f"1. TIA Portal Openness for V{TIA_PORTAL_VERSION} is installed.")
    print("2. The 'siemens_tia_scripting' Python module is installed (pip install ...) or")
    print("   the path to its binaries is set in the 'TIA_SCRIPTING' environment variable.")
    print("3. You are using a compatible Python version (e.g., 3.12.X as per documentation).")
    sys.exit(1)
except Exception as e:
    print(f"An unexpected error occurred during import: {e}")
    traceback.print_exc()
    sys.exit(1)

# --- Functions ---

def select_project_file():
    """Opens a dialog to select a TIA Portal project file."""
    root = tk.Tk()
    root.withdraw() # Hide the main tkinter window
    file_path = filedialog.askopenfilename(
        title="Select TIA Portal Project File",
        filetypes=[(f"TIA Portal V{TIA_PORTAL_VERSION} Projects", f"*.ap{TIA_PORTAL_VERSION.split('.')[0]}")] # e.g. *.ap18
    )
    root.destroy()
    if not file_path:
        print("No project file selected. Exiting.")
        sys.exit(0)
    return file_path

def select_export_directory():
    """Opens a dialog to select the export directory."""
    root = tk.Tk()
    root.withdraw() # Hide the main tkinter window
    dir_path = filedialog.askdirectory(
        title="Select Export Directory"
    )
    root.destroy()
    if not dir_path:
        print("No export directory selected. Exiting.")
        sys.exit(0)
    return dir_path

def export_plc_data(plc, export_base_dir):
    """Exports Blocks, UDTs, and Tag Tables from a given PLC."""
    plc_name = plc.get_name()
    print(f"\n--- Processing PLC: {plc_name} ---")

    # Define base export path for this PLC
    plc_export_dir = os.path.join(export_base_dir, plc_name)
    os.makedirs(plc_export_dir, exist_ok=True)

    # --- Export Program Blocks ---
    blocks_exported = 0
    blocks_skipped = 0
    print(f"\n[PLC: {plc_name}] Exporting Program Blocks...")
    xml_blocks_path = os.path.join(plc_export_dir, "ProgramBlocks_XML")
    scl_blocks_path = os.path.join(plc_export_dir, "ProgramBlocks_SCL")
    os.makedirs(xml_blocks_path, exist_ok=True)
    os.makedirs(scl_blocks_path, exist_ok=True)
    print(f"  XML Target: {xml_blocks_path}")
    print(f"  SCL Target: {scl_blocks_path}")

    try:
        program_blocks = plc.get_program_blocks() #
        print(f"  Found {len(program_blocks)} program blocks.")
        for block in program_blocks:
            block_name = block.get_name() # Assuming get_name() exists
            print(f"    Processing block: {block_name}...")
            try:
                if not block.is_consistent(): #
                    print(f"      Compiling block {block_name}...")
                    block.compile() #
                    if not block.is_consistent():
                       print(f"      WARNING: Block {block_name} inconsistent after compile. Skipping.")
                       blocks_skipped += 1
                       continue

                print(f"      Exporting {block_name} as XML...")
                block.export(target_directory_path=xml_blocks_path, #
                             export_options=EXPORT_OPTIONS, #
                             export_format=ts.Enums.ExportFormats.SimaticML, #
                             keep_folder_structure=KEEP_FOLDER_STRUCTURE) #

                try:
                   prog_language = block.get_property(name="ProgrammingLanguage")
                   if prog_language == "SCL":
                       print(f"      Exporting {block_name} as SCL...")
                       block.export(target_directory_path=scl_blocks_path,
                                    export_options=EXPORT_OPTIONS,
                                    export_format=ts.Enums.ExportFormats.ExternalSource, #
                                    keep_folder_structure=KEEP_FOLDER_STRUCTURE)
                except Exception as prop_ex:
                    print(f"      Could not get ProgrammingLanguage for {block_name}. Skipping SCL. Error: {prop_ex}")

                blocks_exported += 1
            except Exception as block_ex:
                print(f"      ERROR exporting block {block_name}: {block_ex}")
                blocks_skipped += 1
        print(f"  Program Blocks Export Summary: Exported={blocks_exported}, Skipped/Errors={blocks_skipped}")
    except Exception as e:
        print(f"  ERROR processing Program Blocks: {e}")
        traceback.print_exc()

    # --- Export PLC Data Types (UDTs) ---
    udts_exported = 0
    udts_skipped = 0
    print(f"\n[PLC: {plc_name}] Exporting PLC Data Types (UDTs)...")
    udt_export_path = os.path.join(plc_export_dir, "PlcDataTypes")
    os.makedirs(udt_export_path, exist_ok=True)
    print(f"  Target: {udt_export_path}")

    try:
        udts = plc.get_user_data_types() #
        print(f"  Found {len(udts)} UDTs.")
        for udt in udts:
            udt_name = udt.get_name() #
            print(f"    Processing UDT: {udt_name}...")
            try:
                if not udt.is_consistent(): #
                    print(f"      Compiling UDT {udt_name}...")
                    udt.compile() #
                    if not udt.is_consistent():
                        print(f"      WARNING: UDT {udt_name} inconsistent after compile. Skipping.")
                        udts_skipped += 1
                        continue

                print(f"      Exporting {udt_name}...")
                udt.export(target_directory_path=udt_export_path, #
                           export_options=EXPORT_OPTIONS, #
                           # export_format defaults to SimaticML for UDTs
                           keep_folder_structure=KEEP_FOLDER_STRUCTURE) #
                udts_exported += 1
            except Exception as udt_ex:
                print(f"      ERROR exporting UDT {udt_name}: {udt_ex}")
                udts_skipped += 1
        print(f"  UDT Export Summary: Exported={udts_exported}, Skipped/Errors={udts_skipped}")
    except Exception as e:
        print(f"  ERROR processing UDTs: {e}")
        traceback.print_exc()

    # --- Export PLC Tag Tables ---
    tags_exported = 0
    tags_skipped = 0
    print(f"\n[PLC: {plc_name}] Exporting PLC Tag Tables...")
    tags_export_path = os.path.join(plc_export_dir, "PlcTags")
    os.makedirs(tags_export_path, exist_ok=True)
    print(f"  Target: {tags_export_path}")

    try:
        tag_tables = plc.get_plc_tag_tables() #
        print(f"  Found {len(tag_tables)} Tag Tables.")
        for table in tag_tables:
            table_name = table.get_name() #
            print(f"    Processing Tag Table: {table_name}...")
            try:
                # Note: Consistency check might not be available/needed for tag tables like blocks/UDTs
                print(f"      Exporting {table_name}...")
                table.export(target_directory_path=tags_export_path, #
                             export_options=EXPORT_OPTIONS, #
                             # export_format defaults to SimaticML for Tag Tables
                             keep_folder_structure=KEEP_FOLDER_STRUCTURE) #
                tags_exported += 1
            except Exception as table_ex:
                print(f"      ERROR exporting Tag Table {table_name}: {table_ex}")
                tags_skipped += 1
        print(f"  Tag Table Export Summary: Exported={tags_exported}, Skipped/Errors={tags_skipped}")
    except Exception as e:
        print(f"  ERROR processing Tag Tables: {e}")
        traceback.print_exc()

    print(f"\n--- Finished processing PLC: {plc_name} ---")


# --- Main Script ---

if __name__ == "__main__":
    print("--- TIA Portal Data Exporter (Blocks, UDTs, Tags) ---")

    # 1. Select Files/Folders
    project_file = select_project_file()
    export_dir = select_export_directory()

    print(f"\nSelected Project: {project_file}")
    print(f"Selected Export Directory: {export_dir}")

    portal_instance = None
    project_object = None

    try:
        # 2. Connect to TIA Portal
        print(f"\nConnecting to TIA Portal V{TIA_PORTAL_VERSION}...")
        portal_instance = ts.open_portal(
            version=TIA_PORTAL_VERSION,
            portal_mode=ts.Enums.PortalMode.WithGraphicalUserInterface
        )
        print("Connected to TIA Portal.")
        print(f"Portal Process ID: {portal_instance.get_process_id()}") #

        # 3. Open Project
        print(f"Opening project: {os.path.basename(project_file)}...")
        project_object = portal_instance.open_project(project_file_path=project_file) #
        if project_object is None:
             print("Project might already be open, attempting to get handle...")
             project_object = portal_instance.get_project() #
             if project_object is None:
                 raise Exception("Failed to open or get the specified project.")
        print("Project opened successfully.")

        # 4. Get PLCs
        plcs = project_object.get_plcs() #
        if not plcs:
            print("No PLC devices found in the project.")
        else:
            print(f"Found {len(plcs)} PLC(s). Starting export process...")
            
            # 5. Iterate and Export Data for each PLC
            for plc_device in plcs:
                export_plc_data(plc=plc_device, export_base_dir=export_dir)

        print("\nExport process completed.")

    except ts.TiaException as tia_ex:
        print(f"\nTIA Portal Openness Error: {tia_ex}")
        traceback.print_exc()
    except FileNotFoundError:
        print(f"\nERROR: Project file not found at {project_file}")
    except Exception as e:
        print(f"\nAn unexpected error occurred: {e}")
        traceback.print_exc()
    finally:
        # 6. Cleanup
        if portal_instance:
            try:
                print("\nClosing TIA Portal...")
                portal_instance.close_portal() #
                print("TIA Portal closed.")
            except Exception as close_ex:
                print(f"Error during TIA Portal cleanup: {close_ex}")
        
        print("\nScript finished.")