#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# manjaro-rescue
# Based on code from BigLinux project
# Original copyright © 2024 BigLinux
# Updated and adapted for Manjaro by Stefano Capitani, 2026
# License: GNU GPLv3

VERSION = "3.0.0"

import sys
import os
import gi
import subprocess
import shutil
import pwd
import argparse
import threading
import shlex

# Get the real directory of the script (resolves symbolic links like the one in /usr/bin)
APP_DIR = os.path.dirname(os.path.realpath(__file__))

# Ensure the script directory is in the path for local imports
sys.path.insert(0, APP_DIR)

# Ensure we're using the correct versions
try:
    gi.require_version('Gtk', '4.0')
    gi.require_version('Adw', '1')
except (ImportError, ValueError) as e:
    print(f"Error: Missing GTK4/Adwaita dependencies. {e}", file=sys.stderr)
    sys.exit(1)

from gi.repository import Gtk, Adw, GLib, Gio, Gdk
from backend.system_interface import SystemInterface

# Add the 'icons' directory to the GTK icon search path.
# This allows GTK to find "manjaro-rescue.svg" using the name "manjaro-rescue".
icons_dir = os.path.join(APP_DIR, 'icons')
if os.path.isdir(icons_dir) and Gdk.Display.get_default():
    icon_theme = Gtk.IconTheme.get_for_display(Gdk.Display.get_default())
    icon_theme.add_search_path(icons_dir)

# --- GUI LOGIC (Unified) ---

