import json
import subprocess

import gi

gi.require_version("Gtk", "4.0")
gi.require_version("Adw", "1")
from gi.repository import Adw, Gdk, Gio, GObject, Gtk

CSS = """
.part-segment {
    border-radius: 0;
    box-shadow: none;
    border: none;
}
.part-segment:first-child {
    border-top-left-radius: 6px;
    border-bottom-left-radius: 6px;
}
.part-segment:last-child {
    border-top-right-radius: 6px;
    border-bottom-right-radius: 6px;
}

.part-efi {
    background-color: #3584e4; /* Blue */
    color: white;
}
.part-root {
    background-color: #2ec27e; /* Green */
    color: white;
}
.part-swap {
    background-color: #e66100; /* Orange */
    color: white;
}
.part-empty {
    background-color: #deddda; /* Gray */
    color: black;
}
"""


class PartitionSegment(Gtk.Button):
    def __init__(self, part_data, page, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.part_data = part_data
        self.page = page

        self.add_css_class("part-segment")
        self.add_css_class(part_data["style_class"])

        self.set_hexpand(part_data.get("expand", False))

        if "width_request" in part_data:
            self.set_size_request(part_data["width_request"], -1)

        if part_data.get("label"):
            lbl = Gtk.Label(label=part_data["label"])
            lbl.set_ellipsize(3)
            self.set_child(lbl)

        self.connect("clicked", self.on_clicked)

        right_click = Gtk.GestureClick()
        right_click.set_button(3)
        right_click.connect("pressed", self.on_right_click)
        self.add_controller(right_click)

    def on_clicked(self, button):
        self.page.show_details_window(self.part_data)

    def on_right_click(self, gesture, n_press, x, y):
        self.page.show_context_menu(self, x, y)


class PartitioningPage(Adw.Bin):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        provider = Gtk.CssProvider()
        provider.load_from_data(CSS.encode())
        Gtk.StyleContext.add_provider_for_display(
            Gdk.Display.get_default(), provider, Gtk.STYLE_PROVIDER_PRIORITY_APPLICATION
        )

        clamp = Adw.Clamp()
        clamp.set_maximum_size(800)
        self.set_child(clamp)

        box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL)
        box.set_spacing(24)
        box.set_valign(Gtk.Align.CENTER)
        box.set_margin_top(24)
        box.set_margin_bottom(24)
        clamp.set_child(box)

        title = Gtk.Label(label="Manual Partitioning")
        title.add_css_class("title-1")
        box.append(title)

        descr = Gtk.Label(
            label="Modify the partition layout below. Right-click segments for options."
        )
        descr.set_wrap(True)
        box.append(descr)

        disk_info_group = Adw.PreferencesGroup()
        disk_info_group.set_title("Selected Drive")
        box.append(disk_info_group)

        row = Adw.ActionRow()
        row.set_title("Current Disk")
        row.set_subtitle("Detecting...")
        row.set_icon_name("drive-harddisk-symbolic")
        disk_info_group.add(row)
        self.disk_row = row

        bar_box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL)
        bar_box.set_spacing(12)
        box.append(bar_box)

        bar_label = Gtk.Label(label="Partition Layout", xalign=0)
        bar_label.add_css_class("heading")
        bar_box.append(bar_label)

        self.visual_bar = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL)
        self.visual_bar.add_css_class("card")
        self.visual_bar.set_overflow(Gtk.Overflow.HIDDEN)
        self.visual_bar.set_size_request(-1, 80)

        bar_box.append(self.visual_bar)

        legend_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL)
        legend_box.set_spacing(12)
        legend_box.set_halign(Gtk.Align.CENTER)
        bar_box.append(legend_box)

        self.add_legend_item(legend_box, "part-efi", "EFI System")
        self.add_legend_item(legend_box, "part-root", "Root Filesystem")
        self.add_legend_item(legend_box, "part-empty", "Free Space")

        self.partitions = []
        # Initially empty until loaded with a specific disk

    def load_partitions(self, disk_device=None):
        target_disk = None

        try:
            result = subprocess.run(
                ["lsblk", "-J", "-b", "-o", "NAME,SIZE,TYPE,FSTYPE,MOUNTPOINT,PKNAME"],
                capture_output=True,
                text=True,
            )
            if result.returncode == 0:
                data = json.loads(result.stdout)
                devices = data.get("blockdevices", [])

                if disk_device:
                    # Look for the specific device (match name, e.g. "nvme0n1")
                    # disk_device from storage page is usually just the name like "nvme0n1" or "sda"
                    # But if it's full path /dev/sda, we need to strip /dev/
                    dev_name = disk_device.replace("/dev/", "")

                    for dev in devices:
                        if dev.get("name") == dev_name:
                            target_disk = dev
                            break

                # If no specific disk requested or not found, maybe fallback or do nothing
                # The previous logic found the first one. Let's keep strict if device is provided.
                if not target_disk and not disk_device:
                    for dev in devices:
                        if (
                            dev.get("type") == "disk"
                            and not dev.get("name").startswith("zram")
                            and not dev.get("name").startswith("loop")
                        ):
                            target_disk = dev
                            break

        except Exception as e:
            print(f"Error loading disks: {e}")

        if target_disk:
            self.disk_row.set_title(target_disk.get("name", "Unknown Disk"))
            self.disk_row.set_subtitle(
                f"Size: {self.format_size(target_disk.get('size', 0))}"
            )

            raw_parts = target_disk.get("children", [])
            self.partitions = []

            total_disk_size = int(target_disk.get("size", 0))
            covered_size = 0

            for p in raw_parts:
                p_size = int(p.get("size", 0))
                covered_size += p_size

                fstype = p.get("fstype", "").lower()
                style = "part-root"
                if "fat" in fstype or "efi" in fstype:
                    style = "part-efi"
                elif "swap" in fstype:
                    style = "part-swap"

                self.partitions.append(
                    {
                        "type": "partition",
                        "name": p.get("name"),
                        "filesystem": p.get("fstype", "Unknown"),
                        "label": p.get("name"),
                        "mount_point": p.get("mountpoint", ""),
                        "size": self.format_size(p_size),
                        "bytes": p_size,
                        "style_class": style,
                        "expand": False,
                    }
                )

            remaining = total_disk_size - covered_size
            if remaining > 100 * 1024 * 1024:
                self.partitions.append(
                    {
                        "type": "empty",
                        "name": "Unallocated",
                        "size": self.format_size(remaining),
                        "bytes": remaining,
                        "style_class": "part-empty",
                        "expand": False,
                    }
                )

            total_display_units = 1000
            for p in self.partitions:
                ratio = p["bytes"] / total_disk_size if total_disk_size > 0 else 0
                width = int(ratio * 800)
                if width < 50:
                    width = 50
                p["width_request"] = width

        else:
            self.partitions = [
                {
                    "type": "empty",
                    "name": "No Disk Selected",
                    "size": "-",
                    "style_class": "part-empty",
                    "width_request": 300,
                }
            ]

        self.refresh_bar()

    def format_size(self, size_bytes):
        size = float(size_bytes)
        for unit in ["B", "KB", "MB", "GB", "TB"]:
            if size < 1024.0:
                return f"{size:.1f} {unit}"
            size /= 1024.0
        return f"{size:.1f} PB"

    def add_legend_item(self, box, style_class, label_text):
        item = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL)
        item.set_spacing(6)

        swatch = Gtk.Box()
        swatch.set_size_request(16, 16)
        swatch.add_css_class(style_class)
        swatch.add_css_class("circular")

        label = Gtk.Label(label=label_text)

        item.append(swatch)
        item.append(label)
        box.append(item)

    def refresh_bar(self):
        child = self.visual_bar.get_first_child()
        while child:
            next_child = child.get_next_sibling()
            self.visual_bar.remove(child)
            child = next_child

        for part in self.partitions:
            segment = PartitionSegment(part, self)
            self.visual_bar.append(segment)

    def show_details_window(self, data):
        win = Adw.Window(title="Partition Details")
        win.set_transient_for(self.get_root())
        win.set_modal(True)
        win.set_default_size(300, 200)

        box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL)
        box.set_spacing(12)
        box.set_margin_top(24)
        box.set_margin_bottom(24)
        box.set_margin_start(24)
        box.set_margin_end(24)
        win.set_content(box)

        title = Gtk.Label(label=data.get("name", "Unknown"))
        title.add_css_class("title-2")
        box.append(title)

        grid = Gtk.Grid()
        grid.set_column_spacing(12)
        grid.set_row_spacing(6)
        grid.set_halign(Gtk.Align.CENTER)
        box.append(grid)

        rows = [
            ("Type", data.get("type", "-")),
            ("Size", data.get("size", "-")),
        ]
        if data.get("type") == "partition":
            rows.append(("Filesystem", data.get("filesystem", "-")))
            rows.append(("Mount Point", data.get("mount_point") or "None"))

        for i, (key, val) in enumerate(rows):
            k_lbl = Gtk.Label(label=key, xalign=1)
            k_lbl.add_css_class("dim-label")
            v_lbl = Gtk.Label(label=val, xalign=0)

            grid.attach(k_lbl, 0, i, 1, 1)
            grid.attach(v_lbl, 1, i, 1, 1)

        btn = Gtk.Button(label="Close")
        btn.connect("clicked", lambda b: win.close())
        btn.set_halign(Gtk.Align.CENTER)
        box.append(btn)

        win.present()

    def show_context_menu(self, widget, x, y):
        data = widget.part_data

        popover = Gtk.Popover()
        popover.set_parent(widget)

        menu_box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL)
        menu_box.set_spacing(0)
        popover.set_child(menu_box)

        def add_menu_item(label, icon_name=None, destructive=False):
            btn = Gtk.Button(label=label)
            btn.add_css_class("flat")
            btn.set_halign(Gtk.Align.FILL)
            if destructive:
                btn.add_css_class("destructive-action")

            if icon_name:
                content = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL)
                content.set_spacing(12)
                icon = Gtk.Image.new_from_icon_name(icon_name)
                lbl = Gtk.Label(label=label)
                content.append(icon)
                content.append(lbl)
                btn.set_child(content)

            menu_box.append(btn)
            return btn

        if data.get("type") == "partition":
            add_menu_item("Select Mount Point", "folder-open-symbolic")
            add_menu_item("Format", "drive-harddisk-symbolic")
            add_menu_item("Resize", "object-resize-symbolic")
            separator = Gtk.Separator()
            menu_box.append(separator)
            btn_del = add_menu_item("Delete", "user-trash-symbolic", destructive=True)
            btn_del.add_css_class("error")

        elif data.get("type") == "empty":
            add_menu_item("Create Partition", "list-add-symbolic")

        popover.popup()

    def get_config(self):
        return {"partitions": self.partitions}