package replica

import (
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/go-git/go-git/v5"
	"github.com/go-git/go-git/v5/plumbing"
	"github.com/go-git/go-git/v5/plumbing/object"
)

var ErrPruneKeepMissing = errors.New("replica: an object the prune must keep is not in the local repository")

var pruneHook = func() error { return nil }

func Prune(dir string, keep ...plumbing.Hash) (*git.Repository, error) {
	r, err := git.PlainOpen(dir)
	if err != nil {
		return nil, err
	}
	head, err := localHead(r)
	if err != nil {
		return nil, err
	}
	roots := append([]plumbing.Hash{}, keep...)
	if !head.IsZero() {
		roots = append(roots, head)
	}
	critical := map[plumbing.Hash]bool{}
	for _, id := range roots {
		if err := addFull(r, id, critical); err != nil {
			return nil, err
		}
	}
	set := map[plumbing.Hash]bool{}
	for h := range critical {
		set[h] = true
	}
	addCommitChain(r, roots, set)
	objects := filepath.Join(dir, "objects")
	for h := range set {
		p := loosePath(objects, h)
		written := false
		if _, err := os.Stat(p); errors.Is(err, fs.ErrNotExist) {
			obj, err := r.Storer.EncodedObject(plumbing.AnyObject, h)
			if err != nil {
				return nil, err
			}
			if _, err := r.Storer.SetEncodedObject(obj); err != nil {
				return nil, err
			}
			written = true
		} else if err != nil {
			return nil, err
		}
		if !written && !critical[h] {
			continue
		}
		if err := syncPath(p); err != nil {
			return nil, err
		}
		if err := syncPath(filepath.Dir(p)); err != nil {
			return nil, err
		}
	}
	if err := syncPath(objects); err != nil {
		return nil, err
	}
	if err := removeOtherRefs(r); err != nil {
		return nil, err
	}
	if err := removePacks(filepath.Join(objects, "pack")); err != nil {
		return nil, err
	}
	if err := removeLoose(objects, set); err != nil {
		return nil, err
	}
	return git.PlainOpen(dir)
}

func addFull(r *git.Repository, id plumbing.Hash, set map[plumbing.Hash]bool) error {
	c, err := object.GetCommit(r.Storer, id)
	if err != nil {
		return fmt.Errorf("%w: commit %s: %v", ErrPruneKeepMissing, id, err)
	}
	tree, err := c.Tree()
	if err != nil {
		return fmt.Errorf("%w: tree of %s: %v", ErrPruneKeepMissing, id, err)
	}
	for _, e := range tree.Entries {
		if _, err := r.Storer.EncodedObject(plumbing.AnyObject, e.Hash); err != nil {
			return fmt.Errorf("%w: blob of %s: %v", ErrPruneKeepMissing, id, err)
		}
		set[e.Hash] = true
	}
	set[id] = true
	set[tree.Hash] = true
	return nil
}

func addCommitChain(r *git.Repository, roots []plumbing.Hash, set map[plumbing.Hash]bool) {
	seen := map[plumbing.Hash]bool{}
	stack := append([]plumbing.Hash{}, roots...)
	for len(stack) > 0 {
		id := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		if seen[id] {
			continue
		}
		seen[id] = true
		c, err := object.GetCommit(r.Storer, id)
		if err != nil {
			continue
		}
		set[id] = true
		stack = append(stack, c.ParentHashes...)
	}
}

func loosePath(objects string, h plumbing.Hash) string {
	s := h.String()
	return filepath.Join(objects, s[:2], s[2:])
}

func syncPath(p string) error {
	f, err := os.Open(p)
	if err != nil {
		return err
	}
	defer f.Close()
	return f.Sync()
}

func removeOtherRefs(r *git.Repository) error {
	iter, err := r.Storer.IterReferences()
	if err != nil {
		return err
	}
	var names []plumbing.ReferenceName
	if err := iter.ForEach(func(ref *plumbing.Reference) error {
		if ref.Name() != Branch && ref.Name() != plumbing.HEAD {
			names = append(names, ref.Name())
		}
		return nil
	}); err != nil {
		return err
	}
	for _, n := range names {
		if err := r.Storer.RemoveReference(n); err != nil {
			return err
		}
	}
	return nil
}

func removePacks(dir string) error {
	entries, err := os.ReadDir(dir)
	if errors.Is(err, fs.ErrNotExist) {
		return nil
	}
	if err != nil {
		return err
	}
	var packs, rest []string
	for _, e := range entries {
		if strings.HasSuffix(e.Name(), ".pack") {
			packs = append(packs, e.Name())
		} else {
			rest = append(rest, e.Name())
		}
	}
	sort.Strings(packs)
	sort.Strings(rest)
	for _, name := range append(packs, rest...) {
		if err := pruneHook(); err != nil {
			return err
		}
		if err := os.Remove(filepath.Join(dir, name)); err != nil {
			return err
		}
	}
	return nil
}

func removeLoose(objects string, keep map[plumbing.Hash]bool) error {
	dirs, err := os.ReadDir(objects)
	if err != nil {
		return err
	}
	for _, d := range dirs {
		if !d.IsDir() || len(d.Name()) != 2 {
			continue
		}
		sub := filepath.Join(objects, d.Name())
		files, err := os.ReadDir(sub)
		if err != nil {
			return err
		}
		left := 0
		for _, f := range files {
			h := plumbing.NewHash(d.Name() + f.Name())
			if len(d.Name()+f.Name()) != 40 || h.IsZero() || keep[h] {
				left++
				continue
			}
			if err := pruneHook(); err != nil {
				return err
			}
			if err := os.Remove(filepath.Join(sub, f.Name())); err != nil {
				return err
			}
		}
		if left == 0 {
			_ = os.Remove(sub)
		}
	}
	return nil
}