"""Removable media section for the sidebar (Features/Major:
"Show mounted stuff, e.g. USB sticks being plugged in or out").

Talks to udisks2 over D-Bus (dbus-next, already used by PortalService.py),
the same service other file managers use:

- Plugged-in devices show up whether or not they're mounted yet. The old
  version only listed folders under /run/media/<user>, so a stick appeared
  only once something else (e.g. PCManFM) had mounted it, and only on the
  next 3-second poll.
- Changes arrive as udisks2 signals, so the list updates straight away.
- Clicking an unmounted device mounts it (no password for removable media)
  and opens it; the eject button unmounts and powers it off.
- With AutoMountDevices on, a newly plugged-in device is mounted right away.

Falls back to polling /run/media/<user> if udisks2 isn't reachable.
Everything runs on the page's own event loop (Page.run_task), so calling
.update() from here is safe."""

from __future__ import annotations  # the D-Bus type hints below don't exist on Windows

import asyncio
import getpass
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Optional

import flet as ft

IsWindows = sys.platform == "win32"
if not IsWindows:
    # udisks2 is Linux-only; Windows lists its drives itself (ListWindowsDrives).
    from dbus_next import BusType, Message, MessageType, Variant
    from dbus_next.aio import MessageBus

from Core import AppConfig, NotificationStyle, Palette

UDisks = "org.freedesktop.UDisks2"
UDisksPath = "/org/freedesktop/UDisks2"
BlockInterface = f"{UDisks}.Block"
FilesystemInterface = f"{UDisks}.Filesystem"
DriveInterface = f"{UDisks}.Drive"

PollInterval = 3.0  # seconds, for Windows and the no-udisks2 fallback
MountRoots = [Path("/run/media") / getpass.getuser(), Path("/media") / getpass.getuser()]


@dataclass
class Device:
    Key: str  # udisks2 block object path, or the mount path in fallback mode
    Label: str
    MountPoint: Optional[Path]
    DrivePath: str = ""
    CanPowerOff: bool = False
    # True only for an internal drive mounted at a system path ("/", "/home",
    # ...) shown because of Config.ShowSystemDrive — _BuildTiles never shows
    # an eject button for one of these, no matter what CanPowerOff says.
    # Unmounting your own running root filesystem would crash the system.
    IsSystemMount: bool = False


def _HumanSize(Bytes: int) -> str:
    Size = float(Bytes)
    for Unit in ("B", "KB", "MB", "GB", "TB"):
        if Size < 1000 or Unit == "TB":
            return f"{Size:.0f} {Unit}" if Unit == "B" else f"{Size:.1f} {Unit}"
        Size /= 1000
    return f"{Bytes} B"


def _Value(Properties: dict, Name: str, Default=None):
    Item = Properties.get(Name)
    return Item.value if isinstance(Item, Variant) else Default


def _UserMount(MountPoint: Path) -> bool:
    return any(Root in MountPoint.parents for Root in MountRoots)


def ListWindowsDrives(ShowInternal: bool, ShowSystemDrive: bool) -> list[Device]:
    """Windows: removable drives (USB sticks, memory cards, discs with a disc
    in), plus fixed and network drives with ShowInternal, and the system drive
    (C:) with ShowSystemDrive. Windows mounts drives itself, so there's nothing
    to mount; ejecting is left to Windows' own "Safely remove"."""
    import ctypes

    Kernel = ctypes.windll.kernel32
    Removable, Fixed, Remote, CdRom = 2, 3, 4, 5
    SystemDrive = (os.environ.get("SystemDrive") or "C:").upper().rstrip("\\") + "\\"
    Devices: list[Device] = []
    Mask = Kernel.GetLogicalDrives()
    for Index in range(26):
        if not Mask & (1 << Index):
            continue
        Root = f"{chr(65 + Index)}:\\"
        Kind = Kernel.GetDriveTypeW(Root)
        IsSystem = Root.upper() == SystemDrive
        if IsSystem and not ShowSystemDrive:
            continue
        if Kind in (Fixed, Remote) and not (ShowInternal or IsSystem):
            continue
        if Kind not in (Removable, Fixed, Remote, CdRom):
            continue
        Name = ctypes.create_unicode_buffer(261)
        # Fails for an empty card reader or disc drive: nothing to show then.
        if not Kernel.GetVolumeInformationW(Root, Name, 261, None, None, None, None, 0):
            continue
        Fallback = {Removable: "Removable drive", CdRom: "Disc", Remote: "Network drive"}.get(Kind, "Local disk")
        Devices.append(Device(Root, f"{Name.value or Fallback} ({Root[:2]})", Path(Root), IsSystemMount=IsSystem))
    return Devices


