//go:build linux

package storage

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
)

// lsblkOutput matches the top-level JSON from lsblk -J.
type lsblkOutput struct {
	BlockDevices []lsblkDevice `json:"blockdevices"`
}

// lsblkDevice is the raw JSON structure from lsblk.
type lsblkDevice struct {
	Name       string        `json:"name"`
	Path       string        `json:"path"`
	Size       interface{}   `json:"size"` // may be float64 or string
	Type       string        `json:"type"`
	FSType     *string       `json:"fstype"`
	MountPoint *string       `json:"mountpoint"`
	Model      *string       `json:"model"`
	RM         interface{}   `json:"rm"` // removable: bool or "0"/"1"
	Children   []lsblkDevice `json:"children"`
}

func (d *lsblkDevice) sizeBytes() int64 {
	switch v := d.Size.(type) {
	case float64:
		return int64(v)
	}
	return 0
}

func (d *lsblkDevice) sizeHuman() string {
	bytes := d.sizeBytes()
	const (
		GB = 1024 * 1024 * 1024
		TB = GB * 1024
	)
	switch {
	case bytes >= TB:
		return fmt.Sprintf("%.1f TB", float64(bytes)/float64(TB))
	case bytes >= GB:
		return fmt.Sprintf("%.1f GB", float64(bytes)/float64(GB))
	default:
		return fmt.Sprintf("%d MB", bytes/(1024*1024))
	}
}

func (d *lsblkDevice) isRemovable() bool {
	switch v := d.RM.(type) {
	case bool:
		return v
	case float64:
		return v != 0
	case string:
		return v == "1" || strings.EqualFold(v, "true")
	}
	return false
}

func (d *lsblkDevice) fsType() string {
	if d.FSType != nil {
		return *d.FSType
	}
	return ""
}

func (d *lsblkDevice) mountPoint() string {
	if d.MountPoint != nil {
		return *d.MountPoint
	}
	return ""
}

func (d *lsblkDevice) model() string {
	if d.Model != nil {
		return strings.TrimSpace(*d.Model)
	}
	return ""
}

// getSystemDiskNames returns the set of parent disk names (e.g., "sda")
// that contain system partitions (/, /boot, /boot/efi, swap).
// It reads the host's fstab (mounted at /host-fstab in the container)
// and resolves UUIDs to device paths via blkid.
func getSystemDiskNames() map[string]bool {
	systemDisks := map[string]bool{}

	// Step 1: Find and parse fstab
	fstabPath := "/host-fstab"
	if _, err := os.Stat(fstabPath); err != nil {
		fstabPath = "/etc/fstab"
	}

	data, err := os.ReadFile(fstabPath)
	if err != nil {
		return systemDisks
	}

	systemMounts := map[string]bool{"/": true, "/boot": true, "/boot/efi": true}

	var systemUUIDs []string
	var systemDevices []string

	for _, line := range strings.Split(string(data), "\n") {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		fields := strings.Fields(line)
		if len(fields) < 3 {
			continue
		}
		source := fields[0]
		mountPoint := fields[1]
		fsType := fields[2]

		if !systemMounts[mountPoint] && fsType != "swap" {
			continue
		}

		if strings.HasPrefix(source, "UUID=") {
			systemUUIDs = append(systemUUIDs, strings.TrimPrefix(source, "UUID="))
		} else if strings.HasPrefix(source, "/dev/") {
			systemDevices = append(systemDevices, source)
		}
	}

	// Step 2: Resolve UUIDs to device paths via blkid
	for _, uuid := range systemUUIDs {
		out, err := exec.Command("blkid", "-U", uuid).Output()
		if err == nil {
			devPath := strings.TrimSpace(string(out))
			if devPath != "" {
				systemDevices = append(systemDevices, devPath)
			}
		}
	}

	// Step 3: Extract parent disk names from device paths
	for _, devPath := range systemDevices {
		diskName := partitionToParentDisk(devPath)
		if diskName != "" {
			systemDisks[diskName] = true
		}
	}

	return systemDisks
}

// partitionToParentDisk extracts the parent disk name from a partition device path.
// "/dev/sda2" → "sda", "/dev/nvme0n1p2" → "nvme0n1"
func partitionToParentDisk(devPath string) string {
	name := filepath.Base(devPath)

	// NVMe: nvme0n1p2 → nvme0n1
	if strings.Contains(name, "nvme") {
		if idx := strings.LastIndex(name, "p"); idx > 0 {
			if _, err := strconv.Atoi(name[idx+1:]); err == nil {
				return name[:idx]
			}
		}
		return name
	}

	// Standard: sda2 → sda, sdb1 → sdb
	return strings.TrimRight(name, "0123456789")
}

