package main

import (
	"context"
	"encoding/json"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

type DiskMetrics struct {
	Name              string       `json:"name"`
	Path              string       `json:"path"`
	Model             string       `json:"model"`
	Serial            string       `json:"serial"`
	Type              string       `json:"type"`
	Protocol          string       `json:"protocol"`
	SizeBytes         uint64       `json:"sizeBytes"`
	UsedBytes         *uint64      `json:"usedBytes"`
	UsagePercent      *float64     `json:"usagePercent"`
	SMARTStatus       string       `json:"smartStatus"`
	SMARTAvailable    bool         `json:"smartAvailable"`
	AttentionReasons  []string     `json:"attentionReasons"`
	SMARTHistory      []string     `json:"smartHistory"`
	Temperature       *float64     `json:"temperatureCelsius"`
	PowerOnHours      *uint64      `json:"powerOnHours"`
	WearUsedPercent   *float64     `json:"wearUsedPercent"`
	Reallocated       uint64       `json:"reallocated"`
	Pending           uint64       `json:"pending"`
	Uncorrectable     uint64       `json:"uncorrectable"`
	CRCErrors         uint64       `json:"crcErrors"`
	MediaErrors       uint64       `json:"mediaErrors"`
	SMARTErrorCount   uint64       `json:"smartErrorCount"`
	CounterChanges    []string     `json:"counterChanges"`
	LastSelfTest      string       `json:"lastSelfTest"`
	LastSelfTestAt    string       `json:"lastSelfTestAt"`
	SelfTestNeverRun  bool         `json:"selfTestNeverRun"`
	ShortTestStarted  bool         `json:"shortTestStarted"`
	ReadIOPS          float64      `json:"readIops"`
	WriteIOPS         float64      `json:"writeIops"`
	ReadBytesPerSec   float64      `json:"readBytesPerSec"`
	WriteBytesPerSec  float64      `json:"writeBytesPerSec"`
	ReadLatencyMs     float64      `json:"readLatencyMs"`
	WriteLatencyMs    float64      `json:"writeLatencyMs"`
	LatencyMs         float64      `json:"latencyMs"`
	Utilization       float64      `json:"utilizationPercent"`
	QueueDepth        uint64       `json:"queueDepth"`
	AverageQueue      float64      `json:"averageQueue"`
	IOPressureSeconds float64      `json:"ioPressureSeconds"`
	IOCause           string       `json:"ioCause"`
	IOCauseDetail     string       `json:"ioCauseDetail"`
	IOConfidence      int          `json:"ioConfidence"`
	TopConsumers      []IOConsumer `json:"topConsumers"`
}

type diskIOSnapshot struct{ reads, readSectors, readMS, writes, writeSectors, writeMS, inFlight, ioMS, weightedMS uint64 }

type diskCollector struct {
	mu              sync.Mutex
	cached          []DiskMetrics
	updatedAt       time.Time
	previous        map[string]DiskMetrics
	missing         []string
	lastTests       map[string]time.Time
	ioPrevious      map[string]diskIOSnapshot
	ioUpdatedAt     time.Time
	ioPressureSince map[string]time.Time
}

func (c *diskCollector) collect(thresholds AlertThresholds) []DiskMetrics {
	c.mu.Lock()
	defer c.mu.Unlock()
	if time.Since(c.updatedAt) < time.Minute && c.cached != nil {
		c.applyDiskIO(c.cached, thresholds)
		return c.cached
	}
	current := readPhysicalDisks()
	if c.previous == nil {
		c.previous = map[string]DiskMetrics{}
		c.lastTests = map[string]time.Time{}
	}
	seen := map[string]bool{}
	c.missing = nil
	for i := range current {
		disk := &current[i]
		key := disk.Serial
		if key == "" || key == "Не указан" {
			key = disk.Path
		}
		seen[key] = true
		if old, ok := c.previous[key]; ok {
			disk.CounterChanges = diskCounterChanges(old, *disk)
			disk.AttentionReasons = append(disk.AttentionReasons, disk.CounterChanges...)
		}
		if disk.SMARTAvailable && (disk.SelfTestNeverRun || selfTestOlderThanWeek(disk.LastSelfTestAt)) && time.Since(c.lastTests[key]) >= 7*24*time.Hour {
			disk.ShortTestStarted = startShortSMARTTest(disk.Path)
			if disk.ShortTestStarted {
				c.lastTests[key] = time.Now()
			}
		}
		c.previous[key] = *disk
	}
	for key, disk := range c.previous {
		if !seen[key] {
			c.missing = append(c.missing, disk.Model+" ("+disk.Serial+")")
		}
	}
	c.cached = current
	c.applyDiskIO(c.cached, thresholds)
	c.updatedAt = time.Now()
	return c.cached
}

func (c *diskCollector) applyDiskIO(disks []DiskMetrics, thresholds AlertThresholds) {
	now := time.Now()
	snapshots := readDiskIOStats("/proc/diskstats")
	if c.ioPrevious == nil {
		c.ioPrevious = map[string]diskIOSnapshot{}
		c.ioPressureSince = map[string]time.Time{}
	}
	elapsed := now.Sub(c.ioUpdatedAt).Seconds()
	for i := range disks {
		disk := &disks[i]
		disk.ReadIOPS, disk.WriteIOPS, disk.ReadBytesPerSec, disk.WriteBytesPerSec = 0, 0, 0, 0
		disk.ReadLatencyMs, disk.WriteLatencyMs, disk.LatencyMs, disk.Utilization, disk.AverageQueue, disk.IOPressureSeconds = 0, 0, 0, 0, 0, 0
		current, ok := snapshots[disk.Name]
		if !ok {
			continue
		}
		disk.QueueDepth = current.inFlight
		if old, found := c.ioPrevious[disk.Name]; found && elapsed > 0 && diskIOCountersMonotonic(old, current) {
			readOps := current.reads - old.reads
			writeOps := current.writes - old.writes
			disk.ReadIOPS = float64(readOps) / elapsed
			disk.WriteIOPS = float64(writeOps) / elapsed
			disk.ReadBytesPerSec = float64(current.readSectors-old.readSectors) * 512 / elapsed
			disk.WriteBytesPerSec = float64(current.writeSectors-old.writeSectors) * 512 / elapsed
			if readOps > 0 {
				disk.ReadLatencyMs = float64(current.readMS-old.readMS) / float64(readOps)
			}
			if writeOps > 0 {
				disk.WriteLatencyMs = float64(current.writeMS-old.writeMS) / float64(writeOps)
			}
			if readOps+writeOps > 0 {
				disk.LatencyMs = float64(current.readMS-old.readMS+current.writeMS-old.writeMS) / float64(readOps+writeOps)
			}
			disk.Utilization = min(float64(current.ioMS-old.ioMS)/(elapsed*10), 100)
			disk.AverageQueue = float64(current.weightedMS-old.weightedMS) / (elapsed * 1000)
			utilLimit, latencyLimit, queueLimit := thresholds.DiskSSDUtil, thresholds.DiskSSDLatency, thresholds.DiskSSDQueue
			if disk.Type == "HDD" {
				utilLimit, latencyLimit, queueLimit = thresholds.DiskHDDUtil, thresholds.DiskHDDLatency, thresholds.DiskHDDQueue
			} else if strings.Contains(disk.Type, "NVMe") {
				utilLimit, latencyLimit, queueLimit = thresholds.DiskNVMeUtil, thresholds.DiskNVMeLatency, thresholds.DiskNVMeQueue
			}
			pressured := disk.Utilization >= utilLimit || (disk.LatencyMs >= latencyLimit && disk.ReadIOPS+disk.WriteIOPS >= 1) || disk.AverageQueue >= queueLimit
			if pressured {
				if c.ioPressureSince[disk.Name].IsZero() {
					c.ioPressureSince[disk.Name] = now
				}
				disk.IOPressureSeconds = now.Sub(c.ioPressureSince[disk.Name]).Seconds()
			} else {
				delete(c.ioPressureSince, disk.Name)
			}
		}
		c.ioPrevious[disk.Name] = current
	}
	c.ioUpdatedAt = now
}

func diskIOCountersMonotonic(old, current diskIOSnapshot) bool {
	return current.reads >= old.reads && current.readSectors >= old.readSectors && current.readMS >= old.readMS && current.writes >= old.writes && current.writeSectors >= old.writeSectors && current.writeMS >= old.writeMS && current.ioMS >= old.ioMS && current.weightedMS >= old.weightedMS
}

func readDiskIOStats(path string) map[string]diskIOSnapshot {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil
	}
	result := map[string]diskIOSnapshot{}
	for _, line := range strings.Split(string(data), "\n") {
		f := strings.Fields(line)
		if len(f) < 14 || !isPhysicalDiskName(f[2]) {
			continue
		}
		values := make([]uint64, 11)
		for i := 0; i < 11; i++ {
			values[i], _ = strconv.ParseUint(f[i+3], 10, 64)
		}
		result[f[2]] = diskIOSnapshot{reads: values[0], readSectors: values[2], readMS: values[3], writes: values[4], writeSectors: values[6], writeMS: values[7], inFlight: values[8], ioMS: values[9], weightedMS: values[10]}
	}
	return result
}

