import sys
import tkinter as tk
from tkinter import ttk, messagebox
import xml.etree.ElementTree as ET


file_path = '../XML/vasprun.xml'
window_size = '800x600'

nargs = len(sys.argv)
if nargs >= 2:
    file_path = sys.argv[1]
if nargs >= 3:
    window_size = sys.argv[2]


class XMLViewerEditor:
    def __init__(self, root, file_path, window_size):
        self.root = root
        self.root.title("XML Viewer and Editor")
        self.root.geometry(window_size)

        self.tree = ttk.Treeview(self.root)
        self.tree.heading('#0', text='Element')
        self.tree.pack(fill='both', expand=True)

        self.load_xml(file_path)

        self.text_editor = tk.Text(self.root, wrap="none")
        self.text_editor.pack(fill="both", expand=True)
        self.text_editor.grid_remove()

        self.edit_button = tk.Button(self.root, text="Edit XML", command=self.toggle_edit_mode)
        self.edit_button.pack()

    def load_xml(self, file_path):
        try:
            tree = ET.parse(file_path)
            self.root_element = tree.getroot()
            self.display_element(self.root_element, "")
        except ET.ParseError as e:
            messagebox.showerror("Error", str(e))

    def display_element(self, element, parent_id):
        item = self.tree.insert(parent_id, "end", text=element.tag)
        for key, value in element.attrib.items():
            self.tree.insert(item, "end", text=f"{key}: {value}")
        for child in element:
            self.display_element(child, item)

    def toggle_edit_mode(self):
        if self.text_editor.winfo_viewable():
            edited_xml = self.text_editor.get("1.0", "end")
            try:
                new_root = ET.fromstring(edited_xml)
                self.root_element = new_root
                self.tree.delete(*self.tree.get_children())
                self.display_element(self.root_element, "")
                self.text_editor.grid_remove()
                self.edit_button.config(text="Edit XML")
            except ET.ParseError as e:
                messagebox.showerror("Error", str(e))
        else:
            xml_string = ET.tostring(self.root_element).decode()
            self.text_editor.delete("1.0", "end")
            self.text_editor.insert("1.0", xml_string)
            self.text_editor.grid()
            self.edit_button.config(text="Save XML")

if __name__ == "__main__":
    root = tk.Tk()
    app = XMLViewerEditor(root, file_path, window_size)
    root.mainloop()