@Gtk.Template(filename=os.path.join(APP_DIR, 'window.ui'))
class ManjaroRescueWindow(Adw.ApplicationWindow):
    __gtype_name__ = 'ManjaroRescueWindow'

    main_stack = Gtk.Template.Child()
    btn_start = Gtk.Template.Child()
    list_systems = Gtk.Template.Child()
    list_restore_options = Gtk.Template.Child()
    progress_bar = Gtk.Template.Child()
    lbl_status = Gtk.Template.Child()
    box_finished = Gtk.Template.Child()
    btn_back_to_options = Gtk.Template.Child()
    btn_exit = Gtk.Template.Child()
    menu_button = Gtk.Template.Child()
    toast_overlay = Gtk.Template.Child()
    status_page_welcome = Gtk.Template.Child()
    status_page_selection = Gtk.Template.Child()
    status_page_progress = Gtk.Template.Child()
    status_page_options = Gtk.Template.Child()

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        from utils.translation import _
        self._ = _
        self.backend = SystemInterface()
        # The self.current_process variable is used to track the external process
        # in case termination or monitoring features need to be added.
        self.current_process = None
        self.btn_start.connect("clicked", self.on_start_clicked)
        self.list_systems.connect("row-activated", self.on_system_selected)
        self.list_restore_options.connect("row-activated", self.on_option_selected)
        self.btn_back_to_options.connect("clicked", self.on_back_to_options_clicked)
        self.btn_exit.connect("clicked", lambda x: self.get_application().quit())
        
        self.setup_menu()
        self.set_icon_name("manjaro-rescue") # Set the main window icon
        self._set_initial_ui_strings()

    def _set_initial_ui_strings(self):
        # Welcome Page
        self.status_page_welcome.set_title(self._("System Restore"))
        self.status_page_welcome.set_description(
            self._("This tool will help you restore your system.") + "\n\n" +
            self._("This utility should only be used in live mode!") + "\n\n" +
            self._("If the installed system is booting correctly, there are no boot problems, so it's better not to proceed with this tool.")
        )
        self.btn_start.set_label(self._("Start Scan"))

        # Selection Page
        # These will be set again in update_ui_after_detection, but good for initial state
        self.status_page_selection.set_title(self._("Select System"))
        self.status_page_selection.set_description(self._("Choose the partition where the system to be restored is installed."))

        # Progress Page
        self.status_page_progress.set_title(self._("Scan in progress"))
        self.lbl_status.set_label(self._("Initializing..."))
        self.btn_back_to_options.set_label(self._("Back to options"))
        self.btn_exit.set_label(self._("Exit"))

        # Options Page
        # These will be set again in populate_restore_options, but good for initial state
        self.status_page_options.set_title(self._("Restore Options"))
        self.status_page_options.set_description(self._("The restore can be performed in 3 ways. If in doubt, try the first option."))

    def setup_menu(self):
        # Create the menu for the three dots (hamburger menu)
        menu = Gio.Menu.new()
        menu.append(self._("About"), "win.about")
        menu.append(self._("Quit"), "app.quit")
        self.menu_button.set_menu_model(menu)
        
        # Add the 'about' action to the window
        about_action = Gio.SimpleAction.new("about", None)
        about_action.connect("activate", self.show_about_dialog)
        self.add_action(about_action)

    def show_about_dialog(self, action, param):
        about = Adw.AboutWindow(
            transient_for=self,
            application_name=self._("Manjaro Rescue"),
            application_icon="manjaro-rescue",
            developer_name="Stefano Capitani",
            developers=["Stefano Capitani", "Bruno Goncalves", "Tales A. Mendonça"],
            comments=self._("Based on original work by BigLinux Team"),
            version=VERSION,
            copyright=self._("Original code © 2024 BigLinux | Modifications © 2026 Stefano Capitani"),
            website="https://manjaro.org",
            issue_url="https://codeberg.org/Ste74/manjaro-rescue/issues"
        )
        # Connect the signal to handle link opening safely
        about.connect("activate-link", self._on_about_link_activated)
        about.present()

    def _on_about_link_activated(self, window, uri):
        """Handles opening links in the About window as a non-root user"""
        try:
            # Identify the original user (pkexec sets PKEXEC_UID)
            uid = os.environ.get('PKEXEC_UID')
            user = None
            
            if uid and uid != '0':
                try:
                    user = pwd.getpwuid(int(uid)).pw_name
                except:
                    pass
            
            # Specific fallback for Manjaro Live if UID is not found
            if not user and os.path.exists('/home/manjaro'):
                user = 'manjaro'

            if user:
                # Execute xdg-open as the original user
                subprocess.Popen(['sudo', '-u', user, 'xdg-open', uri])
                return True # Signal that we handled the link
        except Exception as e:
            print(f"DEBUG: Error opening link as user: {e}")
        
        # Standard fallback if everything else fails
        return False

    def on_start_clicked(self, button):
        self.main_stack.set_visible_child_name("progress")
        self.status_page_progress.set_title(self._("Scan in progress"))
        self.progress_bar.set_fraction(0.1)
        self.lbl_status.set_text(self._("Scanning systems..."))
        
        # Start detection in a separate thread to avoid blocking the UI
        thread = threading.Thread(target=self.perform_detection)
        thread.daemon = True
        thread.start()

    def perform_detection(self):
        try:
            self.backend.detect_systems()
            # Return to the main thread to update the interface
            GLib.idle_add(self.update_ui_after_detection)
        except Exception as e:
            GLib.idle_add(self.show_error, str(e))

    def update_ui_after_detection(self):
        self.populate_systems()
        # Ensure selection page strings are set when navigating to it
        self.status_page_selection.set_title(self._("Select System"))
        self.status_page_selection.set_description(self._("Choose the partition where the system to be restored is installed."))
        self.main_stack.set_visible_child_name("selection")
        return GLib.SOURCE_REMOVE

    def show_error(self, error_msg):
        # Translate "Error:" prefix
        self.toast_overlay.add_toast(Adw.Toast.new(f"{self._('Error:')} {error_msg}"))
        self.main_stack.set_visible_child_name("welcome")
        return GLib.SOURCE_REMOVE

    def populate_systems(self):
        # Clear the previous list
        while child := self.list_systems.get_first_child():
            self.list_systems.remove(child)
            
        for i, sys in enumerate(self.backend.detected_systems):
            row = Adw.ActionRow(title=sys['name'], subtitle=f"{self._('Partition:')} {sys['partition']} ({self._('Filesystem:')} {sys['filesystem']})")
            # Make the row clickable to emit the row-activated signal
            row.set_activatable(True)
            self.list_systems.append(row)

    def on_system_selected(self, listbox, row):
        # Get the index of the selected row
        index = row.get_index()
        if index != -1:
            selected_sys = self.backend.detected_systems[index]
            
            selected_efi = None
            selected_disk = None

            print(f"DEBUG: Boot mode detected by backend: {self.backend.boot_mode}")
            print(f"DEBUG: EFI partitions detected: {self.backend.efi_partitions}")
            print(f"DEBUG: GRUB disks detected: {self.backend.grub_disks}")

            # Attempt automatic selection to avoid errors in backend scripts
            if self.backend.boot_mode == "EFI":
                if len(self.backend.efi_partitions) == 1:
                    selected_efi = self.backend.efi_partitions[0]
                    print(f"DEBUG: Auto-selected EFI partition: {selected_efi}")
            elif self.backend.boot_mode == "LEGACY":
                if len(self.backend.grub_disks) == 1:
                    selected_disk = self.backend.grub_disks[0]['device']
                    print(f"DEBUG: Auto-selected Disk: {selected_disk}")
                else:
                    print(f"DEBUG: Cannot auto-select GRUB disk. Found {len(self.backend.grub_disks)} GRUB disks.")
            else:
                print(f"DEBUG: Unknown boot mode: {self.backend.boot_mode}")

            # Save the selection in the backend (creates files in /tmp needed by scripts)
            self.backend.save_selection(selected_sys, selected_efi=selected_efi, selected_disk=selected_disk)
            
            self.populate_restore_options()
            self.main_stack.set_visible_child_name("options")
            # Ensure options page strings are set when navigating to it
            self.status_page_options.set_title(self._("Restore Options"))
            self.status_page_options.set_description(self._("The restore can be performed in 3 ways. If in doubt, try the first option."))

    def populate_restore_options(self):
        while child := self.list_restore_options.get_first_child():
            self.list_restore_options.remove(child)
            
        options = [
            {"id": 1, "title": self._("Simple"), "desc": self._("Simple, just writes the Grub again at the beginning of the disk.")},
            {"id": 2, "title": self._("Intermediate"), "desc": self._("Intermediate, reinstalls the grub package on the partition, regenerates the configuration, and updates the initrd.")},
            {"id": 3, "title": self._("Complete"), "desc": self._("Complete, performs the intermediate restore steps, updates the system, and checks if the LTS kernel is installed.")}
        ]
        
        for opt in options:
            row = Adw.ActionRow(title=opt['title'], subtitle=opt['desc'])
            row.set_activatable(True)
            # Add an arrow icon to the right
            row.add_suffix(Gtk.Image.new_from_icon_name("go-next-symbolic"))
            self.list_restore_options.append(row)

        # Optional Tools section via ExpanderRow
        expander = Adw.ExpanderRow(title=self._("Optional Tools"))
        expander.set_icon_name("utilities-system-symbolic")

        # Terminal in Chroot
        target_term = self.backend.get_target_terminal()
        term_row = Adw.ActionRow(
            title=self._("Interactive Terminal"),
            subtitle=self._("Open a terminal inside the target system (using %s).") % target_term
        )
        term_row.set_activatable(True)
        term_row.add_suffix(Gtk.Image.new_from_icon_name("utilities-terminal-symbolic"))
        term_row.connect("activated", lambda r: self._launch_terminal())
        expander.add_row(term_row)

        # Manjaro Log Helper
        if shutil.which('mlh'):
            log_row = Adw.ActionRow(
                title=self._("Manjaro Log Helper"),
                subtitle=self._("Collect system logs for troubleshooting.")
            )
            log_row.set_activatable(True)
            log_row.add_suffix(Gtk.Image.new_from_icon_name("view-list-symbolic"))
            log_row.connect("activated", lambda r: self._launch_log_helper())
            expander.add_row(log_row)

        # Bmenu
        if shutil.which('bmenu'):
            bmenu_row = Adw.ActionRow(
                title=self._("Bmenu"),
                subtitle=self._("Powerful system menu.")
            )
            bmenu_row.set_activatable(True)
            bmenu_row.add_suffix(Gtk.Image.new_from_icon_name("utilities-terminal-symbolic"))
            bmenu_row.connect("activated", lambda r: self._launch_bmenu())
            expander.add_row(bmenu_row)

        # Timeshift (if available)
        if self.backend.is_timeshift_available():
            ts_row = Adw.ActionRow(
                title=self._("Timeshift"),
                subtitle=self._("Restore a previous snapshot using Timeshift.")
            )
            ts_row.set_activatable(True)
            ts_row.add_suffix(Gtk.Image.new_from_icon_name("document-revert-symbolic"))
            ts_row.connect("activated", lambda r: self._launch_timeshift())
            expander.add_row(ts_row)

        self.list_restore_options.append(expander)

    def on_option_selected(self, listbox, row):
        index = row.get_index()
        if index != -1:
            # We only handle the first 3 restore options.
            # Optional tools are now handled via direct signals in the ExpanderRow.
            if index < 3:
                mode = index + 1
                self._start_restore_process(mode)

    def _start_restore_process(self, mode):
            self.main_stack.set_visible_child_name("progress")
            self.status_page_progress.set_title(self._("System Restore"))
            self.box_finished.set_visible(False)
            self.progress_bar.set_fraction(0.0)
            self.progress_bar.set_fraction(0.1)
            self.lbl_status.set_text(self._("Starting restore..."))
            
            print(f"DEBUG: Starting restore thread for mode {mode}")
            thread = threading.Thread(target=self.run_restore, args=(mode,))
            thread.daemon = True
            thread.start()

    def _get_terminal_command(self, command_list):
        """Helper to wrap a command for execution in the detected host terminal."""
        host_term = None
        terminals = [
            'gnome-terminal', 'xfce4-terminal', 'konsole', 'gnome-console', 'kgx',
            'mate-terminal', 'lxterminal', 'qterminal', 'tilix',
            'alacritty', 'foot', 'kitty', 'termite', 'urxvt',
            'terminology', 'x-terminal-emulator', 'st', 'wezterm', 
            'tilda', 'xterm'
        ]
        
        for t in terminals:
            if shutil.which(t):
                host_term = t
                break
        
        if not host_term:
            return None

        if host_term in ['gnome-terminal', 'gnome-console', 'kgx']:
            return [host_term, '--'] + command_list
        elif host_term == 'xfce4-terminal':
            return ['xfce4-terminal', '-e', shlex.join(command_list)]
        elif host_term == 'konsole':
            return ['konsole', '--'] + command_list
        else:
            return [host_term, '-e'] + command_list

    def _launch_terminal(self):
        """Launch the Live terminal to enter chroot."""
        target_shell = self.backend.get_target_terminal()
        cmd = self._get_terminal_command(['manjaro-chroot', '/mnt', target_shell])
        
        if cmd:
            subprocess.Popen(cmd)
        else:
            self.show_error(self._("No host terminal found to launch chroot."))

    def _launch_log_helper(self):
        """Launch Manjaro Log Helper (spawns its own Yad/Zenity GUI)."""
        subprocess.Popen(['mlh'])

    def _launch_bmenu(self):
        """Launch Bmenu in a terminal wrapper (closes terminal on exit)."""
        cmd = self._get_terminal_command(['bmenu'])
        if cmd:
            subprocess.Popen(cmd)
        else:
            # Fallback if no terminal found, though unlikely on Manjaro
            self.show_error(self._("No host terminal found to launch bmenu."))

    def _launch_timeshift(self):
        """Launch Timeshift-GTK as root."""
        if shutil.which('timeshift-gtk'):
            subprocess.Popen(['timeshift-gtk'])
        else:
            print("DEBUG: timeshift-gtk not found in PATH")
            self.show_error(self._("Timeshift is not installed."))

    def run_restore(self, mode):
        full_output = []
        try:
            # If mode is Complete (3), perform interactive update first
            if mode == 3:
                GLib.idle_add(self.status_page_progress.set_title, self._("Complete Restore"))
                GLib.idle_add(self.lbl_status.set_text, self._("Preparing interactive update..."))
                self.backend.prepare_chroot()
                
                cmd = self._get_terminal_command(['manjaro-chroot', '/mnt', 'pacman', '-Syyu'])
                if cmd:
                    GLib.idle_add(self.lbl_status.set_text, self._("System update in progress in terminal..."))
                    # Launch terminal and WAIT for the user to close it
                    process = subprocess.Popen(cmd)
                    process.wait()
                    # After update, proceed with the rest of the operations (Grub/Initrd)
                    # downgrading mode to Intermediate (2) for the backend script.
                    mode = 2

            # Call the backend which returns the Popen process
            print(f"DEBUG: Calling backend.execute_restore({mode})")
            GLib.idle_add(self.status_page_progress.set_title, self._("Applying Changes"))
            process = self.backend.execute_restore(mode)
            self.current_process = process # Save the process reference
            if not process:
                raise Exception(self._("Failed to start restore process (process object is None)."))
            print(f"DEBUG: Process started with PID: {process.pid}")
            
            # Read script output line by line in real time
            try:
                for line in iter(process.stdout.readline, ''):
                    if line:
                        full_output.append(line.strip())
                        # Update the status label in the UI
                        GLib.idle_add(self.lbl_status.set_text, line.strip())
                        GLib.idle_add(self.progress_bar.pulse)
                        print(f"DEBUG: Output script: {line.strip()}") # Also print to terminal for debug
            except Exception as read_e:
                print(f"ERROR: Error reading script output: {read_e}", file=sys.stderr)
                GLib.idle_add(self.lbl_status.set_text, self._("Error reading output."))
                process.kill() # Terminate process if reading fails
                raise # Re-raise the exception to be caught by the outer block
            process.wait()
            # Notify the UI about the process completion
            GLib.idle_add(self.on_restore_finished, process.returncode == 0, "\n".join(full_output))
        except Exception as e:
            GLib.idle_add(self.show_error, str(e))

    def on_back_to_options_clicked(self, button):
        self.box_finished.set_visible(False)
        self.main_stack.set_visible_child_name("options")

    def on_restore_finished(self, success, full_output_text=""):
        self.progress_bar.set_fraction(1.0)
        if success:
            self.lbl_status.set_text(self._("Completed!"))
            self.toast_overlay.add_toast(Adw.Toast.new(self._("Restore completed successfully!")))
        else:
            self.lbl_status.set_text(self._("Restore failed."))
            # If the process failed, try to read the error if available
            if full_output_text:
                print(f"ERROR: Full process output:\n{full_output_text}", file=sys.stderr)
                # Update the status label with the last message or an error summary
                last_line = full_output_text.splitlines()[-1] if full_output_text else ""
                GLib.idle_add(self.lbl_status.set_text, self._("Restore failed. Check terminal for details or last message: ") + last_line)
            self.toast_overlay.add_toast(Adw.Toast.new(self._("Error during restore process.")))
        
        self.box_finished.set_visible(True)
        return GLib.SOURCE_REMOVE