func (c *diskCollector) missingDisks() []string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return append([]string(nil), c.missing...)
}

func diskCounterChanges(old, current DiskMetrics) []string {
	checks := []struct {
		label    string
		old, now uint64
	}{{"Reallocated вырос", old.Reallocated, current.Reallocated}, {"Pending вырос", old.Pending, current.Pending}, {"Uncorrectable вырос", old.Uncorrectable, current.Uncorrectable}, {"CRC-ошибки выросли", old.CRCErrors, current.CRCErrors}, {"Media errors выросли", old.MediaErrors, current.MediaErrors}, {"Ошибки SMART-журнала выросли", old.SMARTErrorCount, current.SMARTErrorCount}}
	var changes []string
	for _, check := range checks {
		if check.now > check.old {
			changes = append(changes, check.label+": "+strconv.FormatUint(check.old, 10)+" → "+strconv.FormatUint(check.now, 10))
		}
	}
	return changes
}

func selfTestOlderThanWeek(value string) bool {
	when, err := time.Parse(time.RFC3339, value)
	return err == nil && time.Since(when) > 7*24*time.Hour
}

func startShortSMARTTest(path string) bool {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	return exec.CommandContext(ctx, "smartctl", "-t", "short", path).Run() == nil
}

func readPhysicalDisks() []DiskMetrics {
	entries, _ := os.ReadDir("/sys/block")
	usage := readDiskUsage()
	var disks []DiskMetrics
	for _, entry := range entries {
		name := entry.Name()
		if !isPhysicalDiskName(name) {
			continue
		}
		base := filepath.Join("/sys/block", name)
		sectors, _ := strconv.ParseUint(readTrimmed(filepath.Join(base, "size")), 10, 64)
		rotational := readTrimmed(filepath.Join(base, "queue/rotational"))
		disk := DiskMetrics{
			Name: name, Path: "/dev/" + name,
			Model:     readTrimmed(filepath.Join(base, "device/model")),
			Serial:    readTrimmed(filepath.Join(base, "device/serial")),
			SizeBytes: sectors * 512, SMARTStatus: "Недоступен",
		}
		if diskUsage, ok := usage[name]; ok && diskUsage.total > 0 {
			disk.UsedBytes = &diskUsage.used
			percent := float64(diskUsage.used) / float64(diskUsage.total) * 100
			disk.UsagePercent = &percent
		}
		if strings.HasPrefix(name, "nvme") {
			disk.Type, disk.Protocol = "NVMe SSD", "NVMe"
		} else if rotational == "1" {
			disk.Type = "HDD"
		} else {
			disk.Type = "SSD"
		}
		applySMART(&disk)
		if disk.Model == "" {
			disk.Model = disk.Name
		}
		if disk.Serial == "" {
			disk.Serial = "Не указан"
		}
		disks = append(disks, disk)
	}
	return disks
}