def ListFallbackDevices() -> list[Device]:
    Devices: list[Device] = []
    for Root in MountRoots:
        if not Root.is_dir():
            continue
        try:
            Devices += [Device(str(P), P.name, P) for P in Root.iterdir() if P.is_dir()]
        except OSError:
            pass
    return sorted(Devices, key=lambda D: D.Label.lower())


class UDisksClient:
    def __init__(self, Bus: MessageBus):
        self.Bus = Bus

    @classmethod
    async def Connect(cls) -> "UDisksClient":
        return cls(await MessageBus(bus_type=BusType.SYSTEM).connect())

    async def _Call(self, ObjectPath: str, Interface: str, Member: str, Signature: str = "", Body: Optional[list] = None):
        Reply = await self.Bus.call(Message(
            destination=UDisks, path=ObjectPath, interface=Interface, member=Member,
            signature=Signature, body=Body or [],
        ))
        if Reply.message_type == MessageType.ERROR:
            raise OSError(Reply.body[0] if Reply.body else Reply.error_name)
        return Reply.body

    async def ListDevices(self, ShowInternal: bool, ShowSystemDrive: bool = False) -> list[Device]:
        Objects = (await self._Call(UDisksPath, "org.freedesktop.DBus.ObjectManager", "GetManagedObjects"))[0]
        Devices = []
        for ObjectPath, Interfaces in Objects.items():
            if FilesystemInterface not in Interfaces or BlockInterface not in Interfaces:
                continue
            Block = Interfaces[BlockInterface]
            DrivePath = _Value(Block, "Drive", "/")
            if _Value(Block, "HintIgnore", False) or DrivePath == "/":
                continue  # hidden, or not backed by a real drive (loop devices etc.)
            MountPoints = [
                Path(bytes(Raw).rstrip(b"\0").decode(errors="replace"))
                for Raw in _Value(Interfaces[FilesystemInterface], "MountPoints", [])
            ]
            MountPoint = MountPoints[0] if MountPoints else None
            IsSystemMount = False
            if _Value(Block, "HintSystem", False):
                # Internal disks: unmounted ones or ones mounted like
                # removable media always count as "internal". One mounted at
                # a system path only counts when ShowSystemDrive is also on,
                # and even then only the root "/" itself — udisks2 reports
                # /boot, /boot/efi, etc. as their own separate block devices,
                # and showing every one of those as its own non-ejectable
                # entry was more than "the main drive" (singular) implied.
                if not ShowInternal:
                    continue
                if MountPoint is not None and not _UserMount(MountPoint):
                    if not ShowSystemDrive or MountPoint != Path("/"):
                        continue
                    IsSystemMount = True
            Drive = Objects.get(DrivePath, {}).get(DriveInterface, {})
            Label = _Value(Block, "IdLabel", "") or _Value(Block, "HintName", "")
            if not Label:
                Model = " ".join(Part for Part in (_Value(Drive, "Vendor", ""), _Value(Drive, "Model", "")) if Part)
                Label = f"{_HumanSize(_Value(Block, 'Size', 0))} {Model or 'Volume'}".strip()
            Devices.append(Device(
                Key=ObjectPath,
                Label=Label,
                MountPoint=MountPoint,
                DrivePath=DrivePath,
                CanPowerOff=bool(_Value(Drive, "CanPowerOff", False)),
                IsSystemMount=IsSystemMount,
            ))
        return sorted(Devices, key=lambda D: D.Label.lower())

    async def Watch(self, OnChange: Callable[[], None]) -> None:
        """Call OnChange for any udisks2 signal: devices added/removed
        (InterfacesAdded/Removed) and mounts changing (PropertiesChanged)."""
        await self.Bus.call(Message(
            destination="org.freedesktop.DBus", path="/org/freedesktop/DBus",
            interface="org.freedesktop.DBus", member="AddMatch", signature="s",
            body=[f"type='signal',sender='{UDisks}'"],
        ))

        def Handler(Received: Message):
            if Received.message_type == MessageType.SIGNAL and (Received.path or "").startswith(UDisksPath):
                OnChange()
            return None

        self.Bus.add_message_handler(Handler)

    async def Mount(self, ObjectPath: str) -> Path:
        return Path((await self._Call(ObjectPath, FilesystemInterface, "Mount", "a{sv}", [{}]))[0])

    async def Unmount(self, ObjectPath: str) -> None:
        await self._Call(ObjectPath, FilesystemInterface, "Unmount", "a{sv}", [{}])

    async def PowerOff(self, DrivePath: str) -> None:
        await self._Call(DrivePath, DriveInterface, "PowerOff", "a{sv}", [{}])


