package smart

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/anatol/smart.go"
)

// SMART data structure for SATA disks
type SATAAttrData struct {
	Id      uint8
	Name    string
	Type    int
	RawVal  uint64
	Current uint8
	Worst   uint8
}

type SMARTData struct {
	ModelNumber  string
	SerialNumber string
	Size         uint64
	Temperature  int

	/* NVME specific fields */
	NVMENameSpace   []smart.NvmeIdentNamespace
	PowerOnHours    uint64
	PowerCycles     uint64
	UnsafeShutdowns uint64
	MediaErrors     uint64

	/* SATA specific fields */
	SATAAttrs []*SATAAttrData
}

// Get SMART data of a particular disk / device
func GetSMARTData(disk string) (*SMARTData, error) {
	if !strings.HasPrefix(disk, "/dev/") {
		disk = filepath.Join("/dev", disk)
	}

	//Check if the disk exists
	if _, err := os.Stat(disk); os.IsNotExist(err) {
		return nil, fmt.Errorf("disk %s does not exist", disk)
	}

	// Check if the disk is NVMe or SATA
	isNVMe := strings.HasPrefix(disk, "/dev/nvme")
	isSATA := strings.HasPrefix(disk, "/dev/sd")

	// If the disk is not NVMe or SATA, return an error
	if !isNVMe && !isSATA {
		return nil, fmt.Errorf("disk %s is not an NVMe or SATA disk", disk)
	}

	if isNVMe {
		return getNVMESMART(disk)
	} else if isSATA {
		return getSATASMART(disk)
	}

	return nil, errors.New("unsupported disk type")
}

func getSATASMART(diskpath string) (*SMARTData, error) {
	dev, err := smart.OpenSata(diskpath)
	if err != nil {
		return nil, err
	}
	defer dev.Close()

	c, err := dev.Identify()
	if err != nil {
		return nil, err
	}

	sm, err := dev.ReadSMARTData()
	if err != nil {
		return nil, err
	}

	_, capacity, _, _, _ := c.Capacity()
	sataAttrs := []*SATAAttrData{}
	for _, attr := range sm.Attrs {
		thisAttrID := attr.Id
		thisAttrName := attr.Name
		thisAttrType := attr.Type
		thisAttrRawVal := attr.ValueRaw
		val, low, high, _, err := attr.ParseAsTemperature()
		if err != nil {
			continue
		}

		fmt.Println("Temperature: ", val, low, high)
		thisAttrData := SATAAttrData{
			Id:      thisAttrID,
			Name:    thisAttrName,
			Type:    thisAttrType,
			RawVal:  thisAttrRawVal,
			Current: attr.Current,
			Worst:   attr.Worst,
		}

		sataAttrs = append(sataAttrs, &thisAttrData)
	}

	smartData := SMARTData{
		ModelNumber:  c.ModelNumber(),
		SerialNumber: c.SerialNumber(),
		Size:         capacity,

		SATAAttrs: sataAttrs,
	}

	return &smartData, nil
}

func getNVMESMART(diskpath string) (*SMARTData, error) {
	dev, err := smart.OpenNVMe(diskpath)
	if err != nil {
		return nil, err
	}
	defer dev.Close()

	c, nss, err := dev.Identify()
	if err != nil {
		return nil, err
	}

	sm, _ := dev.ReadSMART()
	smartData := SMARTData{
		ModelNumber:     c.ModelNumber(),
		SerialNumber:    c.SerialNumber(),
		Size:            c.Tnvmcap.Val[0],
		Temperature:     int(sm.Temperature),
		PowerOnHours:    sm.PowerOnHours.Val[0],
		PowerCycles:     sm.PowerCycles.Val[0],
		UnsafeShutdowns: sm.UnsafeShutdowns.Val[0],
		MediaErrors:     sm.MediaErrors.Val[0],

		NVMENameSpace: nss,
	}

	return &smartData, nil
}