package collect

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

type SmartAttr struct {
	ID         int    `json:"id"`
	Name       string `json:"name"`
	Value      int    `json:"value"`
	Worst      int    `json:"worst"`
	Thresh     int    `json:"thresh"`
	Raw        int64  `json:"raw"`
	RawString  string `json:"raw_string"`
	WhenFailed string `json:"when_failed"`
}

type NVMeHealth struct {
	CriticalWarning int   `json:"critical_warning"`
	AvailableSpare  int   `json:"available_spare"`
	SpareThreshold  int   `json:"spare_threshold"`
	PercentageUsed  int   `json:"percentage_used"`
	MediaErrors     int64 `json:"media_errors"`
	ErrLogEntries   int64 `json:"err_log_entries"`
	UnsafeShutdowns int64 `json:"unsafe_shutdowns"`
}

type SelfTest struct {
	Type       string `json:"type"`
	Status     string `json:"status"`
	Passed     bool   `json:"passed"`
	Hours      int64  `json:"hours"`
	InProgress bool   `json:"in_progress"`
	Remaining  int    `json:"remaining_pct"`
}

// Smart is the normalised view of one smartctl -j -a run.
type Smart struct {
	Device      string      `json:"device"`
	Protocol    string      `json:"protocol"` // ATA | NVMe | SCSI
	Model       string      `json:"model"`
	Serial      string      `json:"serial"`
	Firmware    string      `json:"firmware"`
	Capacity    uint64      `json:"capacity"`
	Standby     bool        `json:"standby"`   // disk asleep, nothing was read
	Supported   bool        `json:"supported"` // SMART available & enabled
	Healthy     bool        `json:"healthy"`
	TempC       int         `json:"temp_c"`
	HasTemp     bool        `json:"has_temp"`
	PowerOnHrs  int64       `json:"power_on_hours"`
	PowerCycles int64       `json:"power_cycles"`
	Attrs       []SmartAttr `json:"attrs,omitempty"`
	NVMe        *NVMeHealth `json:"nvme,omitempty"`
	LastTest    *SelfTest   `json:"last_test,omitempty"`
	Error       string      `json:"error,omitempty"`
	CheckedAt   time.Time   `json:"checked_at"`
}

// Attr returns the ATA attribute with the given id, if present.
func (s *Smart) Attr(id int) (SmartAttr, bool) {
	for _, a := range s.Attrs {
		if a.ID == id {
			return a, true
		}
	}
	return SmartAttr{}, false
}

// ReadSmart runs smartctl for one device. skipStandby adds "-n standby" so a
// spun-down disk is not woken; then Standby=true and nothing else is filled.
func ReadSmart(dev string, skipStandby bool) Smart {
	args := []string{"-j", "-a"}
	if skipStandby {
		args = append(args, "-n", "standby")
	}
	args = append(args, dev)
	out, err := run(60*time.Second, "smartctl", args...)
	s := parseSmart(out)
	s.Device = dev
	s.CheckedAt = time.Now()
	if len(out) == 0 && err != nil {
		s.Error = err.Error()
	}
	return s
}

type smartMsg struct {
	String   string `json:"string"`
	Severity string `json:"severity"`
}

type smartJSON struct {
	Smartctl struct {
		ExitStatus int        `json:"exit_status"`
		Messages   []smartMsg `json:"messages"`
	} `json:"smartctl"`
	Device struct {
		Protocol string `json:"protocol"`
	} `json:"device"`
	ModelName    string `json:"model_name"`
	SerialNumber string `json:"serial_number"`
	Firmware     string `json:"firmware_version"`
	UserCapacity struct {
		Bytes uint64 `json:"bytes"`
	} `json:"user_capacity"`
	SmartSupport struct {
		Available bool `json:"available"`
		Enabled   bool `json:"enabled"`
	} `json:"smart_support"`
	SmartStatus *struct {
		Passed bool `json:"passed"`
	} `json:"smart_status"`
	Temperature struct {
		Current *int `json:"current"`
	} `json:"temperature"`
	PowerOnTime struct {
		Hours int64 `json:"hours"`
	} `json:"power_on_time"`
	PowerCycleCount int64 `json:"power_cycle_count"`
	ATAAttrs        struct {
		Table []struct {
			ID         int    `json:"id"`
			Name       string `json:"name"`
			Value      int    `json:"value"`
			Worst      int    `json:"worst"`
			Thresh     int    `json:"thresh"`
			WhenFailed string `json:"when_failed"`
			Raw        struct {
				Value  int64  `json:"value"`
				String string `json:"string"`
			} `json:"raw"`
		} `json:"table"`
	} `json:"ata_smart_attributes"`
	ATASelfTestLog struct {
		Standard struct {
			Table []struct {
				Type struct {
					String string `json:"string"`
				} `json:"type"`
				Status struct {
					String string `json:"string"`
					Passed bool   `json:"passed"`
				} `json:"status"`
				LifetimeHours int64 `json:"lifetime_hours"`
			} `json:"table"`
		} `json:"standard"`
	} `json:"ata_smart_self_test_log"`
	ATASmartData struct {
		SelfTest struct {
			Status struct {
				String           string `json:"string"`
				RemainingPercent *int   `json:"remaining_percent"`
			} `json:"status"`
		} `json:"self_test"`
	} `json:"ata_smart_data"`
	NVMeLog *struct {
		CriticalWarning  int   `json:"critical_warning"`
		AvailableSpare   int   `json:"available_spare"`
		SpareThreshold   int   `json:"available_spare_threshold"`
		PercentageUsed   int   `json:"percentage_used"`
		MediaErrors      int64 `json:"media_errors"`
		NumErrLogEntries int64 `json:"num_err_log_entries"`
		UnsafeShutdowns  int64 `json:"unsafe_shutdowns"`
		PowerOnHours     int64 `json:"power_on_hours"`
		PowerCycles      int64 `json:"power_cycles"`
	} `json:"nvme_smart_health_information_log"`
	NVMeSelfTestLog struct {
		Current struct {
			Value  int    `json:"value"`
			String string `json:"string"`
		} `json:"current_self_test_operation"`
		CurrentCompletion *int `json:"current_self_test_completion_percent"`
		Table             []struct {
			Code struct {
				String string `json:"string"`
			} `json:"self_test_code"`
			Result struct {
				Value  int    `json:"value"`
				String string `json:"string"`
			} `json:"self_test_result"`
			PowerOnHours int64 `json:"power_on_hours"`
		} `json:"table"`
	} `json:"nvme_self_test_log"`
}