class ManjaroRescueApplication(Adw.Application):
    def __init__(self):
        super().__init__(application_id="org.manjaro.rescue",
                         flags=Gio.ApplicationFlags.FLAGS_NONE | Gio.ApplicationFlags.NON_UNIQUE)

    def do_activate(self):
        win = self.get_active_window()
        if not win:
            win = ManjaroRescueWindow(application=self)
        win.present()

    def do_startup(self):
        Adw.Application.do_startup(self)

def get_user_gsetting(schema, key):
    """Gets a user's gsettings value."""
    try:
        result = subprocess.run(
            ['gsettings', 'get', schema, key],
            capture_output=True, text=True, check=True
        )
        return result.stdout.strip().strip("'")
    except (subprocess.CalledProcessError, FileNotFoundError):
        return None

def main():
    """Main entry point for the application"""
    
    print(f"DEBUG: Executing script from: {os.path.abspath(__file__)}")

    # Set the program name for Wayland association
    GLib.set_prgname("manjaro-rescue")
    
    # Check if running as root
    if os.geteuid() != 0:
        print("This application needs to be run as root. Trying to relaunch with pkexec...")
        try:
            script_path = os.path.abspath(sys.argv[0])
            
            # --- USER SETTINGS DETECTION (RUNNING AS USER) ---
            user_scheme = get_user_gsetting('org.gnome.desktop.interface', 'color-scheme') or 'default'
            user_button_layout = get_user_gsetting('org.gnome.desktop.wm.preferences', 'button-layout') or ''

            env_vars_to_preserve = []
            for var in ['DISPLAY', 'XAUTHORITY', 'DBUS_SESSION_BUS_ADDRESS', 'WAYLAND_DISPLAY']:
                if var in os.environ:
                    env_vars_to_preserve.append(f'{var}={os.environ[var]}')

            args = ['pkexec']
            if env_vars_to_preserve:
                args.extend(['env'] + env_vars_to_preserve)
            
            args.extend([sys.executable, script_path] + sys.argv[1:])
            # Pass detected settings to the root process
            args.append(f'--color-scheme-from-user={user_scheme}')
            args.append(f'--button-layout-from-user={user_button_layout}')
            
            os.execvp('pkexec', args)
        except Exception as e:
            print(f"Failed to relaunch with pkexec: {e}", file=sys.stderr)
            
            # Use local import to avoid issues if dependencies are missing at top level
            # causing immediate crash before this handler
            try:
                from utils.translation import _
            except ImportError:
                _ = lambda x: x

            error_app = Gtk.Application()
            def show_error(app):
                dialog = Gtk.MessageDialog(
                    transient_for=None, modal=True, message_type=Gtk.MessageType.ERROR,
                    buttons=Gtk.ButtonsType.OK, text=_("Root privileges are required."),
                    secondary_text=_("This application must be run as root to function correctly. Please run it with sudo or pkexec.")
                )
                dialog.connect("response", lambda d, r: app.quit())
                dialog.show()
            error_app.connect('activate', show_error)
            error_app.run(None)
            sys.exit(1)

    # --- SETTINGS APPLICATION (RUNNING AS ROOT) ---
    
    # Force Cairo renderer for compatibility, especially in Live environments
    os.environ['GSK_RENDERER'] = 'cairo'
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--color-scheme-from-user', type=str, default='default')
    parser.add_argument('--button-layout-from-user', type=str, default=None)
    args, unknown = parser.parse_known_args()

    # Apply color scheme
    scheme_map = {
        'default': Adw.ColorScheme.DEFAULT,
        'prefer-light': Adw.ColorScheme.FORCE_LIGHT,
        'prefer-dark': Adw.ColorScheme.FORCE_DARK,
    }
    color_scheme_to_apply = scheme_map.get(args.color_scheme_from_user, Adw.ColorScheme.DEFAULT)
    if color_scheme_to_apply != Adw.ColorScheme.DEFAULT:
        style_manager = Adw.StyleManager.get_default()
        style_manager.set_color_scheme(color_scheme_to_apply)

    # Apply button layout
    if args.button_layout_from_user:
        gtk_settings = Gtk.Settings.get_default()
        gtk_settings.set_property('gtk-decoration-layout', args.button_layout_from_user)

    try:
        app = ManjaroRescueApplication()
        # Pass only the filename to avoid GTK seeing argparse flags
        return app.run([sys.argv[0]])
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        
        # Import translation here to ensure it's available after potential path setup
        from utils.translation import _
        
        error_message = _("Failed to load main application:\n") + str(e)
        
        class FallbackApp(Adw.Application):
            def __init__(self, error_msg):
                super().__init__(application_id="org.manjaro.rescue.fallback", flags=Gio.ApplicationFlags.NON_UNIQUE)
                self.error_msg = error_msg
            
            def do_activate(self):
                window = Adw.ApplicationWindow(application=self)
                window.set_title(_("Manjaro Rescue - Error"))
                window.set_default_size(600, 400)
                
                status_page = Adw.StatusPage()
                status_page.set_title(_("Application Error"))
                status_page.set_description(self.error_msg)
                status_page.set_icon_name("dialog-error-symbolic")
                
                window.set_content(status_page)
                window.present()
        
        fallback_app = FallbackApp(error_message)
        return fallback_app.run([sys.argv[0]])

if __name__ == "__main__":
    sys.exit(main())
