from dataclasses import dataclass

import enum
import typing

from dbus_fast.aio import MessageBus, ProxyInterface

from . import SIZE_GIB, SIZE_TEST_POOL, SIZE_TEST_POOL_OVERHEAD

if typing.TYPE_CHECKING:
    from collections.abc import Mapping

    from dbus_fast import Variant


STRATIS_UNIT = 'stratisd.service'
STRATIS_SERVICE = 'org.storage.stratis3'
STRATIS_ROOT = '/org/storage/stratis3'

STRATIS_REV = 6  # AlmaLinux 9 in the VM has stratis 3.6


class TypedVariant[T](typing.Protocol):
    value: T


class PropsPool(typing.TypedDict):
    Name: TypedVariant[str]


class OptionalByBool[T](typing.Protocol):
    @typing.overload
    def __getitem__(self, idx: typing.Literal[0]) -> int: ...
    @typing.overload
    def __getitem__(self, idx: typing.Literal[1]) -> T: ...


def get_optional[T](opt: OptionalByBool[T]) -> T | None:
    return opt[1] if opt[0] else None


class PropsFilesystem(typing.TypedDict):
    Name: TypedVariant[str]
    Pool: TypedVariant[str]
    Size: TypedVariant[str]
    SizeLimit: TypedVariant[OptionalByBool[str]]
    Devnode: TypedVariant[str]


@dataclass(kw_only=True)
class Filesystem:
    INTERFACE: typing.ClassVar[str] = f'org.storage.stratis3.filesystem.r{STRATIS_REV}'

    props: PropsFilesystem
    proxy: ProxyInterface

    @property
    def name(self) -> str:
        return self.props['Name'].value

    @property
    def size(self) -> int:
        return int(self.props['Size'].value)

    @property
    def size_limit(self) -> int | None:
        if (limit := get_optional(self.props['SizeLimit'].value)) is None:
            return None
        return int(limit)

    @property
    def devnode(self) -> str:
        return self.props['Devnode'].value


@dataclass(kw_only=True)
class Pool:
    INTERFACE: typing.ClassVar[str] = f'org.storage.stratis3.pool.r{STRATIS_REV}'

    props: PropsPool
    proxy: ProxyInterface
    filesystems: list[Filesystem]

    @property
    def name(self) -> str:
        return self.props['Name'].value

    async def destroy_filesystems(self, filesystems: list[Filesystem]) -> None:
        results, ret_code, ret_s = typing.cast(
            'tuple[OptionalByBool[list[str]], int, str]',
            await typing.cast(
                'typing.Any',
                self.proxy,
            ).call_destroy_filesystems(  # pyright: ignore[reportAny]
                [fs.proxy.path for fs in filesystems],
            ),
        )
        assert ret_code == 0, ret_s

        results = get_optional(results) or []
        assert len(results) == len(filesystems)

    async def wipe_filesystems(self) -> None:
        if self.filesystems:
            await self.destroy_filesystems(self.filesystems)
            self.filesystems = []


class PoolSelection(enum.Enum):
    DEFAULT = enum.auto()
    PRIMARY = enum.auto()
    SECONDARY = enum.auto()

    PREFIX_POOL = enum.nonmember('topolvm-')
    PREFIX_PROVISIONER = enum.nonmember('topolvm-provisioner')

    @property
    def pool(self) -> str:
        target = self if self != PoolSelection.DEFAULT else PoolSelection.PRIMARY
        return self.PREFIX_POOL + target.name.lower()

    @property
    def device_class(self) -> str:
        return self.name.lower() if self != PoolSelection.DEFAULT else ''

    @property
    def provisioner(self) -> str:
        match self:
            case PoolSelection.DEFAULT | PoolSelection.PRIMARY:
                return self.PREFIX_PROVISIONER
            case PoolSelection.SECONDARY:
                return f'{self.PREFIX_PROVISIONER}-{self.device_class}'

    @property
    def spare_bytes(self) -> int:
        match self:
            case PoolSelection.DEFAULT | PoolSelection.PRIMARY:
                return 2 * SIZE_GIB
            case PoolSelection.SECONDARY:
                return 4 * SIZE_GIB

    @property
    def initial_free_bytes(self) -> int:
        return SIZE_TEST_POOL - SIZE_TEST_POOL_OVERHEAD - self.spare_bytes


@dataclass(kw_only=True)
class StratisState:
    pools_by_name: dict[str, Pool]

    @staticmethod
    async def load(bus: MessageBus) -> 'StratisState':
        introspection = await bus.introspect(STRATIS_SERVICE, STRATIS_ROOT)
        stratisd = bus.get_proxy_object(STRATIS_SERVICE, STRATIS_ROOT, introspection)

        objects = typing.cast(
            'dict[str, dict[str, dict[str, Variant]]]',
            await typing.cast(
                'typing.Any',
                stratisd.get_interface('org.freedesktop.DBus.ObjectManager'),
            ).call_get_managed_objects(),  # pyright: ignore[reportAny]
        )

        pools_by_path: dict[str, Pool] = {}
        pools_by_name: dict[str, Pool] = {}
        for path, ifs in objects.items():
            if (props := ifs.get(Pool.INTERFACE)) is None:
                continue

            props = typing.cast(
                'PropsPool', typing.cast('Mapping[object, object]', props)
            )

            name = props['Name'].value
            if not name.startswith(PoolSelection.PREFIX_POOL):
                continue

            introspection = await bus.introspect(stratisd.bus_name, path)
            pools_by_path[path] = pools_by_name[name] = Pool(
                props=props,
                proxy=bus.get_proxy_object(
                    stratisd.bus_name, path, introspection
                ).get_interface(Pool.INTERFACE),
                filesystems=[],
            )

        if len(pools_by_path) != 2:
            raise ValueError(f'expected 2 test pools, got {len(pools_by_path)}')

        for path, ifs in objects.items():
            if (props := ifs.get(Filesystem.INTERFACE)) is None:
                continue

            props = typing.cast(
                'PropsFilesystem', typing.cast('Mapping[object, object]', props)
            )

            pool = props['Pool'].value
            if pool not in pools_by_path:
                continue

            # TODO: reuse introspection past the first filesystem?
            introspection = await bus.introspect(stratisd.bus_name, path)
            pools_by_path[pool].filesystems.append(
                Filesystem(
                    props=props,
                    proxy=bus.get_proxy_object(
                        stratisd.bus_name, path, introspection
                    ).get_interface(Filesystem.INTERFACE),
                )
            )

        return StratisState(pools_by_name=pools_by_name)


@dataclass
class Stratisd:
    bus: MessageBus
    state: StratisState

    @staticmethod
    async def create(bus: MessageBus) -> 'Stratisd':
        state = await StratisState.load(bus)
        return Stratisd(bus, state)

    async def reload(self) -> None:
        self.state = await StratisState.load(self.bus)

    def pool(self, selection: PoolSelection) -> Pool:
        return self.state.pools_by_name[selection.pool]

    async def wipe_all_filesystems(self) -> None:
        for pool in self.state.pools_by_name.values():
            await pool.wipe_filesystems()