/*
glocate -- ZFS-diff-friendly locate-like utility
Copyright (C) 2022 Sergey Matveev <stargrave@stargrave.org>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, version 3 of the License.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

package main

import (
	"bufio"
	"encoding/gob"
	"flag"
	"fmt"
	"io"
	"io/fs"
	"log"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/klauspost/compress/zstd"
)

type File struct {
	Name  string
	Size  uint64
	Mtime int64
	Files []File
}

type ByName []File

func (a ByName) Len() int {
	return len(a)
}

func (a ByName) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

func (a ByName) Less(i, j int) bool {
	return a[i].Name < a[j].Name
}

func (file *File) IsDir() bool {
	return file.Name[len(file.Name)-1] == '/'
}

func walk(root string) ([]File, uint64, error) {
	fd, err := os.Open(root)
	if err != nil {
		return nil, 0, err
	}
	var files []File
	var size uint64
	var info fs.FileInfo
	for {
		ents, err := fd.ReadDir(1 << 10)
		if err != nil {
			if err == io.EOF {
				break
			}
			fd.Close()
			return nil, 0, err
		}
		for _, ent := range ents {
			file := File{Name: ent.Name()}
			fullPath := path.Join(root, file.Name)
			if ent.IsDir() {
				file.Name += "/"
			}
			info, err = ent.Info()
			if err != nil {
				log.Println("can not stat:", fullPath, ":", err)
				files = append(files, file)
				continue
			}
			file.Mtime = info.ModTime().Unix()
			if ent.IsDir() {
				file.Files, file.Size, err = walk(fullPath)
				if err != nil {
					log.Println("can not walk:", fullPath, ":", err)
					files = append(files, file)
					continue
				}
			} else if info.Mode().IsRegular() {
				file.Size = uint64(info.Size())
			}
			files = append(files, file)
			size += file.Size
		}
	}
	fd.Close()
	sort.Sort(ByName(files))
	return files, size, nil
}

func usage() {
	log.Println("usage")
	os.Exit(1)
}

func load(dbPath string) *File {
	fd, err := os.Open(dbPath)
	if err != nil {
		log.Fatalln(err)
	}
	defer fd.Close()
	comp, err := zstd.NewReader(fd)
	if err != nil {
		log.Fatalln(err)
	}
	dec := gob.NewDecoder(comp)
	var file File
	err = dec.Decode(&file)
	if err != nil {
		log.Fatalln(err)
	}
	comp.Close()
	return &file
}

func (db *File) dump(dbPath string) error {
	tmp, err := os.CreateTemp(path.Dir(dbPath), "glocate")
	if err != nil {
		return err
	}
	defer os.Remove(tmp.Name())
	comp, err := zstd.NewWriter(
		tmp, zstd.WithEncoderLevel(zstd.SpeedBestCompression),
	)
	if err != nil {
		return err
	}
	enc := gob.NewEncoder(comp)
	err = enc.Encode(db)
	if err != nil {
		return err
	}
	err = comp.Close()
	if err != nil {
		return err
	}
	err = tmp.Close()
	if err != nil {
		return err
	}
	umask := syscall.Umask(0)
	syscall.Umask(umask)
	err = os.Chmod(tmp.Name(), os.FileMode(0666&^umask))
	if err != nil {
		return err
	}
	return os.Rename(tmp.Name(), dbPath)
}

func (file *File) listBeauty(indent string, n int, isLast, veryFirst bool) {
	if veryFirst {
		fmt.Printf("[%s]\n", humanize.IBytes(file.Size))
	} else {
		var box string
		if isLast {
			box = "└"
		} else {
			box = "├"
		}
		name := file.Name
		fmt.Printf("%s%s %s\t№%d [%s] %s\n",
			indent, box, name, n, humanize.IBytes(file.Size),
			time.Unix(file.Mtime, 0).Format("2006-01-02"),
		)
		if isLast {
			indent += "  "
		} else {
			indent += "│ "
		}
	}
	for n, f := range file.Files {
		n++
		f.listBeauty(indent, n, n == len(file.Files), false)
	}
}

func (file *File) listSimple(root string, veryFirst bool) {
	name := file.Name
	fmt.Println(
		strconv.FormatUint(file.Size, 10),
		time.Unix(file.Mtime, 0).Format("2006-01-02T15:04:05"),
		root+name,
	)
	if veryFirst {
		name = ""
	}
	for _, f := range file.Files {
		f.listSimple(root+name, false)
	}
}

func (file *File) listFiles(root string, veryFirst bool) {
	name := file.Name
	if veryFirst {
		root = ""
	} else {
		fmt.Println(root + name)
		root += name
	}
	for _, f := range file.Files {
		f.listFiles(root, false)
	}
}

func (db *File) find(p string) (file *File, parents []*File, idx int, err error) {
	file = db
	var f File
Entities:
	for _, ent := range strings.Split(p, "/") {
		for idx, f = range file.Files {
			if (ent == f.Name) || (ent+"/" == f.Name) {
				parents = append(parents, file)
				file = &f
				continue Entities
			}
		}
		err = fmt.Errorf("no entity found: %s", ent)
		return
	}
	return
}

func (db *File) remove(p string) error {
	file, parents, idx, err := db.find(p)
	if err != nil {
		return err
	}
	lastParent := parents[len(parents)-1]
	lastParent.Files = append(
		lastParent.Files[:idx],
		lastParent.Files[idx+1:]...,
	)
	for _, parent := range parents {
		parent.Size -= file.Size
	}
	return nil
}

func (db *File) add(p string) error {
	cols := strings.Split(p, "/")
	cols, name := cols[:len(cols)-1], cols[len(cols)-1]
	var parent *File
	var err error
	if len(cols) != 0 {
		parent, _, _, err = db.find(path.Join(cols...))
		if err != nil {
			return err
		}
	} else {
		parent = db
	}
	info, err := os.Stat(p)
	if err != nil {
		return err
	}
	if info.IsDir() {
		name += "/"
	}
	file := File{
		Name:  name,
		Size:  uint64(info.Size()),
		Mtime: info.ModTime().Unix(),
	}
	parent.Files = append(parent.Files, file)
	sort.Sort(ByName(parent.Files))
	parent.Size += file.Size
	return nil
}

func deoctalize(s string) string {
	chars := make([]byte, 0, len(s))
	for i := 0; i < len(s); i++ {
		if s[i] == '\\' {
			b, err := strconv.ParseUint("0"+s[i+1:i+1+3], 0, 8)
			if err != nil {
				log.Fatalln(err)
			}
			chars = append(chars, byte(b))
			i += 3
		} else {
			chars = append(chars, s[i])
		}
	}
	return string(chars)
}

func main() {
	dbPath := flag.String("db", ".glocate.db", "Path to state file (database)")
	doIndex := flag.Bool("index", false, "Initialize database")
	doUpdate := flag.Bool("update", false, "Update database by zfs-diff's output")
	showBeauty := flag.Bool("show-beauty", false, "Show beauty human-friendly listing")
	showSimple := flag.Bool("show-simple", false, "Show simple listing")
	stripPrefix := flag.String("strip-prefix", "", "Strip prefix from zfs-diff's output")
	flag.Parse()
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds | log.Lshortfile)

	if *doIndex {
		files, size, err := walk(".")
		if err != nil {
			log.Fatalln(err)
		}
		db := File{Name: "./", Size: size, Files: files}
		err = db.dump(*dbPath)
		if err != nil {
			log.Fatalln(err)
		}
		return
	}

	db := load(*dbPath)
	if *doUpdate {
		scanner := bufio.NewScanner(os.Stdin)
		var t string
		for scanner.Scan() {
			t = scanner.Text()
			if len(t) == 0 {
				continue
			}
			cols := strings.Split(t, "\t")
			if len(cols) < 2 {
				log.Fatalln("bad zfs-diff format")
			}
			switch cols[0] {
			case "-":
				name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
				if err := db.remove(name); err != nil {
					log.Println("can not -:", name, ":", err)
				}
			case "+":
				name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
				if err := db.add(name); err != nil {
					log.Println("can not +:", name, ":", err)
				}
			case "M":
				name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
				if name == "" {
					continue
				}
				file, _, _, err := db.find(name)
				if err != nil {
					log.Println("can not M:", name, ":", err)
					continue
				}
				info, err := os.Stat(name)
				if err != nil {
					log.Println("can not M:", name, ":", err)
					continue
				}
				if info.Mode().IsRegular() {
					file.Size = uint64(info.Size())
				}
				file.Mtime = info.ModTime().Unix()
			case "R":
				if len(cols) != 3 {
					log.Fatalln("bad zfs-diff format for R")
				}
				name := deoctalize(strings.TrimPrefix(cols[1], *stripPrefix))
				if err := db.remove(name); err != nil {
					log.Println("can not R-:", name, ":", err)
					continue
				}
				name = deoctalize(strings.TrimPrefix(cols[2], *stripPrefix))
				if err := db.add(name); err != nil {
					log.Println("can not R+:", name, ":", err)
				}
			default:
				log.Fatalln("bad zfs-diff format")
			}
		}
		if err := scanner.Err(); err != nil {
			log.Fatalln(err)
		}
		if err := db.dump(*dbPath); err != nil {
			log.Fatalln(err)
		}
		return
	}

	veryFirst := true
	if len(flag.Args()) > 0 {
		root := flag.Args()[0]
		if root[:2] == "./" {
			root = root[2:]
		}
		if root[len(root)-1:] == "/" {
			root = root[:len(root)-1]
		}
		file, _, _, err := db.find(root)
		if err != nil {
			log.Fatalln(err)
		}
		db = file
		db.Name = root + "/"
		veryFirst = false
	}

	if *showBeauty {
		db.listBeauty("", 0, false, veryFirst)
		return
	}
	if *showSimple {
		db.listSimple("", veryFirst)
		return
	}
	db.listFiles("", veryFirst)
}