from collections import defaultdict
from dataclasses import dataclass, field
from enum import StrEnum
import json
from pathlib import Path
import re
from typing import Any, Iterable, Self

from humanfriendly import parse_size

from kisiac.common import check_type, exists_cmd, run_cmd, UserError


CRYPT_PREFIX = "crypt_"
VGS_DEVICE_REPORT_RE = re.compile(r"^(?P<device>.+)\((?P<info>.+)\)$")
LV_INTERNAL_NAME = re.compile(r"^\[(?P<name>.+)\]$")


def parse_lv_name(name_entry) -> tuple[str, bool]:
    m = LV_INTERNAL_NAME.match(name_entry)
    if m:
        return m.group("name"), True
    else:
        return name_entry, False


@dataclass(frozen=True)
class PV:
    device: Path


class Subvolume(StrEnum):
    ORIG = "corig"
    CACHE = "cvol"


def _subvolume_name(lv_name: str, subvolume: Subvolume) -> str:
    return f"{lv_name}_{subvolume}"


@dataclass
class LV:
    name: str
    layout: set[str]
    size: int | None
    stripes: int
    stripe_size: int  # in bytes
    pv_tag: str | None
    cache_size: int | None
    cache_pv_tag: str | None
    cache_mode: str | None = None

    @property
    def cache_subvolume_name(self) -> str:
        return _subvolume_name(self.name, Subvolume.CACHE)

    @property
    def orig_subvolume_name(self) -> str:
        return _subvolume_name(self.name, Subvolume.ORIG)

    def is_cached(self) -> bool:
        return self.cache_mode is not None

    def is_same_layout(self, other: Self) -> bool:
        # TODO find a better way to compare the layout in the cache case
        if "cache" in other.layout:
            # current LV is cached, original layout is not reported
            return True
        return (
            self.layout <= other.layout
            or other.layout <= self.layout
            and self.stripes == other.stripes
            and self.stripe_size == other.stripe_size
        )

    def is_same_size(self, other: Self) -> bool:
        if self.fills_vg() or other.fills_vg():
            # if both just fill the VG, their sizes are considered unchanged
            return self.fills_vg() and other.fills_vg()

        assert self.size is not None and other.size is not None

        def simplify(size: int) -> int:
            # tens of MB should be precise enough
            return size // 10**7

        return simplify(self.size) == simplify(other.size)

    def fills_vg(self) -> bool:
        return self.size is None

    def size_arg(self) -> list[str]:
        return self._size_arg(Subvolume.ORIG)

    def cache_size_arg(self) -> list[str]:
        return self._size_arg(Subvolume.CACHE)

    def _size_arg(self, subvolume: Subvolume) -> list[str]:
        size = self.size if subvolume is Subvolume.ORIG else self.cache_size
        if size is None:
            assert subvolume is Subvolume.ORIG, "Cache resizing not supported"
            return ["--extents", "+100%FREE"]
        else:
            arg = "--cachesize" if subvolume is Subvolume.CACHE else "--size"
            return [arg, f"{size}B"]

    def stripe_args(self) -> list[str]:
        if self.stripes > 1:
            return [
                "--stripes",
                str(self.stripes),
                "--stripesize",
                f"{self.stripe_size}B",
            ]
        else:
            return []

    def select_arg(self) -> list[str]:
        return self._select_arg(Subvolume.ORIG)

    def cache_select_arg(self) -> list[str]:
        return self._select_arg(Subvolume.CACHE)

    def _select_arg(self, subvolume: Subvolume) -> list[str]:
        pv_tag = self.pv_tag if subvolume is Subvolume.ORIG else self.cache_pv_tag
        if pv_tag:
            return [f"@{pv_tag}"]
        else:
            return []

    def type_arg(self) -> list[str]:
        if self.layout:
            return ["--type", ",".join(self.layout)]
        return []


@dataclass(frozen=True)
class VG:
    name: str
    pvs: dict[str | None, set[PV]] = field(default_factory=dict)
    lvs: dict[str, LV] = field(default_factory=dict)

    def get_lv_device(self, lv_name: str) -> Path:
        return Path("/dev") / self.name / lv_name


