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"`
	Temperature      *float64 `json:"temperatureCelsius"`
	PowerOnHours     *uint64  `json:"powerOnHours"`
	WearUsedPercent  *float64 `json:"wearUsedPercent"`
}

type diskCollector struct {
	mu        sync.Mutex
	cached    []DiskMetrics
	updatedAt time.Time
}

func (c *diskCollector) collect() []DiskMetrics {
	c.mu.Lock()
	defer c.mu.Unlock()
	if time.Since(c.updatedAt) < time.Minute && c.cached != nil {
		return c.cached
	}
	c.cached = readPhysicalDisks()
	c.updatedAt = time.Now()
	return c.cached
}

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"`
	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
		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))
		}
	}
	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 {
		if wearNames[attribute.Name] && attribute.Value >= 0 && attribute.Value <= 100 {
			wear := float64(100 - attribute.Value)
			disk.WearUsedPercent = &wear
		}
		if attribute.WhenFailed != "" && attribute.WhenFailed != "-" {
			reason := attribute.Name + " вышел за порог"
			if attribute.Raw.String != "" {
				reason += " (" + attribute.Raw.String + ")"
			}
			disk.AttentionReasons = append(disk.AttentionReasons, reason)
		}
	}
	if smart.ATAErrorLog.Summary.Count > 0 {
		disk.AttentionReasons = append(disk.AttentionReasons, "Ошибок в 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.AttentionReasons = append(disk.AttentionReasons, "SMART-атрибут ранее находился ниже порога")
	}
	if status&64 != 0 {
		disk.AttentionReasons = append(disk.AttentionReasons, "В журнале SMART обнаружены ошибки")
	}
	if status&128 != 0 {
		disk.AttentionReasons = append(disk.AttentionReasons, "В журнале самотестирования есть ошибки")
	}
}