class DevicesPanel(ft.Container):
    def __init__(
        self,
        OnNavigate: Callable[[Path], None],
        Config: AppConfig,
        OnEmptyChanged: Optional[Callable[[bool], None]] = None,
    ):
        self.OnNavigate = OnNavigate
        self.Config = Config
        # Features/Minor: "Option to hide devices section if there are none,
        # but show if there are" — fired from _Show only on an empty/non-
        # empty transition (not every device-list change), so Sidebar isn't
        # re-toggling visibility on every plug/unplug of an already-nonempty list.
        self.OnEmptyChanged = OnEmptyChanged
        self.Client: Optional[UDisksClient] = None
        self.CurrentDevices: list[Device] = []
        self.CurrentPath: Optional[Path] = None
        self.DeviceTiles: dict[str, ft.ListTile] = {}
        self.RefreshScheduled = False
        self.Busy: set[str] = set()  # devices being mounted/ejected right now
        self.Unmounted = False

        self.List = ft.ListView(spacing=0)
        self.List.controls = self._BuildTiles()
        super().__init__(content=self.List, padding=0)

    def did_mount(self) -> None:
        self.Unmounted = False
        self.page.run_task(self._Start)

    def will_unmount(self) -> None:
        # The sidebar gets rebuilt when its settings change; close this panel's
        # D-Bus connection so old panels don't pile up listening to udisks2.
        self.Unmounted = True
        if self.Client is not None:
            self.Client.Bus.disconnect()
            self.Client = None

    async def _Start(self) -> None:
        if IsWindows:
            await self._PollLoop()  # no udisks2: Windows' own drive list, checked every few seconds
            return
        try:
            self.Client = await UDisksClient.Connect()
            await self.Client.Watch(self._ScheduleRefresh)
        except Exception:
            self.Client = None
            await self._PollLoop()
            return
        await self._Refresh(IsStartup=True)

    async def _PollLoop(self) -> None:
        while not self.Unmounted:
            if IsWindows:
                self._Show(ListWindowsDrives(self.Config.ShowInternalDrives, self.Config.ShowSystemDrive))
            else:
                self._Show(ListFallbackDevices())
            await asyncio.sleep(PollInterval)

    def _ScheduleRefresh(self) -> None:
        # One plug-in fires a burst of signals; refresh once per burst.
        if not self.RefreshScheduled and not self.Unmounted:
            self.RefreshScheduled = True
            self.page.run_task(self._DelayedRefresh)

    async def _DelayedRefresh(self) -> None:
        await asyncio.sleep(0.25)
        self.RefreshScheduled = False
        if self.Client is not None and not self.Unmounted:
            await self._Refresh()

    async def _Refresh(self, IsStartup: bool = False) -> None:
        try:
            Devices = await self.Client.ListDevices(self.Config.ShowInternalDrives, self.Config.ShowSystemDrive)
        except Exception:
            return
        Known = {D.Key for D in self.CurrentDevices}
        self._Show(Devices, Force=IsStartup)
        if IsStartup or not self.Config.AutoMountDevices:
            return
        for NewDevice in Devices:
            if NewDevice.Key not in Known and NewDevice.MountPoint is None:
                await self._MountDevice(NewDevice, Open=False)

    def _Show(self, Devices: list[Device], Force: bool = False) -> None:
        WasEmpty = not self.CurrentDevices
        if Devices == self.CurrentDevices and not Force:
            return
        self.CurrentDevices = Devices
        self.List.controls = self._BuildTiles()
        if self.CurrentPath is not None:
            self._ApplySelection()
        self.update()
        IsEmpty = not Devices
        # Force (startup) always fires too, not just real transitions: an
        # empty list right from the start is still "WasEmpty == IsEmpty",
        # which looks like a no-op transition but is actually the one time
        # Sidebar's listener most needs the initial state explicitly, since
        # DevicesBox otherwise defaults to visible with nothing to show.
        if self.OnEmptyChanged is not None and (IsEmpty != WasEmpty or Force):
            self.OnEmptyChanged(IsEmpty)

    def _BuildTiles(self) -> list[ft.Control]:
        self.DeviceTiles = {}
        if not self.CurrentDevices:
            return [ft.Container(ft.Text("No devices.", color=Palette.Overlay0, size=12), padding=8)]

        Tiles = []
        for Current in self.CurrentDevices:
            Mounted = Current.MountPoint is not None
            Tile = ft.ListTile(
                title=Current.Label,
                leading=ft.Icon(
                    ft.Icons.USB if Mounted else ft.Icons.USB_OUTLINED,
                    color=Palette.Mauve if Mounted else Palette.Overlay0,
                ),
                trailing=ft.IconButton(
                    icon=ft.Icons.EJECT,
                    icon_color=Palette.Overlay0,
                    tooltip="Eject",
                    # Never offer to eject the system/boot drive — unmounting
                    # your own running root filesystem would crash it.
                    visible=self.Client is not None and not Current.IsSystemMount
                    and (Mounted or Current.CanPowerOff),
                    on_click=lambda e, D=Current: self.page.run_task(self._EjectDevice, D),
                ),
                tooltip=str(Current.MountPoint) if Mounted else "Not mounted: click to mount",
                on_click=lambda e, D=Current: self.page.run_task(self._OpenDevice, D),
                selected_tile_color=Palette.Surface1,
                selected_color=Palette.Mauve,
            )
            self.DeviceTiles[Current.Key] = Tile
            Tiles.append(Tile)
        return Tiles

    async def _OpenDevice(self, Target: Device) -> None:
        if Target.MountPoint is not None:
            self.OnNavigate(Target.MountPoint)
        else:
            await self._MountDevice(Target, Open=True)

    async def _MountDevice(self, Target: Device, Open: bool) -> None:
        if self.Client is None or Target.Key in self.Busy:
            return
        self.Busy.add(Target.Key)
        try:
            MountPoint = await self.Client.Mount(Target.Key)
        except OSError as Error:
            self._Notify(f"Couldn't mount {Target.Label}: {Error}")
            return
        finally:
            self.Busy.discard(Target.Key)
        await self._Refresh()
        if Open:
            self.OnNavigate(MountPoint)

    async def _EjectDevice(self, Target: Device) -> None:
        if self.Client is None or Target.Key in self.Busy or Target.IsSystemMount:
            return  # never unmount the system drive, regardless of how this got called
        self.Busy.add(Target.Key)
        try:
            if Target.MountPoint is not None:
                # Step out of the device first, or the unmount is "busy".
                if self.CurrentPath is not None and (
                    self.CurrentPath == Target.MountPoint or Target.MountPoint in self.CurrentPath.parents
                ):
                    self.OnNavigate(Path.home())
                await self.Client.Unmount(Target.Key)
            if Target.CanPowerOff:
                await self.Client.PowerOff(Target.DrivePath)
            self._Notify(f"{Target.Label} can be removed now.", Kind="success")
        except OSError as Error:
            self._Notify(f"Couldn't eject {Target.Label}: {Error}")
        finally:
            self.Busy.discard(Target.Key)
        await self._Refresh()

    def _Notify(self, Message: str, Kind: str = "error") -> None:
        Background, Accent = NotificationStyle(Kind)
        self.page.show_dialog(ft.SnackBar(
            content=ft.Text(Message, color=Accent),
            bgcolor=Background,
            shape=ft.RoundedRectangleBorder(radius=4, side=ft.BorderSide(1, Accent)),
            open=True,
        ))

    def UpdateSelection(self, CurrentPath: Path) -> None:
        self.CurrentPath = CurrentPath
        if self._ApplySelection():
            self.update()

    def _ApplySelection(self) -> bool:
        Changed = False
        for Current in self.CurrentDevices:
            Tile = self.DeviceTiles.get(Current.Key)
            if Tile is None:
                continue
            NewValue = Current.MountPoint is not None and (
                self.CurrentPath == Current.MountPoint or Current.MountPoint in self.CurrentPath.parents
            )
            if Tile.selected != NewValue:
                Tile.selected = NewValue
                Changed = True
        return Changed