// blkidInfo holds filesystem metadata from blkid.
type blkidInfo struct {
	FSType string
	UUID   string
	Label  string
}

// parseBlkidExport parses the output of `blkid -o export` into a map keyed by device path.
// Output format: blocks separated by blank lines, each line "KEY=VALUE".
func parseBlkidExport(data []byte) map[string]blkidInfo {
	result := map[string]blkidInfo{}

	for _, block := range strings.Split(string(data), "\n\n") {
		var devName string
		info := blkidInfo{}
		for _, line := range strings.Split(strings.TrimSpace(block), "\n") {
			parts := strings.SplitN(line, "=", 2)
			if len(parts) != 2 {
				continue
			}
			key, val := parts[0], parts[1]
			switch key {
			case "DEVNAME":
				devName = val
			case "TYPE":
				info.FSType = val
			case "UUID":
				info.UUID = val
			case "LABEL":
				info.Label = val
			}
		}
		if devName != "" {
			result[devName] = info
		}
	}
	return result
}

// enrichWithBlkid fills in missing FSType, UUID, and Label on partitions using blkid.
// blkid directly probes devices, which works in privileged containers where lsblk's
// udev/blkid cache is unavailable.
func enrichWithBlkid(disks []BlockDevice) {
	out, err := exec.Command("blkid", "-o", "export").Output()
	if err != nil {
		return
	}

	blkidMap := parseBlkidExport(out)

	for i := range disks {
		for j := range disks[i].Partitions {
			p := &disks[i].Partitions[j]
			if info, ok := blkidMap[p.Path]; ok {
				if p.FSType == "" {
					p.FSType = info.FSType
				}
				if p.UUID == "" {
					p.UUID = info.UUID
				}
				if p.Label == "" {
					p.Label = info.Label
				}
			}
		}
	}
}

// ScanDisks detects all block devices and classifies them into
// available (not mounted, not system) and system/mounted disks.
func ScanDisks() (*ScanResult, error) {
	out, err := exec.Command(
		"lsblk", "-J", "-b",
		"-o", "NAME,PATH,SIZE,TYPE,FSTYPE,MOUNTPOINT,MODEL,RM",
	).Output()
	if err != nil {
		return nil, fmt.Errorf("lsblk failed: %w", err)
	}

	var parsed lsblkOutput
	if err := json.Unmarshal(out, &parsed); err != nil {
		return nil, fmt.Errorf("lsblk JSON parse failed: %w", err)
	}

	// Get system disk names from host fstab (works correctly inside container)
	systemDiskNames := getSystemDiskNames()

	result := &ScanResult{}

	for _, dev := range parsed.BlockDevices {
		if dev.Type != "disk" {
			continue
		}

		bd := BlockDevice{
			Name:      dev.Name,
			Path:      dev.Path,
			Size:      dev.sizeHuman(),
			SizeBytes: dev.sizeBytes(),
			Model:     dev.model(),
			Type:      dev.Type,
			Removable: dev.isRemovable(),
		}
		if bd.Path == "" {
			bd.Path = "/dev/" + bd.Name
		}

		anyMounted := false
		for _, child := range dev.Children {
			if child.Type != "part" && child.Type != "lvm" && child.Type != "crypt" {
				continue
			}
			part := Partition{
				Name:       child.Name,
				Path:       child.Path,
				Size:       child.sizeHuman(),
				SizeBytes:  child.sizeBytes(),
				FSType:     child.fsType(),
				MountPoint: child.mountPoint(),
			}
			if part.Path == "" {
				part.Path = "/dev/" + part.Name
			}
			bd.Partitions = append(bd.Partitions, part)
			if part.MountPoint != "" {
				anyMounted = true
			}
		}

		// Also check if the disk itself is directly mounted (no partition table)
		if dev.mountPoint() != "" {
			anyMounted = true
		}

		isSystem := systemDiskNames[dev.Name]
		bd.Mounted = anyMounted || isSystem

		if isSystem || anyMounted {
			result.SystemDisks = append(result.SystemDisks, bd)
		} else {
			result.AvailableDisks = append(result.AvailableDisks, bd)
		}
	}

	// Enrich FSType, UUID, Label from blkid (lsblk can't probe fstype in container)
	enrichWithBlkid(result.AvailableDisks)
	enrichWithBlkid(result.SystemDisks)

	return result, nil
}