#!/usr/bin/env python3
"""Waybar disk module: shows disk usage across all mounted partitions."""

import argparse
import json
import os

# Filesystem types to ignore (virtual / non-disk)
IGNORE_FSTYPES = {
    "tmpfs", "devtmpfs", "proc", "sysfs", "cgroup", "cgroup2",
    "pstore", "bpf", "debugfs", "tracefs", "securityfs", "configfs",
    "hugetlbfs", "devpts", "ramfs", "fuse.gvfsd-fuse", "fusectl",
    "mqueue", "rpc_pipefs", "overlay", "squashfs",
}

# Pango color thresholds
COLOR_LOW = "#a6e3a1"    # green  — < 70%
COLOR_MID = "#f9e2af"    # yellow — 70–89%
COLOR_HIGH = "#f38ba8"   # red    — ≥ 90%

DEFAULT_ICON_SIZE = "24pt"


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--icon-size", default=DEFAULT_ICON_SIZE,
                    help="Icon size (Pango size, e.g. 24pt)")
    return p.parse_args()


def _pct_color(pct: float) -> str:
    """Return a Pango color string for the given usage percentage."""
    if pct >= 90:
        return COLOR_HIGH
    if pct >= 70:
        return COLOR_MID
    return COLOR_LOW


def _span(color: str, text: str) -> str:
    """Wrap text in a Pango <span> with the given foreground color."""
    return f"<span foreground='{color}'>{text}</span>"


def get_partitions():
    """Return list of (device, mountpoint, fstype) for real partitions."""
    partitions = []
    seen_devices = set()

    try:
        with open("/proc/mounts") as f:
            for line in f:
                parts = line.split()
                if len(parts) < 3:
                    continue
                device, mountpoint, fstype = parts[:3]

                if fstype in IGNORE_FSTYPES:
                    continue
                if not device.startswith("/dev/"):
                    continue
                if device in seen_devices:
                    continue
                seen_devices.add(device)

                partitions.append((device, mountpoint, fstype))
    except OSError:
        pass

    return partitions


def get_usage(mountpoint):
    """Return (total_bytes, used_bytes, pct) for a mountpoint, or None."""
    try:
        st = os.statvfs(mountpoint)
        total = st.f_frsize * st.f_blocks
        free = st.f_frsize * st.f_bfree
        used = total - free
        pct = (used / total) * 100 if total > 0 else 0.0
        return total, used, pct
    except OSError:
        return None


def fmt_size(bytes_val):
    """Format bytes as human-readable."""
    if bytes_val >= 1 << 40:
        return f"{bytes_val / (1 << 40):.1f}T"
    if bytes_val >= 1 << 30:
        return f"{bytes_val / (1 << 30):.1f}G"
    if bytes_val >= 1 << 20:
        return f"{bytes_val / (1 << 20):.1f}M"
    if bytes_val >= 1 << 10:
        return f"{bytes_val / (1 << 10):.1f}K"
    return f"{bytes_val}B"


def main():
    args = parse_args()
    isize = args.icon_size
    partitions = get_partitions()

    if not partitions:
        print(json.dumps(
            {"text": "  —", "tooltip": "No disk partitions found"},
            ensure_ascii=False))
        return

    tooltip_lines = []
    root_pct = None
    bar_parts = []

    for device, mountpoint, fstype in partitions:
        usage = get_usage(mountpoint)
        if usage is None:
            continue
        total, used, pct = usage

        total_str = fmt_size(total)
        used_str = fmt_size(used)
        color = _pct_color(pct)

        # Icons per mountpoint
        if mountpoint == "/":
            icon = ""
        elif mountpoint == "/home":
            icon = ""
        elif mountpoint == "/boot":
            icon = ""
        elif "/media/" in mountpoint or "/run/media/" in mountpoint:
            icon = ""
        else:
            icon = ""

        # Tooltip line with colored percentage
        pct_span = _span(color, f"{pct:5.1f}%")
        warn = " ⚠" if pct >= 90 else ""
        tooltip_lines.append(
            f"{icon}  {mountpoint:20s}  {used_str:>6s} / {total_str:>6s}  ({pct_span}){warn}"
        )

        # Track root for bar display
        if mountpoint == "/":
            root_pct = pct

    # Build bar text using Pango markup
    if root_pct is not None:
        bar_icon = "⚠" if root_pct >= 90 else ""
        bar_color = _pct_color(root_pct)
        bar_text = (
            f"<span size='{isize}' rise='-1500'>{bar_icon}</span>"
            f"  {_span(bar_color, f'{root_pct:.0f}%')}"
        )
    elif tooltip_lines:
        bar_text = f"<span size='{isize}' rise='-1500'></span>  —"
    else:
        bar_text = "  —"

    tooltip = "\n".join(tooltip_lines) if tooltip_lines else "No partitions"

    print(json.dumps(
        {"text": bar_text, "tooltip": tooltip},
        ensure_ascii=False))


if __name__ == "__main__":
    main()