import os
from tkinter import Tk
from tkinter import filedialog
# Add Z values to gcode
# Revision 1
# Dutcher and A.I. Tool

def read_gcode_file(file_path):
    """
    Read a G-Code file and return its contents as a list of lines.
    
    Args:
        file_path (str): Path to the G-Code file.
    
    Returns:
        list: List of lines in the G-Code file.
    """
    try:
        with open(file_path, 'r') as f:
            return f.readlines()
    except FileNotFoundError:
        print(f"File {file_path} not found.")
        return []

def identify_sections(lines):
    """
    Identify the three sections of the G-Code: header, x-y-z values, and end of the main program.
    
    Args:
        lines (list): List of G-Code lines.
    
    Returns:
        tuple: Tuple containing the header, x-y-z values, and end of the main program sections.
    """
    header = []
    xy_z_values = []
    end_of_program = []
    
    section = 'header'
    
    for line in lines:
        if section == 'header' and line.strip().startswith('; Show'):
            section = 'xy_z_values'
            xy_z_values.append(line)
        elif section == 'header':
            header.append(line)
        elif section == 'xy_z_values' and line.strip().startswith('; End of main program'):
            section = 'end_of_program'
            end_of_program.append(line)
        elif section == 'xy_z_values':
            xy_z_values.append(line)
        elif section == 'end_of_program':
            end_of_program.append(line)
    
    return header, xy_z_values, end_of_program

def insert_z_values(lines):
    """
    Insert missing Z-values into the G-Code lines.
    
    Args:
        lines (list): List of G-Code lines.
    
    Returns:
        list: List of modified G-Code lines with inserted Z-values.
    """
    previous_z = None
    for i, line in enumerate(lines):
        if 'Z' in line:
            # Extract the Z-value from the line
            z_value = float(line.split('Z')[1].split(' ')[0])
            previous_z = z_value
        else:
            if previous_z is not None:
                # Insert the Z-value into the line
                parts = line.strip().split()
                if 'G1' in parts:
                    g1_index = parts.index('G1')
                    last_coord_index = max([j for j, part in enumerate(parts) if part.startswith('X') or part.startswith('Y')])
                    parts.insert(last_coord_index + 1, f'Z{previous_z:.3f}')
                    lines[i] = ' '.join(parts) + '\n'
    return lines

def write_gcode_file(file_path, lines):
    """
    Write the modified G-Code lines to a new file.
    
    Args:
        file_path (str): Path to the output file.
        lines (list): List of modified G-Code lines.
    """
    with open(file_path, 'w') as f:
        f.writelines(lines)

def main():
    # Create a Tkinter window to use the file dialog
    root = Tk()
    root.withdraw()  # Hide the window
    
    # Prompt the user to browse for the input file
    input_file_path = filedialog.askopenfilename(title="Select G-Code File", filetypes=[("G-Code Files", "*.gcode")])
    
    if input_file_path:
        # Get the output file path by appending " Z Mod" to the input file name
        output_file_path = os.path.splitext(input_file_path)[0] + " Z Mod.gcode"
        
        lines = read_gcode_file(input_file_path)
        header, xy_z_values, end_of_program = identify_sections(lines)
        modified_xy_z_values = insert_z_values(xy_z_values)
        combined_lines = header + modified_xy_z_values + end_of_program
        write_gcode_file(output_file_path, combined_lines)
        print(f"Modified G-Code file saved to {output_file_path}")
    else:
        print("No file selected.")

if __name__ == "__main__":
    main()