import json
import os
import re
import subprocess
from pathlib import Path


FILESYSTEMS = {
    "fat32": {
        "mkfs": ["mkfs.fat", "-F", "32"],
        "parted_type": "fat32",
        "parted_flags": {"esp": "esp"},
    },
    "ext4": {
        "mkfs": ["mkfs.ext4", "-F"],
        "parted_type": "ext4",
        "parted_flags": {},
    },
}


def run(*args, check=True):
    """Run a system command and return stdout."""
    result = subprocess.run(
        [str(arg) for arg in args],
        check=check,
        text=True,
        capture_output=True,
    )
    return result.stdout.strip()


def partition_device(disk, number):
    """Return the correct partition path for common Linux devices."""
    disk = str(disk)
    suffix = "p" if re.search(r"\d$", disk) else ""
    return f"{disk}{suffix}{number}"


def parse_size(size, disk_bytes):
    """Convert a size such as 200M or 100% into bytes."""
    if size == "100%":
        return disk_bytes

    match = re.fullmatch(
        r"\s*(\d+(?:\.\d+)?)\s*([KMGT]?)B?\s*",
        size.upper(),
    )
    if not match:
        raise ValueError(f"Invalid partition size: {size!r}")

    value = float(match.group(1))
    unit = match.group(2)
    multiplier = {
        "": 1,
        "K": 1024,
        "M": 1024**2,
        "G": 1024**3,
        "T": 1024**4,
    }[unit]

    return int(value * multiplier)


def get_table(disk):
    """Read the partition table and return its partitions."""
    try:
        data = json.loads(run("sfdisk", "--json", disk))
    except subprocess.CalledProcessError:
        return None

    table = data.get("partitiontable")
    if not table:
        return None

    return table


def get_filesystem(device):
    """Return the detected filesystem type, or None."""
    result = run(
        "blkid", "-p", "-s", "TYPE", "-o", "value",
        device, check=False,
    )
    return result.lower() if result else None


