package collect

import (
	"bufio"
	"bytes"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Vdev struct {
	Name  string `json:"name"`
	Depth int    `json:"depth"`
	State string `json:"state"`
	Read  int64  `json:"read"`
	Write int64  `json:"write"`
	Cksum int64  `json:"cksum"`
	Note  string `json:"note,omitempty"`
}

type Pool struct {
	Name      string     `json:"name"`
	Size      uint64     `json:"size"`
	Alloc     uint64     `json:"alloc"`
	Free      uint64     `json:"free"`
	CapPct    int        `json:"cap_pct"`
	Health    string     `json:"health"`
	FragPct   int        `json:"frag_pct"`
	Scan      string     `json:"scan"`
	ScanErrs  int64      `json:"scan_errors"`
	LastScrub *time.Time `json:"last_scrub,omitempty"`
	Scrubbing bool       `json:"scrubbing"`
	Errors    string     `json:"errors"`
	Status    string     `json:"status,omitempty"`
	Action    string     `json:"action,omitempty"`
	Vdevs     []Vdev     `json:"vdevs"`
}

type Dataset struct {
	Name       string `json:"name"`
	Used       uint64 `json:"used"`
	Avail      uint64 `json:"avail"`
	Refer      uint64 `json:"refer"`
	Mountpoint string `json:"mountpoint"`
}

// ReadPools combines zpool list and zpool status for every pool.
func ReadPools() ([]Pool, error) {
	out, err := run(20*time.Second, "zpool", "list", "-Hp", "-o", "name,size,alloc,free,cap,health,frag")
	if err != nil {
		return nil, err
	}
	pools := parseZpoolList(out)
	for i := range pools {
		st, err := run(20*time.Second, "zpool", "status", pools[i].Name)
		if err != nil && len(st) == 0 {
			continue
		}
		parseZpoolStatus(st, &pools[i])
	}
	return pools, nil
}

func ReadDatasets() ([]Dataset, error) {
	out, err := run(20*time.Second, "zfs", "list", "-Hp", "-t", "filesystem", "-o", "name,used,avail,refer,mountpoint")
	if err != nil {
		return nil, err
	}
	return parseZfsList(out), nil
}

func parseZpoolList(b []byte) []Pool {
	var pools []Pool
	sc := bufio.NewScanner(bytes.NewReader(b))
	for sc.Scan() {
		f := strings.Split(sc.Text(), "\t")
		if len(f) < 7 {
			continue
		}
		p := Pool{Name: f[0], Health: f[5]}
		p.Size, _ = strconv.ParseUint(f[1], 10, 64)
		p.Alloc, _ = strconv.ParseUint(f[2], 10, 64)
		p.Free, _ = strconv.ParseUint(f[3], 10, 64)
		p.CapPct = atoiPct(f[4])
		p.FragPct = atoiPct(f[6])
		if p.CapPct == 0 && p.Size > 0 {
			p.CapPct = int(p.Alloc * 100 / p.Size)
		}
		pools = append(pools, p)
	}
	return pools
}

func atoiPct(s string) int {
	n, _ := strconv.Atoi(strings.TrimSuffix(strings.TrimSpace(s), "%"))
	return n
}

var (
	reSection  = regexp.MustCompile(`^ *(pool|state|status|action|see|scan|scrub|config|errors|remove|checkpoint|expand):(.*)$`)
	reScanDate = regexp.MustCompile(`on (\w{3} \w{3}\s+\d+ \d+:\d+:\d+ \d{4})`)
	reScanErrs = regexp.MustCompile(`with (\d+) errors`)
)

// parseZpoolStatus fills scan/errors/status/vdev tree from the text output of
// "zpool status <pool>" (zfs 2.2 has no JSON output).
func parseZpoolStatus(b []byte, p *Pool) {
	sc := bufio.NewScanner(bytes.NewReader(b))
	section := ""
	var scanLines, statusLines, actionLines []string
	inConfig, headerSeen := false, false
	for sc.Scan() {
		line := sc.Text()
		trim := strings.TrimSpace(line)
		if inConfig {
			if trim == "" {
				if headerSeen {
					inConfig = false
				}
				continue
			}
			if !headerSeen {
				if strings.HasPrefix(trim, "NAME") {
					headerSeen = true
				}
				continue
			}
			p.Vdevs = append(p.Vdevs, parseVdevLine(line))
			continue
		}
		if m := reSection.FindStringSubmatch(line); m != nil {
			key, rest, _ := strings.Cut(trim, ":")
			key = strings.TrimSpace(key)
			rest = strings.TrimSpace(rest)
			section = key
			switch key {
			case "config":
				inConfig, headerSeen = true, false
			case "scan", "scrub":
				scanLines = append(scanLines, rest)
			case "status":
				statusLines = append(statusLines, rest)
			case "action":
				actionLines = append(actionLines, rest)
			case "errors":
				p.Errors = rest
			case "state":
				if p.Health == "" {
					p.Health = rest
				}
			}
			continue
		}
		// continuation line of the current section
		switch section {
		case "scan", "scrub":
			scanLines = append(scanLines, trim)
		case "status":
			statusLines = append(statusLines, trim)
		case "action":
			actionLines = append(actionLines, trim)
		case "errors":
			if trim != "" {
				p.Errors += " " + trim
			}
		}
	}
	p.Scan = strings.TrimSpace(strings.Join(scanLines, " "))
	p.Status = strings.TrimSpace(strings.Join(statusLines, " "))
	p.Action = strings.TrimSpace(strings.Join(actionLines, " "))
	p.Scrubbing = strings.Contains(p.Scan, "scrub in progress")
	if m := reScanErrs.FindStringSubmatch(p.Scan); m != nil {
		p.ScanErrs, _ = strconv.ParseInt(m[1], 10, 64)
	}
	if strings.HasPrefix(p.Scan, "scrub repaired") || strings.HasPrefix(p.Scan, "scrub canceled") {
		if m := reScanDate.FindStringSubmatch(p.Scan); m != nil {
			if t, err := time.ParseInLocation("Mon Jan _2 15:04:05 2006", squashSpaces(m[1]), time.Local); err == nil {
				p.LastScrub = &t
			}
		}
	}
}

func squashSpaces(s string) string { return strings.Join(strings.Fields(s), " ") }

func parseVdevLine(line string) Vdev {
	// zpool indents with a tab, then two spaces per level
	body := strings.TrimLeft(line, "\t")
	indent := len(body) - len(strings.TrimLeft(body, " "))
	f := strings.Fields(body)
	v := Vdev{Depth: indent / 2}
	if len(f) == 0 {
		return v
	}
	v.Name = f[0]
	if len(f) >= 5 {
		v.State = f[1]
		v.Read, _ = strconv.ParseInt(f[2], 10, 64)
		v.Write, _ = strconv.ParseInt(f[3], 10, 64)
		v.Cksum, _ = strconv.ParseInt(f[4], 10, 64)
		if len(f) > 5 {
			v.Note = strings.Join(f[5:], " ")
		}
	} else if len(f) > 1 {
		v.State = f[1]
	}
	return v
}

func parseZfsList(b []byte) []Dataset {
	var out []Dataset
	sc := bufio.NewScanner(bytes.NewReader(b))
	for sc.Scan() {
		f := strings.Split(sc.Text(), "\t")
		if len(f) < 5 {
			continue
		}
		d := Dataset{Name: f[0], Mountpoint: f[4]}
		d.Used, _ = strconv.ParseUint(f[1], 10, 64)
		d.Avail, _ = strconv.ParseUint(f[2], 10, 64)
		d.Refer, _ = strconv.ParseUint(f[3], 10, 64)
		out = append(out, d)
	}
	return out
}