size.go

package typing

import (
	"bytes"
	"encoding/json"
	"fmt"
	"math"
	"strconv"
	"strings"
)

const (
	B  Size = 1
	KB      = B << 10
	MB      = KB << 10
	GB      = MB << 10
	TB      = GB << 10
	PB      = TB << 10
	EB      = PB << 10
)

const (
	unitB  = "B"
	unitKB = "KB"
	unitMB = "MB"
	unitGB = "GB"
	unitTB = "TB"
	unitPB = "PB"
	unitEB = "EB"
)

type Size int64

func (s Size) Bytes() int64 {
	return int64(s)
}

func (s Size) KiloBytes() int64 {
	return int64(s / KB)
}

func (s Size) MegaBytes() int64 {
	return int64(s / MB)
}

func (s Size) GigaBytes() int64 {
	return int64(s / GB)
}

func (s Size) TeraBytes() int64 {
	return int64(s / TB)
}

func (s Size) PetaBytes() int64 {
	return int64(s / PB)
}

func (s Size) ExaBytes() int64 {
	return int64(s / EB)
}

func (s Size) String() string {
	return s.Format(2, "", true)
}

func (s Size) Format(prec int, sep string, unitUpper bool) string {
	var base Size
	switch {
	case s >= EB:
		base = EB
	case s >= PB:
		base = PB
	case s >= TB:
		base = TB
	case s >= GB:
		base = GB
	case s >= MB:
		base = MB
	case s >= KB:
		base = KB
	default:
		base = B
	}
	unit := unitOfSize(base)
	if !unitUpper {
		unit = strings.ToLower(unit)
	}
	if base == B {
		return strconv.FormatInt(int64(s), 10) + sep + unit
	}
	str := strconv.FormatFloat(float64(s)/float64(base), 'f', prec, 64)
	str = strings.TrimRight(strings.TrimRight(str, "0"), ".")
	return str + sep + unit
}

func (s Size) MarshalText() ([]byte, error) {
	return []byte(s.String()), nil
}

func (s *Size) UnmarshalText(b []byte) error {
	v, err := ParseSize(string(b))
	if err != nil {
		return err
	}
	*s = v
	return nil
}

func (s Size) MarshalJSON() ([]byte, error) {
	return json.Marshal(s.String())
}

func (s *Size) UnmarshalJSON(b []byte) error {
	b = bytes.TrimSpace(b)

	if len(b) == 0 || bytes.Equal(b, []byte("null")) {
		*s = 0
		return nil
	}

	switch b[0] {
	case '"':
		// JSON string: "10MB", "1.5GB"
		var str string
		if err := json.Unmarshal(b, &str); err != nil {
			return err
		}
		v, err := ParseSize(str)
		if err != nil {
			return err
		}
		*s = v
		return nil

	default:
		// JSON number: 1024, 1024.5
		var num json.Number
		dec := json.NewDecoder(bytes.NewReader(b))
		dec.UseNumber()

		if err := dec.Decode(&num); err != nil {
			return err
		}

		// 优先 int64（避免 float 精度问题）
		if i, err := num.Int64(); err == nil {
			*s = Size(i)
			return nil
		}

		// 再退化为 float64
		f, err := num.Float64()
		if err != nil {
			return fmt.Errorf("invalid size number: %w", err)
		}
		*s = Size(math.Round(f))
		return nil
	}
}

func unitOfSize(base Size) string {
	switch base {
	case EB:
		return unitEB
	case PB:
		return unitPB
	case TB:
		return unitTB
	case GB:
		return unitGB
	case MB:
		return unitMB
	case KB:
		return unitKB
	default:
		return unitB
	}
}

func splitUnits(s string) (string, string) {
	s = strings.TrimSpace(s)
	i := 0
	sepFound := false

	for i < len(s) {
		b := s[i]
		if b == '.' {
			if sepFound {
				break
			}
			sepFound = true
		} else if b < '0' || b > '9' {
			break
		}
		i++
	}

	return strings.TrimSpace(s[:i]), strings.TrimSpace(s[i:])
}

func ParseSize(size string) (Size, error) {
	s, units := splitUnits(size)
	f, err := strconv.ParseFloat(s, 64)
	if err != nil {
		return 0, err
	}
	if len(units) > 5 {
		return 0, fmt.Errorf("unknown units")
	}

	switch strings.ToLower(units) {
	case "b", "bytes", "":
		return Size(f), nil
	case "k", "kb", "kib":
		return Size(math.Round(f * float64(KB))), nil
	case "m", "mb", "mib":
		return Size(math.Round(f * float64(MB))), nil
	case "g", "gb", "gib":
		return Size(math.Round(f * float64(GB))), nil
	case "t", "tb", "tib":
		return Size(math.Round(f * float64(TB))), nil
	case "p", "pb", "pib":
		return Size(math.Round(f * float64(PB))), nil
	default:
		return 0, fmt.Errorf("unknown size units: %s", units)
	}
}