package main

import (
	"errors"
	"fmt"
	"strconv"
	"time"
)

type duration time.Duration

func (d duration) MarshalYAML() (interface{}, error) {
	return d.String(), nil
}

func (d duration) String() string {
	result := ""
	if weeks := time.Duration(d) / (time.Hour * 24 * 7); weeks > 0 {
		result += fmt.Sprintf("%dw", weeks)
		d -= duration(weeks * time.Hour * 24 * 7)
	}
	if days := time.Duration(d) / (time.Hour * 24); days > 0 {
		result += fmt.Sprintf("%dd", days)
		d -= duration(days * time.Hour * 24)
	}
	return result + time.Duration(d).String()
}

func (d *duration) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var s string
	if err := unmarshal(&s); err != nil {
		return err
	}
	count := ""
	var ttl time.Duration
	for i := range s {
		if s[i] < '0' || s[i] > '9' {
			n, err := strconv.Atoi(count)
			if err != nil {
				return err
			}
			count += s[i : i+1]
			switch s[i] {
			case 'w':
				count = fmt.Sprintf("%dh", n*24*7)
			case 'd':
				count = fmt.Sprintf("%dh", n*24)
			}
			d, err := time.ParseDuration(count)
			if err != nil {
				return err
			}
			ttl += d
			count = ""
		} else {
			count += s[i : i+1]
		}
	}
	if count != "" {
		return errors.New(count)
	}
	*d = duration(ttl)
	return nil
}
