package squashfs

import (
	"bytes"
	"io"
)

type fragment struct {
	w     *Writer
	files []*fileHolder
	sizes []uint32
}

func (f *fragment) sizeLeft() uint32 {
	totalSize := uint32(0)
	for _, siz := range f.sizes {
		totalSize += siz
	}
	return f.w.BlockSize - uint32(totalSize)
}

func (f *fragment) addFragment(fil *fileHolder) {
	//SizeLeft should already be checked
	fil.fragOffset = len(f.files)
	f.files = append(f.files, fil)
	f.sizes = append(f.sizes, fil.blockSizes[len(fil.blockSizes)-1])
}

//TODO: give info about the frags for the frag table.
func (w *Writer) writeFragments(write io.WriterAt, off int64) (curOffset int64, err error) {
	curOffset = off
	var buf bytes.Buffer
	var byts []byte
	var n int
	for _, frag := range w.frags {
		blockOffset := 0
		for i, fil := range frag.files {
			_, err = io.CopyN(&buf, fil.reader, int64(frag.sizes[i]))
			if err != nil {
				return
			}
			fil.fragOffset = blockOffset
			blockOffset += int(frag.sizes[i])
		}
		if !w.Flags.UncompressedFragments && w.compressor != nil {
			byts, err = w.compressor.Compress(buf.Bytes())
			if err != nil {
				return
			}
		} else {
			byts = buf.Bytes()
		}
		n, err = write.WriteAt(byts, curOffset)
		curOffset += int64(n)
		if err != nil {
			return
		}
		buf.Reset()
	}
	return
}

func (w *Writer) addToFragments(fil *fileHolder) {
	fragSize := fil.blockSizes[len(fil.blockSizes)-1]
	//only fragment if the final block is less then 80% of a full block or AlwaysFragment
	if w.Flags.AlwaysFragment || fragSize < uint32(float32(w.BlockSize)*0.8) {
		var possibleFrags []int
		for i := range w.frags {
			left := w.frags[i].sizeLeft()
			if left == fragSize {
				fil.fragIndex = i
				w.frags[i].addFragment(fil)
				return
			} else if left > fragSize {
				possibleFrags = append(possibleFrags, i)
			}
		}
		if len(possibleFrags) > 0 {
			fil.fragIndex = possibleFrags[0]
		} else {
			fil.fragIndex = len(w.frags)
			w.frags = append(w.frags, fragment{
				w:     w,
				files: []*fileHolder{fil},
				sizes: []uint32{fragSize},
			})
		}
	}
}

func (w *Writer) calculateFragsAndBlockSizes() {
	for _, files := range w.structure {
		for i := range files {
			files[i].fragIndex = -1
			files[i].blockSizes = make([]uint32, files[i].size/uint64(w.BlockSize))
			for j := range files[i].blockSizes {
				files[i].blockSizes[j] = w.BlockSize
			}
			fragSize := uint32(files[i].size % uint64(w.BlockSize))
			if fragSize > 0 {
				files[i].blockSizes = append(files[i].blockSizes, fragSize)
				if !w.Flags.NoFragments {
					w.addToFragments(files[i])
				}
			}
		}
	}
}