@dataclass
class LVMSetup:
    pvs: set[PV] = field(default_factory=set)
    vgs: dict[str, VG] = field(default_factory=dict)
    missing_pvs: set[PV] = field(default_factory=set)

    def is_empty(self) -> bool:
        return not self.pvs and not self.vgs

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> Self:
        check_type("lvm key", config, dict)
        entities = cls()
        for pv in config.get("pvs", []):
            check_type("lvm pv entry", pv, str)
            entities.pvs.add(
                PV(
                    device=Path(pv),
                )
            )
        for name, settings in config.get("vgs", {}).items():
            check_type(f"lvm vg {name} entry", settings, dict)

            lvs = settings.get("lvs", {})
            check_type(f"lvm vg {name} lvs entry", lvs, dict)

            lvs_entities = {}
            for lv_name, lv_settings in lvs.items():
                check_type(f"lvm vg {name} lv {lv_name} entry", lv_settings, dict)

                def handle_size(size_entry):
                    if size_entry == "rest":
                        return None
                    else:
                        return parse_size(size_entry, binary=True)

                size = handle_size(lv_settings["size"])
                cache_size = lv_settings.get("cache_size")
                cache_size = handle_size(cache_size) if cache_size is not None else None

                layout = {lv_settings.get("layout")}

                lvs_entities[lv_name] = LV(
                    name=lv_name,
                    layout=layout,
                    size=size,
                    stripes=lv_settings.get("stripes", 1),
                    stripe_size=parse_size(lv_settings.get("stripe_size", "0B")),
                    pv_tag=lv_settings.get("pv_tag"),
                    cache_pv_tag=lv_settings.get("cache_pv_tag"),
                    cache_size=lv_settings.get("cache_size"),
                    cache_mode=lv_settings.get("cache_mode"),
                )

            pvs = {}
            for tag, pvs_entry in settings.get("pvs", {}).items():
                check_type(f"vg {name} pvs entry", pvs_entry, list)
                pvs[tag] = {PV(device=Path(pv)) for pv in pvs_entry}

            entities.vgs[name] = VG(
                name=name,
                pvs=pvs,
                lvs=lvs_entities,
            )
        return entities

    @classmethod
    def from_system(cls, host: str) -> Self:
        # check if lvm2 is installed, return empty LVM entities otherwise
        if not exists_cmd("pvcreate", host=host, sudo=True):
            return cls()

        entities: Self = cls()

        lv_data = json.loads(
            run_cmd(
                [
                    "lvs",
                    "--all",
                    "--units",
                    "b",
                    "--options",
                    "lv_name,vg_name,lv_layout,lv_size,stripes,stripe_size,origin,cache_mode,pool_lv",
                    "--reportformat",
                    "json",
                ],
                host=host,
                sudo=True,
            ).stdout
        )["report"][0]["lv"]

        vg_devices_raw = json.loads(
            run_cmd(
                ["vgs", "--options", "vg_name,devices", "--reportformat", "json"],
                host=host,
                sudo=True,
            ).stdout
        )["report"][0]["vg"]
        vg_devices = defaultdict(list)
        for entry in vg_devices_raw:
            vg_devices[entry["vg_name"]].append(entry["devices"])

        vg_data = json.loads(
            run_cmd(
                ["vgs", "--options", "vg_name", "--reportformat", "json"],
                host=host,
                sudo=True,
            ).stdout
        )["report"][0]["vg"]

        pv_data = json.loads(
            run_cmd(
                [
                    "pvs",
                    "--options",
                    "pv_name,vg_name,lv_name,pv_tags",
                    "--reportformat",
                    "json",
                ],
                host=host,
                sudo=True,
            ).stdout
        )["report"][0]["pv"]

        for entry in vg_data:
            vg_name = entry["vg_name"]
            entities.vgs[vg_name] = VG(name=vg_name)

        for vg_name, device_reports in vg_devices.items():
            entities.missing_pvs.update(get_missing_pvs(device_reports))

        pv_tag_registry = {}
        for entry in pv_data:
            pv_device = entry["pv_name"]

            pv_tags = entry["pv_tags"].split(",")
            if len(pv_tags) > 1:
                raise UserError(
                    "Unsupported number of tags associated with "
                    f"PV {pv_device}: {pv_tags}. Only 1 or 0 allowed."
                )
            elif len(pv_tags) == 1:
                pv_tag = pv_tags[0]
            else:
                pv_tag = None

            pv_obj = PV(device=Path(pv_device))
            entities.pvs.add(pv_obj)
            vg_name = entry["vg_name"]
            if vg_name and vg_name in entities.vgs:
                # in case the PV is assigned to a vg
                vg_pvs = entities.vgs[vg_name].pvs
                if pv_tag not in vg_pvs:
                    vg_pvs[pv_tag] = set()
                vg_pvs[pv_tag].add(pv_obj)

            lv_name, is_internal = parse_lv_name(entry["lv_name"])
            if lv_name:
                pv_tag_registry[lv_name] = pv_tags[0]

        print(pv_tag_registry)

        internal_entries = defaultdict(dict)
        for entry in sorted(lv_data, key=lambda entry: entry["origin"]):
            print(entry)
            vg_name = entry["vg_name"]
            vg = entities.vgs[vg_name]
            lv_name, is_internal = parse_lv_name(entry["lv_name"])
            if is_internal:
                internal_entries[vg_name][lv_name] = entry
                continue

            origin, _ = parse_lv_name(entry["origin"])
            cache_mode = None
            if not origin:
                origin = None
            else:
                cache_mode = entry["cache_mode"]

            lv_orig_name = _subvolume_name(lv_name, Subvolume.ORIG)
            lv_cache_name = _subvolume_name(lv_name, Subvolume.CACHE)
            cache_entry = internal_entries.get(vg_name, {}).get(lv_cache_name, {})
            cache_size = cache_entry.get("lv_size")
            if cache_size is not None:
                cache_size = parse_size(cache_size)

            lv = LV(
                name=lv_name,
                layout=set(entry["lv_layout"].split(",")),
                size=parse_size(entry["lv_size"], binary=True),
                stripes=entry["stripes"],
                stripe_size=parse_size(entry["stripe_size"]),
                pv_tag=pv_tag_registry.get(lv_orig_name),
                cache_pv_tag=pv_tag_registry.get(lv_cache_name),
                cache_size=cache_size,
                cache_mode=cache_mode,
            )

            vg.lvs[lv_name] = lv
        return entities


def get_missing_pvs(device_reports: Iterable[str]) -> Iterable[PV]:
    for device_report in device_reports:
        m = VGS_DEVICE_REPORT_RE.match(device_report)
        assert m is not None, f"Invalid device report: {device_report}"
        info = m.group("info")
        if info == "missing":
            device = m.group("device")
            assert device is not None
            yield PV(device=Path(device))
        else:
            try:
                int(info)
            except ValueError:
                raise ValueError(f"Unexpected vgs device info: {info}")
