package blkid

import (
	"bufio"
	"errors"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

type BlockDevice struct {
	Device    string
	UUID      string
	BlockSize int
	Type      string
	PartUUID  string
	PartLabel string
}

// GetBlockDevices retrieves block devices using the `blkid` command.
func GetPartitionIdInfo() ([]BlockDevice, error) {
	//Check if the current user have superuser privileges
	cmd := exec.Command("id", "-u")
	userIDOutput, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	// Check if the user ID is 0 (root)
	// If not, run blkid without sudo
	if strings.TrimSpace(string(userIDOutput)) == "0" {
		cmd = exec.Command("blkid")
	} else {
		cmd = exec.Command("sudo", "blkid")
	}

	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	scanner := bufio.NewScanner(strings.NewReader(string(output)))
	devices := []BlockDevice{}
	re := regexp.MustCompile(`(\S+):\s+(.*)`)

	for scanner.Scan() {
		line := scanner.Text()
		matches := re.FindStringSubmatch(line)
		if len(matches) != 3 {
			continue
		}

		device := matches[1]
		attributes := matches[2]
		deviceInfo := BlockDevice{Device: device}

		for _, attr := range strings.Split(attributes, " ") {
			kv := strings.SplitN(attr, "=", 2)
			if len(kv) != 2 {
				continue
			}
			key := kv[0]
			value := strings.Trim(kv[1], `"`)

			switch key {
			case "UUID":
				deviceInfo.UUID = value
			case "BLOCK_SIZE":
				// Convert block size to int if possible
				blockSize, err := strconv.Atoi(value)
				if err == nil {
					deviceInfo.BlockSize = blockSize
				} else {
					deviceInfo.BlockSize = 0
				}

			case "TYPE":
				deviceInfo.Type = value
			case "PARTUUID":
				deviceInfo.PartUUID = value
			case "PARTLABEL":
				deviceInfo.PartLabel = value
			}
		}

		devices = append(devices, deviceInfo)
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return devices, nil
}

// GetBlockDeviceIDFromDevicePath retrieves block device information for a given device path.
func GetPartitionIDFromDevicePath(devpath string) (*BlockDevice, error) {
	devpath = strings.TrimPrefix(devpath, "/dev/")
	if strings.Contains(devpath, "/") {
		return nil, errors.New("invalid device path")
	}

	devpath = "/dev/" + devpath

	devices, err := GetPartitionIdInfo()
	if err != nil {
		return nil, err
	}

	for _, device := range devices {
		if device.Device == devpath {
			return &device, nil
		}
	}

	return nil, errors.New("device not found")
}