type diskUsage struct {
	used  uint64
	total uint64
}

type lsblkDevice struct {
	Name     string          `json:"name"`
	FSUsed   json.RawMessage `json:"fsused"`
	FSAvail  json.RawMessage `json:"fsavail"`
	Children []lsblkDevice   `json:"children"`
}

func readDiskUsage() map[string]diskUsage {
	output, err := exec.Command("lsblk", "--json", "--bytes", "--output", "NAME,FSUSED,FSAVAIL").Output()
	if err != nil {
		return nil
	}
	var data struct {
		Devices []lsblkDevice `json:"blockdevices"`
	}
	if json.Unmarshal(output, &data) != nil {
		return nil
	}
	result := make(map[string]diskUsage)
	for _, device := range data.Devices {
		used, total, ok := sumFilesystemUsage(device)
		if ok {
			result[device.Name] = diskUsage{used: used, total: total}
		}
	}
	return result
}

func sumFilesystemUsage(device lsblkDevice) (uint64, uint64, bool) {
	used, usedOK := parseJSONUint(device.FSUsed)
	available, availableOK := parseJSONUint(device.FSAvail)
	if usedOK || availableOK {
		return used, used + available, true
	}
	var totalUsed, totalSize uint64
	found := false
	for _, child := range device.Children {
		childUsed, childTotal, childOK := sumFilesystemUsage(child)
		if childOK {
			totalUsed += childUsed
			totalSize += childTotal
			found = true
		}
	}
	return totalUsed, totalSize, found
}

