package main

import (
	"context"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// PoolColumns are the properties `zpool list` is asked for, in order.
var PoolColumns = []string{"name", "size", "allocated", "free", "capacity", "fragmentation", "health", "dedupratio"}

// PoolsArgv lists every imported pool.
func PoolsArgv() []string {
	return []string{"list", "-Hp", "-o", strings.Join(PoolColumns, ",")}
}

// StatusArgv is every imported pool's status, numbers exact.
func StatusArgv() []string { return []string{"status", "-p"} }

// Pool is one pool as `zpool list` answers it.
type Pool struct {
	Name      string `json:"name"`
	Size      uint64 `json:"size"`
	Allocated uint64 `json:"allocated"`
	Free      uint64 `json:"free"`
	// Capacity is the share of the pool allocated, in %.
	Capacity int `json:"capacity_pct"`
	// Fragmentation is the free space's fragmentation in %, nil when the pool cannot say.
	Fragmentation *int    `json:"fragmentation_pct,omitempty"`
	Health        string  `json:"health"`
	DedupRatio    float64 `json:"dedupratio"`
}

// ParsePools reads `zpool list -Hp -o <PoolColumns>`.
func ParsePools(out string) ([]Pool, error) {
	var ps []Pool
	for _, line := range lines(out) {
		f := strings.Split(line, "\t")
		if len(f) != len(PoolColumns) {
			return nil, fmt.Errorf("zpool list answered a line with %d columns, not %d: %q", len(f), len(PoolColumns), line)
		}
		p := Pool{Name: f[0], Health: f[6]}
		var err error
		if p.Size, err = number(f[1]); err != nil {
			return nil, fmt.Errorf("zpool list: size of %s: %v", p.Name, err)
		}
		if p.Allocated, err = number(f[2]); err != nil {
			return nil, fmt.Errorf("zpool list: allocated of %s: %v", p.Name, err)
		}
		if p.Free, err = number(f[3]); err != nil {
			return nil, fmt.Errorf("zpool list: free of %s: %v", p.Name, err)
		}
		if p.Capacity, err = strconv.Atoi(strings.TrimSuffix(f[4], "%")); err != nil {
			return nil, fmt.Errorf("zpool list: capacity of %s is %q", p.Name, f[4])
		}
		if fr, e := strconv.Atoi(strings.TrimSuffix(f[5], "%")); e == nil {
			p.Fragmentation = &fr
		}
		if p.DedupRatio, err = ratio(f[7]); err != nil {
			return nil, fmt.Errorf("zpool list: dedupratio of %s is %q", p.Name, f[7])
		}
		ps = append(ps, p)
	}
	return ps, nil
}

// PoolRow is one pool's space as the tools answer it.
type PoolRow struct {
	Name          string `json:"name"`
	Health        string `json:"health"`
	Size          uint64 `json:"size_bytes"`
	Allocated     uint64 `json:"allocated_bytes"`
	Free          uint64 `json:"free_bytes"`
	Capacity      int    `json:"capacity_pct"`
	Fragmentation *int   `json:"fragmentation_pct,omitempty"`
	Said          string `json:"said"`
}

// Row says a pool's space in a line.
func (p Pool) Row() PoolRow {
	r := PoolRow{Name: p.Name, Health: p.Health, Size: p.Size, Allocated: p.Allocated, Free: p.Free, Capacity: p.Capacity,
		Fragmentation: p.Fragmentation}
	frag := "fragmentation not known"
	if p.Fragmentation != nil {
		frag = fmt.Sprintf("%d %% fragmented", *p.Fragmentation)
	}
	r.Said = fmt.Sprintf("pool %s (%s): %s of %s allocated (%d %%), %s free, %s", p.Name, p.Health, bytesSaid(p.Allocated),
		bytesSaid(p.Size), p.Capacity, bytesSaid(p.Free), frag)
	return r
}

// Pools are every imported pool.
func (m *Machine) Pools(ctx context.Context) ([]Pool, error) {
	out, err := m.ask(ctx, "zpool", PoolsArgv()...)
	if err != nil {
		return nil, err
	}
	return ParsePools(out)
}

// Device is one line of a pool's configuration: the pool itself, a vdev (mirror, raidz), a disk, or a
// section (logs, cache, spares, special, dedup), with its children.
type Device struct {
	Name string `json:"name"`
	// Kind is pool, vdev (a group of disks), device (a disk or a partition), or section.
	Kind string `json:"kind"`
	// Role is data, log, cache, spare, special or dedup: the section a device is in.
	Role     string    `json:"role,omitempty"`
	State    string    `json:"state,omitempty"`
	Read     uint64    `json:"read_errors"`
	Write    uint64    `json:"write_errors"`
	Checksum uint64    `json:"checksum_errors"`
	Note     string    `json:"note,omitempty"`
	Children []*Device `json:"children,omitempty"`
}

// Errors is the error count status gives this line, read, write and checksum together: a vdev's line
// counts what was found at the vdev, not the sum of its disks.
func (d *Device) Errors() uint64 { return d.Read + d.Write + d.Checksum }

// Scan is a pool's last scan, as status says it: ZFS keeps only the latest, a scrub's or a resilver's.
type Scan struct {
	// Kind is scrub or resilver; "" when none was ever asked.
	Kind string `json:"kind,omitempty"`
	// State is finished, running, paused, canceled or none.
	State    string     `json:"state"`
	Started  *time.Time `json:"started,omitempty"`
	Ended    *time.Time `json:"ended,omitempty"`
	Took     string     `json:"took,omitempty"`
	Repaired string     `json:"repaired,omitempty"`
	Errors   *int       `json:"errors,omitempty"`
	Done     *float64   `json:"done_pct,omitempty"`
	ToGo     string     `json:"to_go,omitempty"`
	Said     string     `json:"said"`
}

// PoolStatus is one pool as `zpool status` says it.
type PoolStatus struct {
	Name   string `json:"name"`
	State  string `json:"state"`
	Status string `json:"status,omitempty"`
	Action string `json:"action,omitempty"`
	See    string `json:"see,omitempty"`
	Scan   Scan   `json:"scan"`
	// Other are the lines status gives for a removal, a checkpoint or an expansion.
	Other  map[string]string `json:"other,omitempty"`
	Root   *Device           `json:"layout,omitempty"`
	Errors string            `json:"errors"`
}

var sections = map[string]string{"logs": "log", "cache": "cache", "spares": "spare", "special": "special", "dedup": "dedup"}

// ParseStatus reads `zpool status -p` for every pool it shows.
func ParseStatus(out string) ([]PoolStatus, error) {
	var pools []PoolStatus
	var cur *PoolStatus
	key := ""
	inConfig := false
	var stack []*Device
	section := ""
	flush := func() {
		if cur != nil {
			cur.Scan = parseScan(cur.Scan.Said)
			pools = append(pools, *cur)
		}
	}
	for _, raw := range strings.Split(out, "\n") {
		line := strings.TrimRight(raw, "\r ")
		trimmed := strings.TrimSpace(line)
		if k, v, ok := statusKey(line); ok {
			inConfig = false
			key = k
			switch k {
			case "pool":
				flush()
				cur = &PoolStatus{Name: v}
				stack, section = nil, ""
			case "config":
				inConfig = true
			}
			if cur == nil {
				continue
			}
			switch k {
			case "state":
				cur.State = v
			case "status":
				cur.Status = v
			case "action":
				cur.Action = v
			case "see":
				cur.See = v
			case "scan":
				cur.Scan.Said = v
			case "errors":
				cur.Errors = v
			case "pool", "config":
			default:
				if cur.Other == nil {
					cur.Other = map[string]string{}
				}
				cur.Other[k] = v
			}
			continue
		}
		if cur == nil || trimmed == "" {
			continue
		}
		if !inConfig {
			// A continuation of the key before it, indented by a tab.
			add := func(s *string) { *s = strings.TrimSpace(*s + " " + trimmed) }
			switch key {
			case "status":
				add(&cur.Status)
			case "action":
				add(&cur.Action)
			case "scan":
				cur.Scan.Said += "\n" + trimmed
			case "errors":
				add(&cur.Errors)
			default:
				if cur.Other != nil {
					if v, ok := cur.Other[key]; ok {
						cur.Other[key] = strings.TrimSpace(v + " " + trimmed)
					}
				}
			}
			continue
		}
		// The configuration: a tab, two spaces per level, then the columns.
		body := strings.TrimPrefix(line, "\t")
		depth := (len(body) - len(strings.TrimLeft(body, " "))) / 2
		f := strings.Fields(body)
		if f[0] == "NAME" && len(f) > 1 && f[1] == "STATE" {
			continue
		}
		d := &Device{Name: f[0]}
		if depth == 0 {
			if role, ok := sections[f[0]]; ok && len(f) == 1 {
				d.Kind, d.Role, section = "section", role, role
				if cur.Root != nil {
					cur.Root.Children = append(cur.Root.Children, d)
				}
				// A section stands at the pool's level and holds what is indented under it.
				stack = []*Device{d}
				continue
			}
		}
		if len(f) > 1 {
			d.State = f[1]
		}
		rest := f[min(len(f), 2):]
		if len(rest) >= 3 && isCount(rest[0]) && isCount(rest[1]) && isCount(rest[2]) {
			d.Read, d.Write, d.Checksum = count(rest[0]), count(rest[1]), count(rest[2])
			rest = rest[3:]
		}
		d.Note = strings.Join(rest, " ")
		if depth == 0 {
			d.Kind, d.Role = "pool", ""
			cur.Root = d
			stack, section = []*Device{d}, ""
			continue
		}
		role := section
		if role == "" {
			role = "data"
		}
		d.Role = role
		if depth < len(stack) {
			stack = stack[:depth]
		}
		if len(stack) == 0 {
			return nil, fmt.Errorf("zpool status: the device %s of pool %s has no parent", d.Name, cur.Name)
		}
		parent := stack[len(stack)-1]
		parent.Children = append(parent.Children, d)
		stack = append(stack, d)
	}
	flush()
	for i := range pools {
		kinds(pools[i].Root)
	}
	return pools, nil
}

// kinds names each device by its place: a device with children is a vdev, one without a device.
func kinds(d *Device) {
	if d == nil {
		return
	}
	for _, c := range d.Children {
		if c.Kind == "" {
			if len(c.Children) > 0 {
				c.Kind = "vdev"
			} else {
				c.Kind = "device"
			}
		}
		kinds(c)
	}
}

var keyLine = regexp.MustCompile(`^\s{0,8}([a-z]+):(?:\s(.*))?$`)

// statusKey is a `key: value` line of status: pool, state, status, action, see, scan, config, errors, …
func statusKey(line string) (string, string, bool) {
	if strings.HasPrefix(line, "\t") {
		return "", "", false
	}
	m := keyLine.FindStringSubmatch(line)
	if m == nil {
		return "", "", false
	}
	return m[1], strings.TrimSpace(m[2]), true
}

func isCount(s string) bool {
	if s == "" {
		return false
	}
	_, err := strconv.ParseFloat(strings.TrimRight(s, "KMGTP"), 64)
	return err == nil
}

// count is an error count: exact under -p, else 1.2K and the like.
func count(s string) uint64 {
	mult := 1.0
	switch s[len(s)-1] {
	case 'K':
		mult = 1e3
	case 'M':
		mult = 1e6
	case 'G':
		mult = 1e9
	case 'T':
		mult = 1e12
	case 'P':
		mult = 1e15
	}
	f, _ := strconv.ParseFloat(strings.TrimRight(s, "KMGTP"), 64)
	return uint64(f * mult)
}

var (
	scanFinished = regexp.MustCompile(`^(scrub repaired|resilvered) (\S+) in (.+?) with (\d+) errors on (.+)$`)
	scanRunning  = regexp.MustCompile(`^(scrub|resilver) in progress since (.+)$`)
	scanPaused   = regexp.MustCompile(`^(scrub|resilver) paused since (.+)$`)
	scanCanceled = regexp.MustCompile(`^(scrub|resilver) canceled on (.+)$`)
	scanDone     = regexp.MustCompile(`([\d.]+)% done`)
	scanToGo     = regexp.MustCompile(`(\d+ days \d+:\d+:\d+|\d+:\d+:\d+) to go`)
	scanRepaired = regexp.MustCompile(`(\S+) repaired`)
)

// parseScan reads status's scan line and its continuation.
func parseScan(said string) Scan {
	s := Scan{State: "none", Said: said}
	first, more, _ := strings.Cut(said, "\n")
	first = strings.TrimSpace(first)
	switch {
	case first == "" || first == "none requested":
		return s
	case scanFinished.MatchString(first):
		m := scanFinished.FindStringSubmatch(first)
		s.Kind, s.State, s.Repaired, s.Took = "scrub", "finished", m[2], m[3]
		if m[1] == "resilvered" {
			s.Kind = "resilver"
		}
		n, _ := strconv.Atoi(m[4])
		s.Errors = &n
		s.Ended = when(m[5])
		if s.Ended != nil {
			if d, ok := took(m[3]); ok {
				st := s.Ended.Add(-d)
				s.Started = &st
			}
		}
	case scanRunning.MatchString(first):
		m := scanRunning.FindStringSubmatch(first)
		s.Kind, s.State, s.Started = m[1], "running", when(m[2])
	case scanPaused.MatchString(first):
		m := scanPaused.FindStringSubmatch(first)
		s.Kind, s.State, s.Started = m[1], "paused", when(m[2])
	case scanCanceled.MatchString(first):
		m := scanCanceled.FindStringSubmatch(first)
		s.Kind, s.State, s.Ended = m[1], "canceled", when(m[2])
	default:
		s.State = "unknown"
	}
	if more != "" {
		if m := scanDone.FindStringSubmatch(more); m != nil {
			v, _ := strconv.ParseFloat(m[1], 64)
			s.Done = &v
		}
		if m := scanToGo.FindStringSubmatch(more); m != nil {
			s.ToGo = m[1]
		}
		if m := scanRepaired.FindStringSubmatch(more); m != nil && s.Repaired == "" {
			s.Repaired = m[1]
		}
	}
	return s
}

// when reads status's dates, as ctime writes them: "Sun Sep 13 05:36:34 2026", in local time.
func when(s string) *time.Time {
	s = strings.TrimSpace(s)
	for _, layout := range []string{time.ANSIC, "Mon Jan 2 15:04:05 2006"} {
		if t, err := time.ParseInLocation(layout, s, time.Local); err == nil {
			return &t
		}
	}
	return nil
}

// took reads a scan's duration: "05:12:33" or "1 days 05:12:33".
func took(s string) (time.Duration, bool) {
	days := 0
	if d, rest, ok := strings.Cut(s, " days "); ok {
		n, err := strconv.Atoi(d)
		if err != nil {
			return 0, false
		}
		days, s = n, rest
	}
	p := strings.Split(s, ":")
	if len(p) != 3 {
		return 0, false
	}
	var v [3]int
	for i := range p {
		n, err := strconv.Atoi(p[i])
		if err != nil {
			return 0, false
		}
		v[i] = n
	}
	return time.Duration(days)*24*time.Hour + time.Duration(v[0])*time.Hour + time.Duration(v[1])*time.Minute +
		time.Duration(v[2])*time.Second, true
}

// Status is every pool's status.
func (m *Machine) Status(ctx context.Context) ([]PoolStatus, error) {
	out, err := m.ask(ctx, "zpool", StatusArgv()...)
	if err != nil {
		return nil, err
	}
	return ParseStatus(out)
}

// PoolsAnswer is what zfs_pools answers.
type PoolsAnswer struct {
	Applicable bool         `json:"applicable"`
	Pools      []PoolHealth `json:"pools"`
	Said       []string     `json:"said"`
}

// PoolHealth is one pool's health and space.
type PoolHealth struct {
	PoolRow
	State      string `json:"state"`
	Status     string `json:"status,omitempty"`
	Action     string `json:"action,omitempty"`
	Errors     string `json:"errors"`
	DeviceErrs uint64 `json:"device_errors"`
	Scan       string `json:"last_scan"`
}

// PoolsNow is every pool's health in a line, with its space.
func (m *Machine) PoolsNow(ctx context.Context) (any, error) {
	ps, err := m.Pools(ctx)
	if err != nil {
		return notHere(err)
	}
	st, err := m.Status(ctx)
	if err != nil {
		return nil, err
	}
	a := PoolsAnswer{Applicable: true, Pools: []PoolHealth{}, Said: []string{}}
	for _, p := range ps {
		h := PoolHealth{PoolRow: p.Row(), State: p.Health}
		if s := statusOf(st, p.Name); s != nil {
			h.State, h.Status, h.Action, h.Errors, h.Scan = s.State, s.Status, s.Action, s.Errors, firstLine(s.Scan.Said)
			h.DeviceErrs = sumErrors(s.Root)
		}
		line := h.Said + fmt.Sprintf("; %d device errors; errors: %s", h.DeviceErrs, h.Errors)
		if h.Status != "" {
			line += "; status: " + h.Status
		}
		a.Pools = append(a.Pools, h)
		a.Said = append(a.Said, line)
	}
	if len(ps) == 0 {
		a.Said = append(a.Said, "ZFS is here, and no pool is imported.")
	}
	return a, nil
}

// LayoutAnswer is what zfs_pool_layout answers.
type LayoutAnswer struct {
	Applicable bool         `json:"applicable"`
	Pools      []PoolStatus `json:"pools"`
	Said       []string     `json:"said"`
}

// Layout is every pool's vdevs and devices with their state and errors.
func (m *Machine) Layout(ctx context.Context) (any, error) {
	st, err := m.Status(ctx)
	if err != nil {
		return notHere(err)
	}
	a := LayoutAnswer{Applicable: true, Pools: st, Said: []string{}}
	for _, p := range st {
		a.Said = append(a.Said, fmt.Sprintf("pool %s: %s", p.Name, p.State))
		walk(p.Root, func(d *Device, depth int) {
			if d.Kind == "pool" {
				return
			}
			line := strings.Repeat("  ", depth) + d.Name
			if d.Kind == "section" {
				a.Said = append(a.Said, line+" ("+d.Role+")")
				return
			}
			line += " " + d.State
			if d.Errors() > 0 {
				line += fmt.Sprintf(" — ERRORS read %d, write %d, checksum %d", d.Read, d.Write, d.Checksum)
			}
			if d.Note != "" {
				line += " " + d.Note
			}
			a.Said = append(a.Said, line)
		}, 0)
	}
	return a, nil
}

func walk(d *Device, f func(*Device, int), depth int) {
	if d == nil {
		return
	}
	f(d, depth)
	for _, c := range d.Children {
		walk(c, f, depth+1)
	}
}

func sumErrors(d *Device) uint64 {
	var n uint64
	walk(d, func(x *Device, _ int) {
		if x.Kind != "section" {
			n += x.Errors()
		}
	}, 0)
	return n
}

func statusOf(st []PoolStatus, name string) *PoolStatus {
	for i := range st {
		if st[i].Name == name {
			return &st[i]
		}
	}
	return nil
}