def apply_partition_config(
    disk,
    config,
    *,
    dry_run=True,
    format_mismatches=False,
):
    """
    Apply a partition configuration to a disk.

    Safety behavior:
      - Dry-run by default.
      - Creates GPT only if no partition table exists.
      - Does not automatically resize, delete, or reorder partitions.
      - Refuses to format filesystems unless format_mismatches=True.
      - Verifies the expected partition count and filesystem types.

    Requires root privileges and: parted, sfdisk, blkid, blockdev,
    mkfs.fat, mkfs.ext4, and partprobe.
    """
    disk = str(Path(disk))

    if os.geteuid() != 0:
        raise PermissionError("Run this function as root.")

    if not os.path.exists(disk):
        raise FileNotFoundError(disk)

    if run("lsblk", "-dn", "-o", "TYPE", disk) != "disk":
        raise ValueError(f"{disk} is not a whole-disk device")

    partitions = config["partitions"]
    if not partitions:
        raise ValueError("At least one partition is required")

    for p in partitions:
        fs = p["type"].lower()
        if fs not in FILESYSTEMS:
            raise ValueError(f"Unsupported filesystem: {fs}")
        if not isinstance(p.get("flags", []), list):
            raise ValueError("'flags' must be a list")

    disk_bytes = int(run("blockdev", "--getsize64", disk))
    sector_size = int(run("blockdev", "--getss", disk))
    alignment = 1024 * 1024  # 1 MiB

    def align_up(n):
        return ((n + alignment - 1) // alignment) * alignment

    # Calculate partition boundaries. Reserve 1 MiB at the start.
    layout = []
    start = alignment

    for i, p in enumerate(partitions):
        requested = p["size"]
        if requested == "100%":
            end = disk_bytes - alignment
        else:
            size_bytes = parse_size(requested, disk_bytes)
            end = align_up(start + size_bytes)

        if i == len(partitions) - 1 and requested == "100%":
            end = disk_bytes - alignment

        if end <= start:
            raise ValueError(f"Invalid size for partition {i + 1}")

        if end > disk_bytes - alignment:
            raise ValueError(
                f"Partition {i + 1} exceeds disk capacity"
            )

        layout.append({
            "number": i + 1,
            "start": start,
            "end": end,
            "type": p["type"].lower(),
            "flags": p.get("flags", []),
        })
        start = end

    def parted(*args):
        if dry_run:
            print("+ parted", "-s", disk, *args)
        else:
            run("parted", "-s", disk, *args)

    table = get_table(disk)

    # Create GPT only when the disk has no partition table.
    if table is None:
        print(f"No readable partition table found on {disk}.")
        parted("mklabel", "gpt")
        table = None

    elif table.get("label") != "gpt":
        raise RuntimeError(
            f"{disk} has a {table.get('label')} table, not GPT. "
            "Refusing to overwrite it."
        )

    existing = table.get("partitions", []) if table else []

    # Never guess how to repair an incomplete or mismatched layout.
    if existing and len(existing) != len(layout):
        raise RuntimeError(
            f"Found {len(existing)} partitions; "
            f"configuration requires {len(layout)}. "
            "Automatic resizing/deletion is disabled."
        )

    if not existing:
        for item in layout:
            fs = item["type"]
            start_mib = item["start"] / (1024**2)
            end_mib = item["end"] / (1024**2)

            # parted accepts MiB boundaries.
            parted(
                "mkpart",
                "primary",
                FILESYSTEMS[fs]["parted_type"],
                f"{start_mib:.6f}MiB",
                f"{end_mib:.6f}MiB",
            )

            for flag in item["flags"]:
                if flag not in FILESYSTEMS[fs]["parted_flags"]:
                    raise ValueError(
                        f"Unsupported flag {flag!r} for {fs}"
                    )
                parted(
                    "set",
                    str(item["number"]),
                    FILESYSTEMS[fs]["parted_flags"][flag],
                    "on",
                )

    elif not dry_run:
        # Re-read the on-disk table before inspecting filesystems.
        table = get_table(disk)
        existing = table.get("partitions", [])

        # Validate the existing boundaries. Do not modify them.
        for expected, actual in zip(layout, existing):
            actual_start = int(actual["start"]) * sector_size
            actual_end = (
                int(actual["start"]) + int(actual["size"])
            ) * sector_size

            if (
                abs(actual_start - expected["start"]) > alignment
                or abs(actual_end - expected["end"]) > alignment
            ):
                raise RuntimeError(
                    f"Partition {expected['number']} has incorrect "
                    "boundaries. Automatic resizing is disabled."
                )

    # Ask the kernel to re-read the partition table.
    if not dry_run:
        run("partprobe", disk)
        run("udevadm", "settle")

    results = []

    for item in layout:
        device = partition_device(disk, item["number"])
        expected_fs = item["type"]

        if dry_run:
            print(
                f"Would inspect {device}; expected filesystem "
                f"{expected_fs}"
            )
            if format_mismatches:
                print(f"Would format {device} if needed")
            results.append({
                "device": device,
                "expected": expected_fs,
                "status": "dry-run",
            })
            continue

        if not os.path.exists(device):
            raise RuntimeError(
                f"Partition device {device} did not appear"
            )

        actual_fs = get_filesystem(device)

        if actual_fs != expected_fs:
            if not format_mismatches:
                raise RuntimeError(
                    f"{device}: expected {expected_fs}, "
                    f"found {actual_fs!r}. "
                    "Set format_mismatches=True to format it."
                )

            # Avoid accidentally formatting a mounted filesystem.
            mounts = run(
                "lsblk", "-n", "-o", "MOUNTPOINTS", device
            )
            if mounts.strip():
                raise RuntimeError(
                    f"{device} appears to be mounted; refusing "
                    "to format it."
                )

            fs = FILESYSTEMS[expected_fs]
            run(*fs["mkfs"], device)

            actual_fs = get_filesystem(device)
            if actual_fs != expected_fs:
                raise RuntimeError(
                    f"Filesystem verification failed for {device}"
                )

        results.append({
            "device": device,
            "expected": expected_fs,
            "actual": actual_fs,
            "status": "ok",
        })

    return results