func parseJSONUint(raw json.RawMessage) (uint64, bool) {
	value := strings.Trim(strings.TrimSpace(string(raw)), `"`)
	if value == "" || value == "null" {
		return 0, false
	}
	parsed, err := strconv.ParseUint(value, 10, 64)
	return parsed, err == nil
}

func isPhysicalDiskName(name string) bool {
	return strings.HasPrefix(name, "sd") ||
		(strings.HasPrefix(name, "nvme") && strings.Contains(name, "n"))
}

func readTrimmed(path string) string {
	data, err := os.ReadFile(path)
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(data))
}

type smartctlJSON struct {
	Smartctl struct {
		ExitStatus int `json:"exit_status"`
		Messages   []struct {
			String   string `json:"string"`
			Severity string `json:"severity"`
		} `json:"messages"`
	} `json:"smartctl"`
	ModelName    string `json:"model_name"`
	SerialNumber string `json:"serial_number"`
	Device       struct {
		Protocol string `json:"protocol"`
	} `json:"device"`
	SmartStatus *struct {
		Passed bool `json:"passed"`
	} `json:"smart_status"`
	Temperature struct {
		Current float64 `json:"current"`
	} `json:"temperature"`
	PowerOnTime struct {
		Hours uint64 `json:"hours"`
	} `json:"power_on_time"`
	ATAAttributes struct {
		Table []struct {
			Name       string `json:"name"`
			Value      int    `json:"value"`
			Threshold  int    `json:"thresh"`
			WhenFailed string `json:"when_failed"`
			Raw        struct {
				String string `json:"string"`
				Value  int64  `json:"value"`
			} `json:"raw"`
		} `json:"table"`
	} `json:"ata_smart_attributes"`
	ATAErrorLog struct {
		Summary struct {
			Count int `json:"count"`
		} `json:"summary"`
	} `json:"ata_smart_error_log"`
	ATASelfTestLog struct {
		Standard struct {
			Table []struct {
				Status struct {
					String string `json:"string"`
				} `json:"status"`
				LifetimeHours uint64 `json:"lifetime_hours"`
			} `json:"table"`
		} `json:"standard"`
	} `json:"ata_smart_self_test_log"`
	NVMeHealth *struct {
		Temperature     float64 `json:"temperature"`
		PercentageUsed  float64 `json:"percentage_used"`
		PowerOnHours    uint64  `json:"power_on_hours"`
		CriticalWarning int     `json:"critical_warning"`
		MediaErrors     uint64  `json:"media_errors"`
	} `json:"nvme_smart_health_information_log"`
}

func applySMART(disk *DiskMetrics) {
	if _, err := exec.LookPath("smartctl"); err != nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	output, _ := exec.CommandContext(ctx, "smartctl", "-a", "-j", disk.Path).Output()
	if len(output) == 0 {
		return
	}
	var smart smartctlJSON
	if json.Unmarshal(output, &smart) != nil {
		return
	}
	if smart.SmartStatus != nil {
		disk.SMARTAvailable = true
		disk.SMARTStatus = "Исправен"
		if !smart.SmartStatus.Passed {
			disk.SMARTStatus = "Ошибка"
			disk.AttentionReasons = append(disk.AttentionReasons, "Общая проверка SMART завершилась ошибкой")
		}
	}
	if smart.ModelName != "" {
		disk.Model = smart.ModelName
	}
	if smart.SerialNumber != "" {
		disk.Serial = smart.SerialNumber
	}
	if smart.Device.Protocol != "" {
		disk.Protocol = smart.Device.Protocol
	}
	if smart.Temperature.Current > 0 {
		disk.Temperature = &smart.Temperature.Current
	} else if smart.NVMeHealth != nil && smart.NVMeHealth.Temperature > 0 {
		disk.Temperature = &smart.NVMeHealth.Temperature
	}
	if smart.PowerOnTime.Hours > 0 {
		disk.PowerOnHours = &smart.PowerOnTime.Hours
	} else if smart.NVMeHealth != nil && smart.NVMeHealth.PowerOnHours > 0 {
		disk.PowerOnHours = &smart.NVMeHealth.PowerOnHours
	}
	if smart.NVMeHealth != nil {
		disk.WearUsedPercent = &smart.NVMeHealth.PercentageUsed
		disk.MediaErrors = smart.NVMeHealth.MediaErrors
		if smart.NVMeHealth.CriticalWarning != 0 {
			disk.AttentionReasons = append(disk.AttentionReasons, "NVMe сообщает критическое предупреждение")
		}
		if smart.NVMeHealth.MediaErrors > 0 {
			disk.AttentionReasons = append(disk.AttentionReasons, "Ошибки целостности носителя: "+strconv.FormatUint(smart.NVMeHealth.MediaErrors, 10))
		}
	}
	if len(smart.ATASelfTestLog.Standard.Table) > 0 {
		last := smart.ATASelfTestLog.Standard.Table[0]
		disk.LastSelfTest = firstNonEmpty(last.Status.String, "Результат неизвестен")
		if disk.PowerOnHours != nil && *disk.PowerOnHours >= last.LifetimeHours {
			disk.LastSelfTestAt = time.Now().Add(-time.Duration(*disk.PowerOnHours-last.LifetimeHours) * time.Hour).Format(time.RFC3339)
		}
	} else {
		disk.SelfTestNeverRun = true
		disk.LastSelfTest = "Никогда не запускался"
		disk.AttentionReasons = append(disk.AttentionReasons, "SMART self-test ещё ни разу не запускался")
	}
	applyATAAttributes(disk, smart)
	applySMARTExitStatus(disk, smart.Smartctl.ExitStatus)
	if len(disk.AttentionReasons) > 0 && disk.SMARTStatus == "Исправен" {
		disk.SMARTStatus = "Требует внимания"
	}
}

