package main

import (
	"bufio"
	"encoding/binary"
	"io"
	"log"
	"os"

	"github.com/klauspost/compress/zstd"
)

func dirSizer(dirSizes *[]int64, depth int, sinkBack, sinkIn, sinkOut chan Ent) (curSize int64) {
	var ent Ent
	var opened bool
	var dirIdx int
	for {
		select {
		case ent = <-sinkBack:
			goto Got
		default:
		}
		ent, opened = <-sinkIn
		if !opened {
			break
		}
	Got:
		if len(ent.name) < depth {
			sinkBack <- ent
			return
		}
		sinkOut <- ent
		curSize += ent.size
		if !ent.IsDir() {
			continue
		}
		dirIdx = len(*dirSizes)
		(*dirSizes) = append(*dirSizes, 0)
		dirSize := dirSizer(dirSizes, depth+1, sinkBack, sinkIn, sinkOut)
		(*dirSizes)[dirIdx] = dirSize
		curSize += dirSize
	}
	return
}

func applyDirSizes(src *os.File, dirSizes []int64) *os.File {
	_, err := src.Seek(0, io.SeekStart)
	if err != nil {
		log.Fatalln(err)
	}
	tmp, err := os.CreateTemp(TmpDir, "glocate-idx")
	if err != nil {
		log.Fatalln(err)
	}

	compR, err := zstd.NewReader(src)
	if err != nil {
		log.Fatalln(err)
	}
	br := bufio.NewReaderSize(compR, 1<<17)

	compW, err := zstd.NewWriter(tmp,
		zstd.WithEncoderLevel(zstd.SpeedBestCompression))
	if err != nil {
		log.Fatalln(err)
	}
	bw := bufio.NewWriterSize(compW, 1<<17)

	num := make([]byte, 8)
	var nameLen int
	name := make([]byte, 0, 1<<16)
	var dirIdx int
	for {
		if _, err = io.ReadFull(br, num[:2]); err != nil {
			if err == io.EOF {
				break
			}
			log.Fatalln(err)
		}
		mustWrite(bw, num[:2])
		nameLen = int(binary.BigEndian.Uint16(num[:2]))
		name = name[:nameLen]
		mustReadFull(br, name)
		mustWrite(bw, name)
		if _, err = io.CopyN(bw, br, 1+8); err != nil {
			log.Fatalln(err)
		}
		if name[len(name)-1] == byte('/') {
			if _, err = br.Discard(8); err != nil {
				log.Fatalln(err)
			}
			binary.BigEndian.PutUint64(num, uint64(dirSizes[dirIdx]))
			mustWrite(bw, num)
			dirIdx++
		} else {
			if _, err = io.CopyN(bw, br, 8); err != nil {
				log.Fatalln(err)
			}
		}
	}
	if err = bw.Flush(); err != nil {
		log.Fatalln(err)
	}
	if err = compW.Close(); err != nil {
		log.Fatalln(err)
	}
	compR.Close()
	return tmp
}