package collect

import (
	"encoding/json"
	"strconv"
	"strings"
	"time"
)

type Disk struct {
	Name   string `json:"name"` // sda
	Path   string `json:"path"` // /dev/sda
	Model  string `json:"model"`
	Serial string `json:"serial"`
	Size   uint64 `json:"size"`
	Rota   bool   `json:"rota"`
	Tran   string `json:"tran"` // sata, nvme, usb
}

type lsblkOut struct {
	Blockdevices []struct {
		Name   string   `json:"name"`
		Path   string   `json:"path"`
		Type   string   `json:"type"`
		Model  string   `json:"model"`
		Serial string   `json:"serial"`
		Size   flexUint `json:"size"`
		Rota   flexBool `json:"rota"`
		Tran   string   `json:"tran"`
	} `json:"blockdevices"`
}

// flexBool accepts true/false as well as the "1"/"0" strings older lsblk prints.
type flexBool bool

func (b *flexBool) UnmarshalJSON(data []byte) error {
	switch strings.Trim(string(data), `"`) {
	case "true", "1":
		*b = true
	default:
		*b = false
	}
	return nil
}

// flexUint accepts a JSON number or a numeric string.
type flexUint uint64

func (u *flexUint) UnmarshalJSON(data []byte) error {
	n, err := strconv.ParseUint(strings.Trim(string(data), `"`), 10, 64)
	if err != nil {
		*u = 0
		return nil
	}
	*u = flexUint(n)
	return nil
}

// ListDisks returns physical block devices (whole disks, no loop/zram/dm).
func ListDisks() ([]Disk, error) {
	b, err := run(10*time.Second, "lsblk", "-J", "-b", "-d", "-o", "NAME,PATH,TYPE,MODEL,SERIAL,SIZE,ROTA,TRAN")
	if err != nil {
		return nil, err
	}
	return parseLsblk(b)
}

func parseLsblk(b []byte) ([]Disk, error) {
	var o lsblkOut
	if err := json.Unmarshal(b, &o); err != nil {
		return nil, err
	}
	var out []Disk
	for _, d := range o.Blockdevices {
		if d.Type != "disk" {
			continue
		}
		if strings.HasPrefix(d.Name, "zram") || strings.HasPrefix(d.Name, "loop") ||
			strings.HasPrefix(d.Name, "dm-") || strings.HasPrefix(d.Name, "zd") || strings.HasPrefix(d.Name, "sr") {
			continue
		}
		path := d.Path
		if path == "" {
			path = "/dev/" + d.Name
		}
		out = append(out, Disk{
			Name: d.Name, Path: path, Model: strings.TrimSpace(d.Model),
			Serial: strings.TrimSpace(d.Serial), Size: uint64(d.Size), Rota: bool(d.Rota), Tran: d.Tran,
		})
	}
	return out, nil
}