#!/usr/bin/env python3
# Genera el disk_config de archinstall (>=4.x) dimensionado al disco REAL y lo
# inyecta en el JSON base. archinstall 4.x quitó la unidad "Percent", así que un
# JSON con tamaños fijos no es portable entre la VM (20G/10G) y el hardware
# (250G/1TB): por eso se calcula en tiempo de instalación con los propios modelos
# de archinstall (válido ante cambios de schema entre versiones).
#
#   gen-disk-config.py BASE.json [NVME_DEV] [SDA_DEV] > merged.json
#
# Layout (replica el hardware): NVMe = ESP 512M + /boot 1G ext4 + / btrfs (resto);
# SDA = /home ext4 (todo). swap = zram (lo pone "swap": true del base).
# Si SDA no existe (portatil de 1 disco), se omite: /home queda dentro del / btrfs.
import json, sys
from pathlib import Path
from archinstall.lib.disk.device_handler import device_handler
from archinstall.lib.models.device import (
    DiskLayoutConfiguration, DiskLayoutType, DeviceModification,
    PartitionModification, ModificationStatus, PartitionType,
    FilesystemType, PartitionFlag, Size, Unit,
)

MiB = 1024 ** 2
END_MARGIN = MiB  # deja 1 MiB al final para la cabecera GPT de respaldo

def part(status, ptype, start_b, len_b, ss, fs=None, mnt=None, flags=()):
    return PartitionModification(
        status=status, type=ptype,
        start=Size(start_b, Unit.B, ss),
        length=Size(len_b, Unit.B, ss),
        fs_type=fs, mountpoint=Path(mnt) if mnt else None,
        flags=list(flags),
    )

def nvme_mods(dev):
    ss = dev.device_info.sector_size
    total = dev.device_info.total_size._normalize()
    esp_start, esp_len = 1 * MiB, 511 * MiB           # 1 -> 512 MiB
    boot_start, boot_len = 512 * MiB, 1024 * MiB        # 512 -> 1536 MiB
    root_start = boot_start + boot_len                  # 1536 MiB
    root_len = total - root_start - END_MARGIN
    return DeviceModification(device=dev, wipe=True, partitions=[
        part(ModificationStatus.CREATE, PartitionType.PRIMARY, esp_start, esp_len, ss,
             FilesystemType.FAT32, "/boot/efi", (PartitionFlag.BOOT, PartitionFlag.ESP)),
        part(ModificationStatus.CREATE, PartitionType.PRIMARY, boot_start, boot_len, ss,
             FilesystemType.EXT4, "/boot"),
        part(ModificationStatus.CREATE, PartitionType.PRIMARY, root_start, root_len, ss,
             FilesystemType.BTRFS, "/"),
    ])

def sda_mods(dev):
    ss = dev.device_info.sector_size
    total = dev.device_info.total_size._normalize()
    start = 1 * MiB
    return DeviceModification(device=dev, wipe=True, partitions=[
        part(ModificationStatus.CREATE, PartitionType.PRIMARY, start,
             total - start - END_MARGIN, ss, FilesystemType.EXT4, "/home"),
    ])

def main():
    base_path = sys.argv[1]
    nvme = sys.argv[2] if len(sys.argv) > 2 else "/dev/nvme0n1"
    sda = sys.argv[3] if len(sys.argv) > 3 else "/dev/sda"

    dev_nvme = device_handler.get_device(Path(nvme))
    if not dev_nvme:
        sys.exit(f"No encuentro el disco de sistema {nvme} (revisa lsblk)")

    # 2 discos -> layout hardware: NVMe=sistema, SDA=/home aparte.
    # 1 disco (portatil) -> todo en el NVMe; /home vive dentro del btrfs raiz.
    dev_sda = device_handler.get_device(Path(sda))
    mods = [nvme_mods(dev_nvme)]
    if dev_sda:
        mods.append(sda_mods(dev_sda))

    layout = DiskLayoutConfiguration(
        config_type=DiskLayoutType.Manual,
        device_modifications=mods,
    )
    cfg = json.loads(Path(base_path).read_text())
    cfg["disk_config"] = layout.json()
    json.dump(cfg, sys.stdout, indent=2)

if __name__ == "__main__":
    main()