func parseSmart(b []byte) Smart {
	var s Smart
	var j smartJSON
	if err := json.Unmarshal(b, &j); err != nil {
		s.Error = fmt.Sprintf("parse smartctl output: %v", err)
		return s
	}
	for _, m := range j.Smartctl.Messages {
		if strings.Contains(strings.ToUpper(m.String), "STANDBY") {
			s.Standby = true
			return s
		}
	}
	// bit 0: command line error, bit 1: device open failed
	if j.Smartctl.ExitStatus&3 != 0 {
		for _, m := range j.Smartctl.Messages {
			if m.Severity == "error" {
				s.Error = m.String
				break
			}
		}
		if s.Error == "" {
			s.Error = fmt.Sprintf("smartctl exit status %d", j.Smartctl.ExitStatus)
		}
	}
	s.Protocol = j.Device.Protocol
	s.Model = j.ModelName
	s.Serial = j.SerialNumber
	s.Firmware = j.Firmware
	s.Capacity = j.UserCapacity.Bytes
	s.Supported = (j.SmartSupport.Available && j.SmartSupport.Enabled) || j.SmartStatus != nil
	if j.SmartStatus != nil {
		s.Healthy = j.SmartStatus.Passed
	}
	if j.Temperature.Current != nil {
		s.TempC, s.HasTemp = *j.Temperature.Current, true
	}
	s.PowerOnHrs = j.PowerOnTime.Hours
	s.PowerCycles = j.PowerCycleCount
	for _, a := range j.ATAAttrs.Table {
		raw := a.Raw.Value
		if a.ID == 188 {
			// Seagate packs three 16-bit counters into 188; the low word is the timeout count
			raw &= 0xFFFF
		}
		s.Attrs = append(s.Attrs, SmartAttr{
			ID: a.ID, Name: a.Name, Value: a.Value, Worst: a.Worst, Thresh: a.Thresh,
			Raw: raw, RawString: a.Raw.String, WhenFailed: a.WhenFailed,
		})
	}
	if j.NVMeLog != nil {
		s.NVMe = &NVMeHealth{
			CriticalWarning: j.NVMeLog.CriticalWarning,
			AvailableSpare:  j.NVMeLog.AvailableSpare,
			SpareThreshold:  j.NVMeLog.SpareThreshold,
			PercentageUsed:  j.NVMeLog.PercentageUsed,
			MediaErrors:     j.NVMeLog.MediaErrors,
			ErrLogEntries:   j.NVMeLog.NumErrLogEntries,
			UnsafeShutdowns: j.NVMeLog.UnsafeShutdowns,
		}
		if s.PowerOnHrs == 0 {
			s.PowerOnHrs = j.NVMeLog.PowerOnHours
		}
		if s.PowerCycles == 0 {
			s.PowerCycles = j.NVMeLog.PowerCycles
		}
	}
	// self-test: an in-progress test wins, otherwise the newest log entry
	if rp := j.ATASmartData.SelfTest.Status.RemainingPercent; rp != nil {
		s.LastTest = &SelfTest{Status: j.ATASmartData.SelfTest.Status.String, InProgress: true, Remaining: *rp, Passed: true}
	} else if t := j.ATASelfTestLog.Standard.Table; len(t) > 0 {
		s.LastTest = &SelfTest{Type: t[0].Type.String, Status: t[0].Status.String, Passed: t[0].Status.Passed, Hours: t[0].LifetimeHours}
	} else if j.NVMeSelfTestLog.Current.Value != 0 {
		rem := 0
		if j.NVMeSelfTestLog.CurrentCompletion != nil {
			rem = 100 - *j.NVMeSelfTestLog.CurrentCompletion
		}
		s.LastTest = &SelfTest{Status: j.NVMeSelfTestLog.Current.String, InProgress: true, Remaining: rem, Passed: true}
	} else if t := j.NVMeSelfTestLog.Table; len(t) > 0 {
		s.LastTest = &SelfTest{Type: t[0].Code.String, Status: t[0].Result.String, Passed: t[0].Result.Value == 0, Hours: t[0].PowerOnHours}
	}
	return s
}

// StartSelfTest launches a short or long SMART self-test on the device.
func StartSelfTest(dev, kind string) error {
	if kind != "short" && kind != "long" {
		return fmt.Errorf("unknown self-test kind %q", kind)
	}
	out, err := run(30*time.Second, "smartctl", "-j", "-t", kind, dev)
	if err != nil {
		var j smartJSON
		if json.Unmarshal(out, &j) == nil {
			for _, m := range j.Smartctl.Messages {
				if m.Severity == "error" {
					return fmt.Errorf("smartctl -t %s %s: %s", kind, dev, m.String)
				}
			}
		}
		return fmt.Errorf("smartctl -t %s %s: %w", kind, dev, err)
	}
	return nil
}