import gi
import re
import json
import os
import shutil
import subprocess
import threading
from typing import List, Dict, Any, Optional

gi.require_version("Gtk", "4.0")
from gi.repository import Gtk, Gdk, GLib

FS_CHOICES = ["ext4", "btrfs", "xfs"]
BTRFS_DEFAULT_SUBVOLS = [
    ("@", "/", True),
    ("@home", "/home", True),
    ("@snapshots", "/.snapshots", True),
    ("@var_log", "/var/log", True),
    ("@var_cache", "/var/cache", True),
    ("@var_tmp", "/var/tmp", True),
    ("@tmp", "/tmp", True),
    ("@opt", "/opt", False),
    ("@srv", "/srv", False),
    ("@var_lib_libvirt_images", "/var/lib/libvirt/images", False),
]

MODE_ERASE = "erase"
MODE_EXISTING = "existing"
MODE_FREE = "free"

class PartitioningView(Gtk.Box):
    def __init__(self):
        super().__init__(orientation=Gtk.Orientation.VERTICAL, spacing=8)
        self.set_margin_top(6)
        self.set_margin_bottom(6)
        self.set_margin_start(6)
        self.set_margin_end(6)
        
        self.is_uefi = os.path.isdir("/sys/firmware/efi")
        self.disks = []
        self.selected_disk = None
        self.mode = MODE_ERASE
        self.free_space_mib = 0
        self.partition_paths = []
        
        # Lock für Thread-Sicherheit
        self.scan_lock = threading.Lock()

        self._setup_ui()
        
        # Laden der Disks im Hintergrund starten, damit die GUI nicht einfriert
        threading.Thread(target=self._load_disks_bg, daemon=True).start()
        
        self._update_mode_visibility()

    def _set_margins(self, widget, value=12):
        widget.set_margin_top(value)
        widget.set_margin_bottom(value)
        widget.set_margin_start(value)
        widget.set_margin_end(value)

    def _setup_ui(self):
        title = Gtk.Label.new(_("Festplatten-Partitionierung"))
        title.add_css_class("title-1")
        title.set_halign(Gtk.Align.START)
        self.append(title)
        
        sc = Gtk.ScrolledWindow()
        sc.set_vexpand(True)
        sc.set_policy(Gtk.PolicyType.NEVER, Gtk.PolicyType.AUTOMATIC)
        self.append(sc)

        main_box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=12)
        sc.set_child(main_box)

        # 1. Disk Selection
        disk_frame = Gtk.Frame(label=_("1. Ziel-Datenträger"))
        disk_frame.add_css_class("card")
        main_box.append(disk_frame)

        disk_vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=8)
        self._set_margins(disk_vbox)
        disk_frame.set_child(disk_vbox)

        disk_row = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=8)
        disk_row.append(Gtk.Label.new(_("Festplatte:")))
        self.device_combo = Gtk.ComboBoxText()
        self.device_combo.set_hexpand(True)
        self.device_combo.connect("changed", self._on_device_changed)
        disk_row.append(self.device_combo)
        
        self.btn_refresh = Gtk.Button.new_from_icon_name("view-refresh-symbolic")
        self.btn_refresh.connect("clicked", lambda _: threading.Thread(target=self._load_disks_bg, daemon=True).start())
        disk_row.append(self.btn_refresh)
        disk_vbox.append(disk_row)

        self.btn_gparted = Gtk.Button.new_with_label(_("GParted öffnen"))
        self.btn_gparted.add_css_class("suggested-action")
        self.btn_gparted.connect("clicked", self._on_open_gparted)
        
        self.gparted_spinner = Gtk.Spinner()
        g_box = Gtk.Box(spacing=8)
        g_box.append(self.btn_gparted)
        g_box.append(self.gparted_spinner)
        disk_vbox.append(g_box)

        # 2. Modes
        mode_frame = Gtk.Frame(label=_("2. Installations-Methode"))
        mode_frame.add_css_class("card")
        main_box.append(mode_frame)

        mode_box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=10)
        self._set_margins(mode_box)
        mode_frame.set_child(mode_box)

        self.rb_erase = Gtk.CheckButton.new_with_label(_("Gesamte Festplatte löschen"))
        self.rb_free = Gtk.CheckButton.new_with_label(_("Freien Speicherplatz verwenden (Dual-Boot)"))
        self.rb_existing = Gtk.CheckButton.new_with_label(_("Vorhandene Partition manuell wählen"))
        self.rb_free.set_group(self.rb_erase)
        self.rb_existing.set_group(self.rb_erase)
        self.rb_erase.set_active(True)

        for rb in [self.rb_erase, self.rb_free, self.rb_existing]:
            rb.connect("toggled", self._on_mode_changed)
            mode_box.append(rb)

        # 3. Size details
        self.mode_detail_box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=12)
        main_box.append(self.mode_detail_box)

        self.size_options_frame = Gtk.Frame(label=_("Speicher-Zuweisung"))
        self.size_options_frame.add_css_class("card")
        size_vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=10)
        self._set_margins(size_vbox)
        
        self.lbl_free_info = Gtk.Label()
        self.lbl_free_info.set_halign(Gtk.Align.START)
        size_vbox.append(self.lbl_free_info)
        
        # Spinner für den Free-Space Scan
        self.scan_spinner = Gtk.Spinner()
        size_vbox.append(self.scan_spinner)

        row_root = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=12)
        row_root.append(Gtk.Label.new(_("Größe für Root (/) (GiB):")))
        self.spin_root_size = Gtk.SpinButton.new_with_range(10, 10000, 1)
        self.spin_root_size.set_value(40)
        row_root.append(self.spin_root_size)
        size_vbox.append(row_root)

        self.cb_sep_home = Gtk.CheckButton.new_with_label(_("Separate /home Partition erstellen"))
        self.cb_sep_home.connect("toggled", self._update_home_visibility)
        size_vbox.append(self.cb_sep_home)

        self.row_home_size = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=12)
        self.row_home_size.append(Gtk.Label.new(_("Größe für /home (GiB):")))
        self.spin_home_size = Gtk.SpinButton.new_with_range(1, 10000, 1)
        self.spin_home_size.set_value(20)
        self.row_home_size.append(self.spin_home_size)
        size_vbox.append(self.row_home_size)

        self.size_options_frame.set_child(size_vbox)
        self.mode_detail_box.append(self.size_options_frame)

        self.existing_frame = Gtk.Frame(label=_("Partition wählen"))
        self.existing_frame.add_css_class("card")
        ex_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=12)
        self._set_margins(ex_box)
        self.partition_model = Gtk.StringList.new([])
        self.partition_dropdown = Gtk.DropDown.new(self.partition_model, None)
        self.partition_dropdown.connect("notify::selected", self._on_partition_selected)
        ex_box.append(Gtk.Label.new(_("Partition:")))
        ex_box.append(self.partition_dropdown)
        self.existing_frame.set_child(ex_box)
        self.mode_detail_box.append(self.existing_frame)

        # 4. FS Options
        self.layout_frame = Gtk.Frame(label=_("3. Layout & Dateisystem"))
        self.layout_frame.add_css_class("card")
        layout_box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=12)
        self._set_margins(layout_box)
        self.layout_frame.set_child(layout_box)
        main_box.append(self.layout_frame)

        fs_row = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=12)
        fs_row.append(Gtk.Label.new(_("Dateisystem:")))
        self.fs_combo = Gtk.ComboBoxText()
        for fs in FS_CHOICES: self.fs_combo.append_text(fs)
        self.fs_combo.set_active(0)
        self.fs_combo.connect("changed", self._on_fs_changed)
        fs_row.append(self.fs_combo)
        layout_box.append(fs_row)

        self.cb_swap = Gtk.CheckButton.new_with_label(_("Swap-Partition anlegen"))
        self.cb_swap.connect("toggled", lambda cb: self.row_swap_size.set_visible(cb.get_active()))
        layout_box.append(self.cb_swap)

        self.row_swap_size = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=12)
        self.row_swap_size.append(Gtk.Label.new(_("Swap-Größe (GiB):")))
        self.spin_swap = Gtk.SpinButton.new_with_range(1, 128, 1)
        self.spin_swap.set_value(4)
        self.row_swap_size.append(self.spin_swap)
        self.row_swap_size.set_visible(False)
        layout_box.append(self.row_swap_size)

        self.subvol_frame = Gtk.Frame(label=_("Btrfs-Subvolumes"))
        self.subvol_frame.add_css_class("card")
        self.subvol_list = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=4)
        self._set_margins(self.subvol_list)
        self.subvol_frame.set_child(self.subvol_list)
        main_box.append(self.subvol_frame)
        self._build_subvols_ui()

    def _load_disks_bg(self):
        # Disks im Hintergrund laden
        found_disks = []
        try:
            out = subprocess.check_output(["lsblk", "-p", "-J", "-o", "NAME,SIZE,TYPE,PATH"], text=True)
            devices = json.loads(out).get("blockdevices", [])
            for dev in devices:
                if dev.get("type") == "disk":
                    found_disks.append(dev)
        except: pass
        
        GLib.idle_add(self._update_disk_combo, found_disks)

    def _update_disk_combo(self, found_disks):
        self.device_combo.remove_all()
        self.disks = found_disks
        for dev in self.disks:
            self.device_combo.append_text(f"{dev['path']} ({dev['size']})")
        
        if self.disks:
            self.device_combo.set_active(0)
        else:
            self.lbl_free_info.set_text(_("Keine Festplatten gefunden."))

    def _update_free_space_info(self):
        if not self.selected_disk: return
        
        self.lbl_free_info.set_text(_("Analysiere Festplatte..."))
        self.spin_root_size.set_sensitive(False)
        self.scan_spinner.start()
        
        # Scan im Thread starten
        threading.Thread(target=self._scan_free_space_task, args=(self.selected_disk,), daemon=True).start()

    def _scan_free_space_task(self, disk_path):
        size_mib = 0
        try:
            from installation_backend import InstallationBackend
            backend = InstallationBackend()
            segs = backend._parted_free_segments_mib(disk_path)
            if segs:
                unused_start, unused_end, s_mib = max(segs, key=lambda t: t[2])
                size_mib = s_mib
        except Exception as e:
            print(f"Fehler beim Scan: {e}")
        
        GLib.idle_add(self._on_scan_finished, size_mib)

    def _on_scan_finished(self, size_mib):
        self.scan_spinner.stop()
        self.free_space_mib = size_mib
        size_gib = size_mib / 1024.0
        
        if size_gib > 0:
            self.lbl_free_info.set_markup(_("Verfügbarer Platz: <b>{:.2f} GiB</b>").format(size_gib))
            max_val = int(size_gib)
            if max_val >= 10:
                self.spin_root_size.set_range(10, max_val)
                self.spin_root_size.set_value(min(40, max_val))
                self.spin_root_size.set_sensitive(True)
            else:
                self.spin_root_size.set_range(0, max_val)
                self.spin_root_size.set_value(max_val)
                self.spin_root_size.set_sensitive(True)
        else:
            self.lbl_free_info.set_text(_("Kein unpartitionierter Platz gefunden."))
            self.spin_root_size.set_sensitive(False)

    def _on_device_changed(self, combo):
        idx = combo.get_active()
        if idx >= 0:
            self.selected_disk = self.disks[idx]["path"]
            if self.mode == MODE_FREE:
                self._update_free_space_info()
            self._refresh_partitions()

    def _on_mode_changed(self, rb):
        if not rb.get_active(): return
        if self.rb_erase.get_active(): self.mode = MODE_ERASE
        elif self.rb_free.get_active(): self.mode = MODE_FREE
        elif self.rb_existing.get_active(): self.mode = MODE_EXISTING
        self._update_mode_visibility()
        if self.mode == MODE_FREE: self._update_free_space_info()

    def _update_mode_visibility(self):
        is_ex = (self.mode == MODE_EXISTING)
        is_fr = (self.mode == MODE_FREE)
        is_er = (self.mode == MODE_ERASE)
        self.size_options_frame.set_visible(is_fr or is_er)
        self.existing_frame.set_visible(is_ex)
        self.layout_frame.set_visible(not is_ex)
        self.spin_root_size.set_sensitive(is_fr) 
        self._update_home_visibility()
        self._on_fs_changed(self.fs_combo)

    def _update_home_visibility(self, *args):
        active = self.cb_sep_home.get_active()
        self.row_home_size.set_visible(active and self.size_options_frame.get_visible())

    def _on_fs_changed(self, combo):
        self.fs_choice = combo.get_active_text() or "ext4"
        self.subvol_frame.set_visible(self.fs_choice == "btrfs" and self.mode != MODE_EXISTING)

    def _on_open_gparted(self, widget):
        if shutil.which("gparted"):
            subprocess.Popen(["gparted", self.selected_disk or ""])
        else:
            self.btn_gparted.set_sensitive(False)
            self.gparted_spinner.start()
            threading.Thread(target=self._install_gparted_task, daemon=True).start()

    def _install_gparted_task(self):
        res = subprocess.run(["xbps-install", "-y", "gparted"])
        GLib.idle_add(self._on_gparted_install_done, res.returncode == 0)

    def _on_gparted_install_done(self, success):
        self.gparted_spinner.stop()
        self.btn_gparted.set_sensitive(True)
        if success:
            subprocess.Popen(["gparted", self.selected_disk or ""])

    def _on_partition_selected(self, dropdown, pspec):
        idx = dropdown.get_selected()
        if idx != Gtk.INVALID_LIST_POSITION and idx < len(self.partition_paths):
            self.selected_partition = self.partition_paths[idx]
        else:
            self.selected_partition = None

    def _refresh_partitions(self):
        # Auch partitions scan im Thread machen um UI Blockaden zu vermeiden
        threading.Thread(target=self._refresh_partitions_bg, daemon=True).start()

    def _refresh_partitions_bg(self):
        if not self.selected_disk: return
        paths = []
        model_items = []
        try:
            out = subprocess.check_output(["lsblk", "-p", "-J", "-o", "NAME,SIZE,FSTYPE,TYPE,PATH,MOUNTPOINT", self.selected_disk], text=True)
            data = json.loads(out).get("blockdevices", [])
            for disk in data:
                for part in disk.get("children", []):
                    if part.get("type") == "part" and not part.get("mountpoint"):
                        path = part["path"]
                        paths.append(path)
                        model_items.append(f"{path} ({part['size']}, {part.get('fstype') or 'empty'})")
        except: pass
        GLib.idle_add(self._update_partition_dropdown, paths, model_items)

    def _update_partition_dropdown(self, paths, items):
        self.partition_paths = paths
        self.partition_model = Gtk.StringList.new(items)
        self.partition_dropdown.set_model(self.partition_model)
        if self.partition_paths:
            self.partition_dropdown.set_selected(0)

    def _build_subvols_ui(self):
        self.subvol_rows = []
        for name, mnt, active in BTRFS_DEFAULT_SUBVOLS:
            cb = Gtk.CheckButton.new_with_label(f"{name} -> {mnt}")
            cb.set_active(active)
            self.subvol_list.append(cb)
            self.subvol_rows.append((name, cb))

    def get_plan(self) -> Dict[str, Any]:
        plan = {
            "device": self.selected_disk,
            "mode": self.mode,
            "uefi": self.is_uefi,
            "auto_layout": {
                "filesystem": self.fs_choice,
                "use_swap_partition": self.cb_swap.get_active(),
                "swap_partition_gib": int(self.spin_swap.get_value()),
                "use_separate_home": self.cb_sep_home.get_active(),
                "home_size_gib": int(self.spin_home_size.get_value()),
            }
        }
        if self.mode in [MODE_FREE, MODE_ERASE]:
            plan["free_space"] = {"root_size_gib": int(self.spin_root_size.get_value())}
        elif self.mode == MODE_EXISTING:
            plan["existing_partition"] = self.selected_partition
            plan["format_existing"] = True
            
        if self.fs_choice == "btrfs":
            plan["auto_layout"]["subvolumes"] = [n for n, cb in self.subvol_rows if cb.get_active()]
        return plan

    def validate_plan(self):
        if not self.selected_disk: raise ValueError(_("Bitte wähle eine Festplatte aus."))
        if self.mode == MODE_EXISTING and not self.selected_partition:
            raise ValueError(_("Bitte wähle eine Partition aus."))
        if self.mode == MODE_FREE:
            # Wir verlassen uns auf den gecachten Wert, oder prüfen kurz
            needed_gib = int(self.spin_root_size.get_value())
            if self.cb_sep_home.get_active(): needed_gib += int(self.spin_home_size.get_value())
            if self.cb_swap.get_active(): needed_gib += int(self.spin_swap.get_value())
            
            # Sicherheitscheck, falls Scan noch nicht fertig war
            if self.free_space_mib == 0:
                 pass # Warnung könnte hier optional sein
            elif (needed_gib * 1024) > self.free_space_mib:
                raise ValueError(_("Gewählter Platz überschreitet verfügbaren Speicher!"))