import json
import os
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;
}
"""


def get_total_memory() -> int:
    """Returns total system memory in bytes."""
    try:
        return os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
    except (ValueError, OSError):
        with open("/proc/meminfo", "r") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    parts = line.split()
                    return int(parts[1]) * 1024
    return 0


def calculate_auto_partitions(disk_device):
    """
    Generates an automatic partition layout.
      - 2GB EFI
      - RAM + 2GB Swap
      - Rest Root (ext4)
    """
    disk_size = 0
    try:
        # Get disk size in bytes
        result = subprocess.run(
            ["lsblk", "-J", "-b", "-o", "NAME,SIZE,TYPE"],
            capture_output=True,
            text=True,
        )
        if result.returncode == 0:
            data = json.loads(result.stdout)
            devices = data.get("blockdevices", [])
            dev_name = disk_device.replace("/dev/", "")

            for dev in devices:
                if dev.get("name") == dev_name:
                    disk_size = int(dev.get("size", 0))
                    break
    except Exception as e:
        print(f"Error getting disk size for auto partitioning: {e}")
        return []

    if disk_size == 0:
        return []

    ram_size = get_total_memory()

    # Sizes in bytes
    efi_size = 2 * 1024 * 1024 * 1024
    swap_size = ram_size + (2 * 1024 * 1024 * 1024)

    # Check if disk is large enough
    min_root = 10 * 1024 * 1024 * 1024  # minimum 10GB for root
    if disk_size < (efi_size + swap_size + min_root):
        print("Disk too small for automatic partitioning scheme.")
        return []

    root_size = disk_size - efi_size - swap_size

    partitions = [
        {
            "type": "partition",
            "name": "EFI System",
            "filesystem": "vfat",
            "mount_point": "/boot/efi",
            "size": f"{efi_size / (1024**3):.1f} GB",
            "bytes": efi_size,
            "style_class": "part-efi",
        },
        {
            "type": "partition",
            "name": "Swap",
            "filesystem": "swap",
            "mount_point": "[SWAP]",
            "size": f"{swap_size / (1024**3):.1f} GB",
            "bytes": swap_size,
            "style_class": "part-swap",
        },
        {
            "type": "partition",
            "name": "Root",
            "filesystem": "ext4",
            "mount_point": "/",
            "size": f"{root_size / (1024**3):.1f} GB",
            "bytes": root_size,
            "style_class": "part-root",
        },
    ]

    return partitions


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):
        self.current_disk_path = disk_device
        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") or "").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") or "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
        self.selected_disk_path = self.current_disk_path # Assumes we store this

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

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

        def add_menu_item(label, icon_name, callback, destructive=False):
            btn = Gtk.Button(label=label)
            btn.add_css_class("flat")
            if destructive:
                btn.add_css_class("destructive-action")
            
            content = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, 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)
            
            btn.connect("clicked", lambda b: [popover.popdown(), callback()])
            menu_box.append(btn)

        if data.get("type") == "partition":
            add_menu_item("Select Mount Point", "folder-open-symbolic", 
                         lambda: self.select_mount_point(data))
            separator = Gtk.Separator()
            menu_box.append(separator)
            add_menu_item("Delete", "user-trash-symbolic", 
                         lambda: self.delete_part(data), destructive=True)

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

        popover.popup()

    def select_mount_point(self, data):
        win = Adw.Window(title="Select Mount Point", modal=True, transient_for=self.get_root())
        box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=12, margin_top=24, margin_bottom=24, margin_start=24, margin_end=24)
        win.set_content(box)

        options = ["/", "/boot/efi", "[SWAP]", "None"]
        dropdown = Gtk.DropDown.new_from_strings(options)
        box.append(Gtk.Label(label=f"Mount point for {data['name']}:"))
        box.append(dropdown)

        def on_apply(b):
            selected = options[dropdown.get_selected()]
            data["mount_point"] = selected if selected != "None" else ""
            win.close()
            self.refresh_bar()

        btn = Gtk.Button(label="Apply")
        btn.add_css_class("suggested-action")
        btn.connect("clicked", on_apply)
        box.append(btn)
        win.present()

    def delete_part(self, data):
        from ...backend.disk import delete_partition
        # Extract partition number from name (e.g., nvme0n1p3 -> 3)
        import re
        match = re.search(r"(\d+)$", data["name"])
        if match:
            part_num = int(match.group(1))
            delete_partition(self.current_disk_path, part_num)
            self.load_partitions(self.current_disk_path)

    def create_part_dialog(self, data):
        win = Adw.Window(title="Create Partition", modal=True, transient_for=self.get_root())
        box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=12, margin_top=24, margin_bottom=24, margin_start=24, margin_end=24)
        win.set_content(box)

        # Name entry
        name_entry = Gtk.Entry(text="Linux filesystem")
        box.append(Gtk.Label(label="Partition Name:"))
        box.append(name_entry)

        # Size entry
        size_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=12)
        size_entry = Gtk.Entry(text=str(int(data["bytes"] / (1024*1024))))
        size_box.append(size_entry)
        size_box.append(Gtk.Label(label="MB"))
        box.append(Gtk.Label(label="Size:"))
        box.append(size_box)

        # Type dropdown
        types = {"Root (Linux)": "8300", "EFI System": "ef00", "Swap": "8200"}
        type_names = list(types.keys())
        type_dropdown = Gtk.DropDown.new_from_strings(type_names)
        box.append(Gtk.Label(label="Partition Type:"))
        box.append(type_dropdown)

        # Filesystem dropdown
        fs_options = ["ext4", "fat32", "swap"]
        fs_dropdown = Gtk.DropDown.new_from_strings(fs_options)
        box.append(Gtk.Label(label="Filesystem:"))
        box.append(fs_dropdown)

        # Auto-update FS and Name based on Type
        def on_type_changed(dropdown, pspec):
            selected_type = type_names[dropdown.get_selected()]
            if selected_type == "EFI System":
                fs_dropdown.set_selected(1) # fat32
                name_entry.set_text("EFI System")
            elif selected_type == "Swap":
                fs_dropdown.set_selected(2) # swap
                name_entry.set_text("Swap")
            else:
                fs_dropdown.set_selected(0) # ext4
                name_entry.set_text("Linux filesystem")

        type_dropdown.connect("notify::selected", on_type_changed)

        error_label = Gtk.Label(label="")
        error_label.add_css_class("error")
        box.append(error_label)

        def on_create(b):
            from ...backend.disk import create_partition
            name = name_entry.get_text()
            selected_type = type_names[type_dropdown.get_selected()]
            
            if selected_type == "EFI System" and name == "Linux filesystem":
                error_label.set_text("Can't set partition name to Linux filesystem")
                return

            size_mb = int(size_entry.get_text())
            type_code = types[selected_type]
            fstype = fs_options[fs_dropdown.get_selected()]
            
            create_partition(self.current_disk_path, size_mb, type_code, name, fstype)
            win.close()
            self.load_partitions(self.current_disk_path)

        btn = Gtk.Button(label="Create")
        btn.add_css_class("suggested-action")
        btn.connect("clicked", on_create)
        box.append(btn)
        win.present()

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