func applyATAAttributes(disk *DiskMetrics, smart smartctlJSON) {
	wearNames := map[string]bool{
		"Percent_Lifetime_Remain": true, "SSD_Life_Left": true,
		"Media_Wearout_Indicator": true, "Remaining_Lifetime_Perc": true,
		"Wear_Leveling_Count": true,
	}
	for _, attribute := range smart.ATAAttributes.Table {
		switch attribute.Name {
		case "Reallocated_Sector_Ct", "Reallocated_Event_Count":
			disk.Reallocated += uint64(max(attribute.Raw.Value, 0))
		case "Current_Pending_Sector":
			disk.Pending = uint64(max(attribute.Raw.Value, 0))
		case "Offline_Uncorrectable", "Reported_Uncorrect":
			disk.Uncorrectable += uint64(max(attribute.Raw.Value, 0))
		case "UDMA_CRC_Error_Count", "Interface_CRC_Error_Count":
			disk.CRCErrors += uint64(max(attribute.Raw.Value, 0))
		}
		if wearNames[attribute.Name] && attribute.Value >= 0 && attribute.Value <= 100 {
			wear := float64(100 - attribute.Value)
			disk.WearUsedPercent = &wear
		}
		if strings.EqualFold(attribute.WhenFailed, "FAILING_NOW") {
			reason := attribute.Name + " вышел за порог"
			if attribute.Raw.String != "" {
				reason += " (" + attribute.Raw.String + ")"
			}
			disk.AttentionReasons = append(disk.AttentionReasons, reason)
		} else if attribute.WhenFailed != "" && attribute.WhenFailed != "-" {
			reason := attribute.Name + " ранее выходил за порог"
			if attribute.Raw.String != "" {
				reason += " (" + attribute.Raw.String + ")"
			}
			disk.SMARTHistory = append(disk.SMARTHistory, reason)
		}
	}
	if smart.ATAErrorLog.Summary.Count > 0 {
		disk.SMARTErrorCount = uint64(smart.ATAErrorLog.Summary.Count)
		disk.SMARTHistory = append(disk.SMARTHistory, "Записей в SMART-журнале: "+strconv.Itoa(smart.ATAErrorLog.Summary.Count))
	}
}

func applySMARTExitStatus(disk *DiskMetrics, status int) {
	if status&8 != 0 {
		disk.SMARTStatus = "Ошибка"
		disk.AttentionReasons = append(disk.AttentionReasons, "SMART сообщает возможный скорый отказ диска")
	}
	if status&16 != 0 {
		disk.SMARTStatus = "Ошибка"
		disk.AttentionReasons = append(disk.AttentionReasons, "Критический SMART-атрибут достиг порога")
	}
	if status&32 != 0 {
		disk.SMARTHistory = append(disk.SMARTHistory, "SMART-атрибут ранее находился ниже порога")
	}
	if status&64 != 0 {
		disk.SMARTHistory = append(disk.SMARTHistory, "В журнале SMART есть старые записи об ошибках")
	}
	if status&128 != 0 {
		disk.SMARTHistory = append(disk.SMARTHistory, "В журнале самотестирования есть старые ошибки")
	}
}