package main

import (
	"errors"
	"fmt"
	"github.com/prometheus/tsdb/fileutil"
	"os"
	"sort"
	"strconv"

	kitlog "github.com/go-kit/kit/log"
	"github.com/prometheus/tsdb/wal"
)

type segmentRef struct {
	name  string
	index int
}

// copied from wal.go
func segments(dir string) (first, last int, err error) {
	refs, err := listSegments(dir)
	if err != nil {
		return 0, 0, err
	}
	if len(refs) == 0 {
		return -1, -1, nil
	}
	return refs[0].index, refs[len(refs)-1].index, nil
}

// copied from wal.go
func listSegments(dir string) (refs []segmentRef, err error) {
	files, err := fileutil.ReadDir(dir)
	if err != nil {
		return nil, err
	}
	var last int
	for _, fn := range files {
		k, err := strconv.Atoi(fn)
		if err != nil {
			continue
		}
		if len(refs) > 0 && k > last+1 {
			return nil, errors.New("segments are not sequential")
		}
		refs = append(refs, segmentRef{name: fn, index: k})
		last = k
	}
	sort.Slice(refs, func(i, j int) bool {
		return refs[i].index < refs[j].index
	})
	return refs, nil
}

func main() {
	var err error
	var readErr *wal.CorruptionErr
	var seg *wal.Segment

	first, last, err := segments("bad-data")
	checkErr(err)

segmentLoop:
	for i := first; i <= last; i++ {
		fmt.Println("looking for corruption in segment: ", i)
		seg, err = wal.OpenReadSegment(wal.SegmentName("bad-data", i))
		checkErr(err)
		reader := wal.NewReader(seg)

		checkErr(err)
		for reader.Next() {
			// Nothing.
		}
		if reader.Err() != nil {
			fmt.Println("r.total: ", reader.Offset())
			fmt.Println("error: ", reader.Err())
			readErr = &wal.CorruptionErr{
				Dir:     "bad-data",
				Segment: i,
				Offset:  reader.Offset(),
				Err:     reader.Err(),
			}
			// Encountered error, start repair.
			break segmentLoop
		}
	}

	if readErr == nil {
		fmt.Println("no corruption in any segments")
		return
	}
	fmt.Printf("corruption error: %+v\n", readErr)

	logger := kitlog.NewLogfmtLogger(kitlog.NewSyncWriter(os.Stdout))
	logger = kitlog.With(logger, "caller", kitlog.DefaultCaller)
	wally, err := wal.NewSize(logger, nil, "bad-data", wal.DefaultSegmentSize)
	checkErr(err)
	checkErr(wally.Repair(readErr))
}

func checkErr(err error) {
	if err != nil {
		panic(err)
	}
}