package main

import (
	"flag"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"time"

	"code.justin.tv/common/chitin"
	_ "code.justin.tv/common/golibs/bininfo"
)

func main() {
	var period time.Duration
	flag.DurationVar(&period, "period", 10*time.Second, "Delay between cleanup cycles")
	var reportCount int
	flag.IntVar(&reportCount, "keep", 30, "Number of reports to keep")
	var byteCount int64
	flag.Int64Var(&byteCount, "bytes", 100*1000*1000*1000, "Number of raw data bytes to keep")
	var dataDir string
	flag.StringVar(&dataDir, "data", "", "Path to data directory")
	var reportDir string
	flag.StringVar(&reportDir, "reports", "", "Path to reports directory")
	flag.Parse()

	err := chitin.ExperimentalTraceProcessOptIn()
	if err != nil {
		log.Fatalf("trace enable err=%q", err)
	}

	// we'll do the cleanup periodically
	tick := time.NewTimer(period)
	defer tick.Stop()
	// the first run should happen immediately
	tick.Reset(1 * time.Nanosecond)

	for {
		// time to go!
		<-tick.C
		// TODO: jitter
		tick.Reset(period)

		task := cleanupTask{
			reportCount: reportCount,
			byteCount:   byteCount,
			dataDir:     dataDir,
			reportDir:   reportDir,
		}

		err := task.run()
		if err != nil {
			log.Printf("cleanup err=%q", err)
		}
	}
}

type cleanupTask struct {
	reportCount int
	byteCount   int64
	dataDir     string
	reportDir   string
}

func (task *cleanupTask) run() error {
	log.Printf("cleanup-state=%q", "begin")
	defer log.Printf("cleanup-state=%q", "end")

	log.Printf("visit dir=%q", task.dataDir)
	if task.dataDir != "" {
		err := cleanupBySize(task.dataDir, task.byteCount)
		if err != nil {
			return err
		}
	}

	log.Printf("visit dir=%q", task.reportDir)
	if task.reportDir != "" {
		err := cleanupByCount(task.reportDir, task.reportCount)
		if err != nil {
			return err
		}
	}

	return nil
}

func cleanupBySize(dir string, byteCount int64) error {
	infos, err := ioutil.ReadDir(dir)
	if err != nil {
		return err
	}
	sort.Sort(sort.Reverse(fileInfoByTime(infos)))

	// We have to keep at least the two most recent files - one is being
	// written, and the other is the most recent complete file (which is used
	// for running reports).
	const minCount = 2

	var keeping int
	var bytesSaved int64
	for _, info := range infos {
		fullname := filepath.Join(dir, info.Name())
		bytesSaved += info.Size()

		if keeping < minCount {
			log.Printf("keep=%q", fullname)
			keeping++
			continue
		}
		if bytesSaved <= byteCount {
			log.Printf("keep=%q", fullname)
			keeping++
			continue
		}

		log.Printf("delete=%q", fullname)
		err := os.RemoveAll(fullname)
		if err != nil {
			return err
		}
	}

	return nil
}

func cleanupByCount(dir string, keep int) error {
	infos, err := ioutil.ReadDir(dir)
	if err != nil {
		return err
	}
	sort.Sort(sort.Reverse(fileInfoByTime(infos)))

	var keeping int
	for _, info := range infos {
		fullname := filepath.Join(dir, info.Name())
		if keeping < keep {
			log.Printf("keep=%q", fullname)
			keeping++
			continue
		}
		log.Printf("delete=%q", fullname)
		err := os.RemoveAll(fullname)
		if err != nil {
			return err
		}
	}

	return nil
}

type fileInfoByTime []os.FileInfo

func (l fileInfoByTime) Len() int           { return len(l) }
func (l fileInfoByTime) Less(i, j int) bool { return l[i].ModTime().Before(l[j].ModTime()) }
func (l fileInfoByTime) Swap(i, j int)      { l[i], l[j] = l[j], l[i] }
