package main

import (
	"errors"
	"fmt"
	"sort"

	helpers "git.bullercodeworks.com/brian/adventofcode/helpers"
)

func main() {
	inp := helpers.StdinToString()
	part1(inp)
	fmt.Println()
	part2(inp)
}

func part1(inp string) {
	fs := decompress(inp)
	fs.Fragment()
	fmt.Println("# Part 1")
	fmt.Println(fs.Checksum())
}

func part2(inp string) {
	fs := decompress(inp)
	fs.SortFilesDesc()
	fmt.Println("# Part 2")
	for i := range fs.Files {
		fmt.Print(fs.Files[i])
	}
	fmt.Println()
	fs.Defrag()
	fmt.Println(fs)
	fmt.Println(fs.Checksum())
}

type Filesystem struct {
	Files  []*FSFile
	Blocks []*FSBlock
}

func (fs *Filesystem) SortFilesDesc() {
	sort.Slice(fs.Files, func(i, j int) bool { return fs.Files[i].ID > fs.Files[j].ID })
}

func (fs *Filesystem) Checksum() int {
	var ret int
	for i := range fs.Blocks {
		if fs.Blocks[i].ID != -1 {
			ret = ret + (fs.Blocks[i].ID * i)
		}
	}
	return ret
}

func (fs *Filesystem) Defrag() bool {
	for i := range fs.Files {
		fl := fs.Files[i]
		loc := fs.FindFile(fl)
		mv := fs.FindFirstSpaceWithSize(fl.Size())
		if mv != -1 && mv < loc {
			fs.MoveFile(fl, mv)
		}
	}
	return false
}

func (fs *Filesystem) FindFile(fl *FSFile) int {
	if len(fl.Blocks) <= 0 {
		return -1
	}
	return fs.FindBlock(fl.Blocks[0])
}

func (fs *Filesystem) FindBlock(bl *FSBlock) int {
	for i := range fs.Blocks {
		if fs.Blocks[i] == bl {
			return i
		}
	}
	return -1
}

func (fs *Filesystem) MoveFile(fl *FSFile, to int) error {
	if fs.CalcSpaceAt(to) < fl.Size() {
		return errors.New("not enough space to move file")
	}
	for i := 0; i < fl.Size(); i++ {
		fs.MoveBlock(fl.Blocks[i], to+i)
	}
	return nil
}

func (fs *Filesystem) MoveBlock(bl *FSBlock, to int) {
	for i := range fs.Blocks {
		if fs.Blocks[i] == bl {
			// Swap i with to
			fs.SwapBlocks(i, to)
			return
		}
	}
}

func (fs *Filesystem) SwapBlocks(bl1, bl2 int) {
	fs.Blocks[bl1], fs.Blocks[bl2] = fs.Blocks[bl2], fs.Blocks[bl1]
}

func (fs *Filesystem) FindFirstSpaceWithSize(size int) int {
	for i := 0; i < len(fs.Blocks); i++ {
		if fs.CalcSpaceAt(i) >= size {
			return i
		}
	}
	return -1
}

func (fs *Filesystem) CalcFileSize(id int) int {
	startIdx, endIdx := -1, -1
	for i := range fs.Blocks {
		if fs.Blocks[i].ID == id && startIdx == -1 {
			// Found it
			startIdx = i
		} else if fs.Blocks[i].ID != id && startIdx != -1 {
			endIdx = i
			return endIdx - startIdx
		}
	}
	// Couldn't find the file
	return -1
}

func (fs *Filesystem) CalcSpaceAt(idx int) int {
	if fs.Blocks[idx].ID != -1 {
		return 0
	}
	var startIdx, endIdx int
	for i := idx; i >= 0; i-- {
		if fs.Blocks[i].ID == -1 {
			startIdx = i
		} else {
			break
		}
	}
	for i := idx; i < len(fs.Blocks); i++ {
		if fs.Blocks[i].ID == -1 {
			endIdx = i
		} else {
			break
		}
	}
	return (endIdx - startIdx) + 1
}

func (fs *Filesystem) Fragment() bool {
	// Find the last block
	lastBlock, firstSpace := fs.FindLastBlock(), fs.FindFirstSpace()
	if firstSpace < lastBlock {
		// Swap these two
		fs.SwapBlocks(firstSpace, lastBlock)
		// Recurse
		return fs.Fragment()
	} else {
		return true
	}
}

func (fs *Filesystem) FindLastBlock() int {
	for i := len(fs.Blocks) - 1; i >= 0; i-- {
		if fs.Blocks[i].ID != -1 {
			return i
		}
	}
	return -1
}

func (fs *Filesystem) FindFirstSpace() int {
	// Now find the first space
	for i := 0; i < len(fs.Blocks); i++ {
		if fs.Blocks[i].ID == -1 {
			return i
		}
	}
	return -1
}

func (fs *Filesystem) AddFile(id, size int) {
	fl := &FSFile{ID: id}
	for i := 0; i < size; i++ {
		c := &FSBlock{ID: id}
		fl.Blocks = append(fl.Blocks, c)
		fs.Blocks = append(fs.Blocks, c)
	}
	fs.Files = append(fs.Files, fl)
}

func (fs *Filesystem) AddSpace(size int) {
	for i := 0; i < size; i++ {
		c := &FSBlock{ID: -1}
		fs.Blocks = append(fs.Blocks, c)
	}
}

func (fs Filesystem) String() string {
	var ret string
	for i := range fs.Blocks {
		ret = fmt.Sprintf("%s%s", ret, fs.Blocks[i])
	}
	return ret
}

type FSFile struct {
	ID     int
	Blocks []*FSBlock
}

func (f *FSFile) Size() int {
	return len(f.Blocks)
}

func (f FSFile) String() string {
	return fmt.Sprintf("[%d;%d]", f.ID, f.Size())
}

type FSBlock struct {
	ID int
}

func (fsc FSBlock) String() string {
	if fsc.ID == -1 {
		return "[.]"
	}
	return fmt.Sprintf("[%d]", fsc.ID)
}

func decompress(inp string) Filesystem {
	ret := Filesystem{}
	latestId := 0
	file := true
	for i := range inp {
		size := int(inp[i] - '0')
		if file {
			ret.AddFile(latestId, size)
			latestId = latestId + 1
		} else {
			ret.AddSpace(size)
		}
		file = !file
	}
	return ret
}

func add(out []byte, v byte, count int) []byte {
	for i := 0; i < count; i++ {
		out = append(out, v)
	}
	return out
}