/*
 * vdev_root_sys_block.c
 * Copyright (C) 2025  Aitor C.Z. <aitor_czr@gnuinos.org>
 *
 * This program is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the
 * Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * See the COPYING file.
 */

#include <ctype.h>
#include <dirent.h>
#include <errno.h>
#include <mntent.h>
#include <stdio.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/sysmacros.h>
#include <unistd.h>

#include "libvdev/misc.h"
#include "libvdev/sbuf.h"

/* Function to find the actual physical device of a mapper device (LVM) */
void resolve_physical_devices(const char *sysfs_root, const char *dev_path,
								   sbuf_t *s)
{
	char sys_path[512];
	struct stat st;

	if (stat(dev_path, &st) != 0)
		return;

	/* If the device starts with /dev/dm- or is located in /dev/mapper/
	 * look for its "slaves" in <sysfs>/dev/block/MAJOR:MINOR/slaves */
	snprintf(sys_path, sizeof(sys_path), "%s/dev/block/%d:%d/slaves", sysfs_root,
			 major(st.st_rdev), minor(st.st_rdev));

	DIR *d = opendir(sys_path);
	if (d) {
		struct dirent *dir;
		int first = 1;
		while ((dir = readdir(d)) != NULL) {
			if (dir->d_name[0] != '.') {
				sbuf_concat(s, 2, first ? "" : ", ", dir->d_name);
				first = 0;
			}
		}
		closedir(d);
	}
}

/* Normalize the physical device name (SATA, NVMe, RAID, ZFS). */
void normalize_disk_name(char *dev)
{
	int len = strlen(dev);

	/* NVMe: nvme0n1p1 -> nvme0n1 */
	char *p_part = strstr(dev, "p");

	if (p_part && p_part > dev && isdigit(*(p_part - 1)) &&
		isdigit(*(p_part + 1))) {
		*p_part = '\0';
		return;
	}

	/* Software RAID (mdadm): md127p1 -> md127 */
	if (strncmp(dev, "md", 2) == 0) {
		if (p_part && isdigit(*(p_part + 1))) {
			*p_part = '\0';
			return;
		}
	}

	/* ZFS or datasets: if it contains '/', keep the pool/base name */
	char *slash = strchr(dev, '/');
	if (slash) {
		*slash = '\0';
		return;
	}

	/* NVMe namespace without partition suffix: nvme0n1 (do not strip the "1"). */
	if (len >= 6 && strncmp(dev, "nvme", 4) == 0 && dev[4] != '\0') {
		char *n = strchr(dev + 4, 'n');
		if (n != NULL && n[1] != '\0' && strchr(n + 1, 'p') == NULL) {
			while (n[1] != '\0')
				n++;
			if (n > dev && isdigit(*n))
				return;
		}
	}

	/* Standard case (SATA/USB): sdb1 -> sdb */
	while (len > 0 && isdigit(dev[len - 1]))
		dev[--len] = '\0';
}

static void usage(char const *argv0)
{
	fprintf(stderr, "Usage: %s [SYSFS_MOUNTPOINT]\n", argv0);
}

int main(int argc, char **argv)
{
	struct mntent *ent;
	FILE *mtab;
	DIR *dirp;
	struct dirent *E;
	sbuf_t sysdir;
	sbuf_t devdir;
	sbuf_t rootdir;
	sbuf_t path;
	char *str = NULL;
	char *aux = NULL;

	if (argc > 2) {
		usage(argv[0]);
		return 1;
	}
	if (argc == 2 && (!strcmp(argv[1], "-h") || !strcmp(argv[1], "--help"))) {
		usage(argv[0]);
		return 0;
	}

	mtab = setmntent("/etc/mtab", "r");
	if (mtab == NULL) {
		vdev_error("Cannot open /etc/mtab");
		return 1;
	}

	sbuf_init(&sysdir);
	sbuf_init(&devdir);
	sbuf_init(&rootdir);

	if (argc == 2)
		sbuf_addstr(&sysdir, argv[1]);

	while ((ent = getmntent(mtab)) != NULL) {
		if (sysdir.len == 0 && !strcmp(ent->mnt_type, "sysfs")) {
			sbuf_addstr(&sysdir, ent->mnt_dir);
		} else if (devdir.len == 0 && ((!strcmp(ent->mnt_fsname, "udev") &&
										!strcmp(ent->mnt_type, "devtmpfs")) ||
									   (!strcmp(ent->mnt_fsname, "none") &&
										!strcmp(ent->mnt_type, "tmpfs")))) {
			sbuf_addstr(&devdir, ent->mnt_dir);
		} else if (rootdir.len == 0 && !strcmp(ent->mnt_dir, "/")) {
			if (strstr(ent->mnt_fsname, "/dev/mapper/") ||
				strstr(ent->mnt_fsname, "/dev/dm-"))
				resolve_physical_devices(sysdir.buf, ent->mnt_fsname, &rootdir);
			else
				sbuf_addstr(&rootdir, ent->mnt_fsname);
		} else if (rootdir.len == 0 && !strcmp(ent->mnt_dir, "/run/live/medium")) {
			sbuf_addstr(&rootdir, ent->mnt_fsname);
		}
	}

	// Close the mount table file
	endmntent(mtab);

	if (sysdir.len == 0 || devdir.len == 0 || rootdir.len == 0) {
		sbuf_free(&sysdir);
		sbuf_free(&devdir);
		sbuf_free(&rootdir);
		return 1;
	}

	str = strstr(rootdir.buf, devdir.buf);
	if (str == NULL) {
		sbuf_free(&sysdir);
		sbuf_free(&devdir);
		sbuf_free(&rootdir);
		return 1;
	}
	str = str + devdir.len + 1;
	aux = strdup(str);
	normalize_disk_name(aux);

	sbuf_init(&path);
	sbuf_concat(&path, 6, sysdir.buf, "/block/", aux, "/", str, "/partition");

	// Does '/sys/block/sdb/sdb1/partition' exist?
	if (access(path.buf, F_OK) == 0)
		printf("%s\n", aux);

	sbuf_free(&path);
	free(aux);

	sbuf_free(&sysdir);
	sbuf_free(&devdir);
	sbuf_free(&rootdir);

	return 0;
}
