import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import cProfile
import pstats
import subprocess
import os
import threading
import snakeviz.main  # Import snakeviz
import sys


class PythonProfilerGUI:
    """
    A GUI application for profiling Python scripts using cProfile,
    with advanced features like argument passing, profiling options,
    report visualization, progress bar, and threading.
    """

    def __init__(self, master):
        """
        Initializes the main application window.

        Args:
            master (tk.Tk): The root Tk window.
        """
        self.master = master
        master.title("Python Script Profiler")

        # --- GUI Elements ---
        # Script Path
        self.script_path_label = tk.Label(master, text="Script Path:")
        self.script_path_label.grid(row=0, column=0, padx=5, pady=5, sticky="e")

        self.script_path_entry = tk.Entry(master, width=50)
        self.script_path_entry.grid(row=0, column=1, padx=5, pady=5, sticky="ew")

        self.browse_button = tk.Button(master, text="Browse", command=self.browse_file)
        self.browse_button.grid(row=0, column=2, padx=5, pady=5)

        # Arguments
        self.args_label = tk.Label(master, text="Arguments:")
        self.args_label.grid(row=1, column=0, padx=5, pady=5, sticky="e")

        self.args_entry = tk.Entry(master, width=50)
        self.args_entry.grid(row=1, column=1, padx=5, pady=5, sticky="ew")
        self.args_entry.insert(0, "")  # Default arguments

        # Output File
        self.output_path_label = tk.Label(master, text="Output File:")
        self.output_path_label.grid(row=2, column=0, padx=5, pady=5, sticky="e")

        self.output_path_entry = tk.Entry(master, width=50)
        self.output_path_entry.grid(row=2, column=1, padx=5, pady=5, sticky="ew")
        self.output_path_entry.insert(0, "profile_output.prof")  # Default output file

        # Profiling Options (currently placeholder)
        self.profiling_options_label = tk.Label(master, text="Profiling Options:")
        self.profiling_options_label.grid(row=3, column=0, padx=5, pady=5, sticky="e")

        # TODO: Add actual profiling options (e.g., sample rate) here.
        # For now, a simple placeholder:
        self.profiling_options_entry = tk.Entry(master, width=50)
        self.profiling_options_entry.grid(row=3, column=1, padx=5, pady=5, sticky="ew")
        self.profiling_options_entry.insert(0, "")  # Placeholder, no options yet

        # Profile Button
        self.profile_button = tk.Button(master, text="Profile Script", command=self.start_profiling)
        self.profile_button.grid(row=4, column=1, padx=5, pady=10)

        # Progress Bar
        self.progress_bar = ttk.Progressbar(master, orient="horizontal", length=300, mode="indeterminate")
        self.progress_bar.grid(row=5, column=1, padx=5, pady=5)

        # --- Menu ---
        self.menu_bar = tk.Menu(master)
        self.file_menu = tk.Menu(self.menu_bar, tearoff=0)
        self.file_menu.add_command(label="Exit", command=master.quit)
        self.menu_bar.add_cascade(label="File", menu=self.file_menu)

        self.view_menu = tk.Menu(self.menu_bar, tearoff=0)
        self.view_menu.add_command(label="View Report (Snakeviz)", command=self.view_report)
        self.menu_bar.add_cascade(label="View", menu=self.view_menu)

        self.help_menu = tk.Menu(self.menu_bar, tearoff=0)
        self.help_menu.add_command(label="About", command=self.show_about)
        self.menu_bar.add_cascade(label="Help", menu=self.help_menu)

        master.config(menu=self.menu_bar)

        self.profiling_thread = None  # Store the profiling thread

    def browse_file(self):
        """
        Opens a file dialog to select the Python script to profile.
        """
        filename = filedialog.askopenfilename(
            initialdir=os.getcwd(),
            title="Select a Python Script",
            filetypes=(("Python files", "*.py"), ("All files", "*.*"))
        )
        self.script_path_entry.delete(0, tk.END)
        self.script_path_entry.insert(0, filename)

    def start_profiling(self):
        """
        Starts the profiling process in a separate thread.
        """
        script_path = self.script_path_entry.get()
        output_file = self.output_path_entry.get()

        if not script_path:
            messagebox.showerror("Error", "Please select a script to profile.")
            return

        if not os.path.exists(script_path):
            messagebox.showerror("Error", "Script file not found.")
            return

        # Disable the profile button and start the progress bar
        self.profile_button.config(state="disabled")
        self.progress_bar.start()

        # Create a thread for profiling to avoid blocking the GUI
        self.profiling_thread = threading.Thread(target=self.profile_script,
                                                 args=(script_path, output_file))
        self.profiling_thread.start()

    def profile_script(self, script_path, output_file):
        """
        Profiles the selected Python script using cProfile and saves the output to a file.
        Displays error messages if the script path is invalid or if an error occurs during profiling.
        This function runs in a separate thread.

        Args:
            script_path (str): The path to the Python script.
            output_file (str): The path to save the profiling results.
        """
        try:
            args_str = self.args_entry.get()
            args = args_str.split()  # Split arguments by space

            # Use subprocess to execute the script, capturing output for error reporting.
            with cProfile.Profile() as pr:
                command = ["python", script_path] + args  # Script path + arguments
                process = subprocess.Popen(command,
                                         stdout=subprocess.PIPE,
                                         stderr=subprocess.PIPE,
                                         text=True)
                stdout, stderr = process.communicate()
                return_code = process.returncode

            if return_code != 0:
                self.master.after(0, lambda: messagebox.showerror("Error", f"Script execution failed:\n{stderr}"))  # Use after to update GUI from thread
                return

            stats = pstats.Stats(pr)
            stats.sort_stats(pstats.SortKey.TIME)
            stats.dump_stats(output_file)

            self.master.after(0, lambda: messagebox.showinfo("Success", f"Profiling complete. Results saved to {output_file}"))  # Use after to update GUI from thread

        except Exception as e:
            self.master.after(0, lambda: messagebox.showerror("Error", f"An error occurred during profiling: {e}"))  # Use after to update GUI from thread
        finally:
            # Re-enable the profile button and stop the progress bar
            self.master.after(0, lambda: self.profile_button.config(state="normal"))  # Use after to update GUI from thread
            self.master.after(0, lambda: self.progress_bar.stop())  # Use after to update GUI from thread

    def view_report(self):
        """
        Opens the profiling report using Snakeviz.
        """
        output_file = self.output_path_entry.get()

        if not os.path.exists(output_file):
            messagebox.showerror("Error", "Profiling output file not found. Profile the script first.")
            return

        try:
            # Use subprocess to run snakeviz
            subprocess.run(["snakeviz", output_file], check=True)

        except FileNotFoundError:
            messagebox.showerror("Error", "Snakeviz is not installed or not in your PATH.")
        except subprocess.CalledProcessError as e:
            messagebox.showerror("Error", f"An error occurred while viewing the report with Snakeviz: {e}")

    def show_about(self):
        """
        Displays an "About" dialog with information about the application.
        """
        messagebox.showinfo("About", "Python Script Profiler\nSimple tool for profiling Python scripts using cProfile.")

# --- Main ---
if __name__ == "__main__":
    root = tk.Tk()
    gui = PythonProfilerGUI(root)
    root.mainloop()