package idxfile

import (
	"crypto"
	encbin "encoding/binary"
	"io"
	"sort"
	"sync"

	"github.com/go-git/go-git/v6/plumbing"
)

const (
	// VersionSupported is the only idx version supported.
	VersionSupported = 2

	noMapping = -1
)

var idxHeader = []byte{255, 't', 'O', 'c'}

// Index represents an index of a packfile.
type Index interface {
	// Contains checks whether the given hash is in the index.
	Contains(h plumbing.Hash) (bool, error)
	// FindOffset finds the offset in the packfile for the object with
	// the given hash.
	FindOffset(h plumbing.Hash) (int64, error)
	// FindCRC32 finds the CRC32 of the object with the given hash.
	FindCRC32(h plumbing.Hash) (uint32, error)
	// FindHash finds the hash for the object with the given offset.
	FindHash(o int64) (plumbing.Hash, error)
	// Count returns the number of entries in the index.
	Count() (int64, error)
	// Entries returns an iterator to retrieve all index entries.
	Entries() (EntryIter, error)
	// EntriesByOffset returns an iterator to retrieve all index entries ordered
	// by offset.
	EntriesByOffset() (EntryIter, error)
}

// MemoryIndex is the in memory representation of an idx file.
type MemoryIndex struct {
	// Version is the version of the index file.
	Version uint32
	// Fanout is a table where the Nth entry is the cumulative count of objects with the first byte of their name <= N.
	Fanout [256]uint32
	// FanoutMapping maps the position in the fanout table to the position
	// in the Names, Offset32 and CRC32 slices. This improves the memory
	// usage by not needing an array with unnecessary empty slots.
	FanoutMapping [256]int
	// Names is the list of object names.
	Names [][]byte
	// Offset32 is the list of 32-bit offsets.
	Offset32 [][]byte
	// CRC32 is the list of CRC32 checksums.
	CRC32 [][]byte
	// Offset64 is the list of 64-bit offsets.
	Offset64 []byte
	// PackfileChecksum is the checksum of the packfile.
	PackfileChecksum plumbing.Hash
	// IdxChecksum is the checksum of the index file.
	IdxChecksum plumbing.Hash

	offsetHash      map[int64]plumbing.Hash
	offsetBuildOnce sync.Once
	mu              sync.RWMutex

	objectIDSize int
}

var _ Index = (*MemoryIndex)(nil)

// NewMemoryIndex returns an instance of a new MemoryIndex.
func NewMemoryIndex(objectIDSize int) *MemoryIndex {
	m := &MemoryIndex{objectIDSize: objectIDSize}
	m.IdxChecksum.ResetBySize(objectIDSize)
	m.PackfileChecksum.ResetBySize(objectIDSize)
	return m
}

func (idx *MemoryIndex) findHashIndex(h plumbing.Hash) (int, bool) {
	k := idx.FanoutMapping[h.Bytes()[0]]
	if k == noMapping {
		return 0, false
	}

	if len(idx.Names) <= k {
		return 0, false
	}

	data := idx.Names[k]
	high := uint64(len(idx.Offset32[k])) >> 2
	if high == 0 {
		return 0, false
	}

	low := uint64(0)
	for {
		mid := (low + high) >> 1
		offset := mid * uint64(idx.idSize())

		cmp := h.Compare(data[offset : offset+uint64(idx.idSize())])
		switch {
		case cmp < 0:
			high = mid
		case cmp == 0:
			return int(mid), true
		default:
			low = mid + 1
		}

		if low >= high {
			break
		}
	}

	return 0, false
}

// Contains implements the Index interface.
func (idx *MemoryIndex) Contains(h plumbing.Hash) (bool, error) {
	_, ok := idx.findHashIndex(h)
	return ok, nil
}

// FindOffset implements the Index interface.
func (idx *MemoryIndex) FindOffset(h plumbing.Hash) (int64, error) {
	fo := h.Bytes()[0]
	if len(idx.FanoutMapping) <= int(fo) {
		return 0, plumbing.ErrObjectNotFound
	}

	k := idx.FanoutMapping[fo]
	i, ok := idx.findHashIndex(h)
	if !ok {
		return 0, plumbing.ErrObjectNotFound
	}

	offset := idx.getOffset(k, i)

	// Save the offset for reverse lookup
	idx.mu.Lock()
	if idx.offsetHash == nil {
		idx.offsetHash = make(map[int64]plumbing.Hash)
	}
	idx.offsetHash[int64(offset)] = h
	idx.mu.Unlock()

	return int64(offset), nil
}

const isO64Mask = uint64(1) << 31

func (idx *MemoryIndex) getOffset(firstLevel, secondLevel int) uint64 {
	offset := secondLevel << 2
	ofs := encbin.BigEndian.Uint32(idx.Offset32[firstLevel][offset : offset+4])

	if (uint64(ofs) & isO64Mask) != 0 {
		offset := 8 * (uint64(ofs) & ^isO64Mask)
		n := encbin.BigEndian.Uint64(idx.Offset64[offset : offset+8])
		return n
	}

	return uint64(ofs)
}

// FindCRC32 implements the Index interface.
func (idx *MemoryIndex) FindCRC32(h plumbing.Hash) (uint32, error) {
	k := idx.FanoutMapping[h.Bytes()[0]]
	i, ok := idx.findHashIndex(h)
	if !ok {
		return 0, plumbing.ErrObjectNotFound
	}

	return idx.getCRC32(k, i), nil
}

func (idx *MemoryIndex) getCRC32(firstLevel, secondLevel int) uint32 {
	offset := secondLevel << 2
	return encbin.BigEndian.Uint32(idx.CRC32[firstLevel][offset : offset+4])
}

// FindHash implements the Index interface.
func (idx *MemoryIndex) FindHash(o int64) (plumbing.Hash, error) {
	var hash plumbing.Hash
	var ok bool

	idx.mu.RLock()
	if idx.offsetHash != nil {
		if hash, ok = idx.offsetHash[o]; ok {
			idx.mu.RUnlock()
			return hash, nil
		}
	}
	idx.mu.RUnlock()

	var genErr error
	idx.offsetBuildOnce.Do(func() {
		genErr = idx.genOffsetHash()
	})
	if genErr != nil {
		return plumbing.ZeroHash, genErr
	}

	idx.mu.RLock()
	hash, ok = idx.offsetHash[o]
	idx.mu.RUnlock()

	if !ok {
		return plumbing.ZeroHash, plumbing.ErrObjectNotFound
	}

	return hash, nil
}

// genOffsetHash generates the offset/hash mapping for reverse search.
func (idx *MemoryIndex) genOffsetHash() error {
	count, err := idx.Count()
	if err != nil {
		return err
	}

	offsetHash := make(map[int64]plumbing.Hash, count)

	var hash plumbing.Hash
	hash.ResetBySize(idx.objectIDSize)

	i := uint32(0)
	for firstLevel, fanoutValue := range idx.Fanout {
		mappedFirstLevel := idx.FanoutMapping[firstLevel]
		for secondLevel := uint32(0); i < fanoutValue; i++ {
			_, _ = hash.Write(idx.Names[mappedFirstLevel][secondLevel*uint32(idx.idSize()):])
			offset := int64(idx.getOffset(mappedFirstLevel, int(secondLevel)))
			offsetHash[offset] = hash
			secondLevel++
		}
	}

	idx.mu.Lock()
	idx.offsetHash = offsetHash
	idx.mu.Unlock()

	return nil
}

// Count implements the Index interface.
func (idx *MemoryIndex) Count() (int64, error) {
	return int64(idx.Fanout[fanout-1]), nil
}

// Entries implements the Index interface.
func (idx *MemoryIndex) Entries() (EntryIter, error) {
	return &idxfileEntryIter{idx, 0, 0, 0}, nil
}

// EntriesByOffset implements the Index interface.
func (idx *MemoryIndex) EntriesByOffset() (EntryIter, error) {
	count, err := idx.Count()
	if err != nil {
		return nil, err
	}

	iter := &idxfileEntryOffsetIter{
		entries: make(entriesByOffset, count),
	}

	entries, err := idx.Entries()
	if err != nil {
		return nil, err
	}

	for pos := 0; int64(pos) < count; pos++ {
		entry, err := entries.Next()
		if err != nil {
			return nil, err
		}

		iter.entries[pos] = entry
	}

	sort.Sort(iter.entries)

	return iter, nil
}

func (idx *MemoryIndex) idSize() int {
	if idx.objectIDSize != 0 {
		return idx.objectIDSize
	}
	return crypto.SHA1.Size()
}

// EntryIter is an iterator that will return the entries in a packfile index.
type EntryIter interface {
	// Next returns the next entry in the packfile index.
	Next() (*Entry, error)
	// Close closes the iterator.
	Close() error
}

type idxfileEntryIter struct {
	idx                     *MemoryIndex
	total                   int
	firstLevel, secondLevel int
}

func (i *idxfileEntryIter) Next() (*Entry, error) {
	for {
		if i.firstLevel >= fanout {
			return nil, io.EOF
		}

		if i.total >= int(i.idx.Fanout[i.firstLevel]) {
			i.firstLevel++
			i.secondLevel = 0
			continue
		}

		mappedFirstLevel := i.idx.FanoutMapping[i.firstLevel]
		entry := new(Entry)
		entry.Hash.ResetBySize(i.idx.idSize())
		_, _ = entry.Hash.Write(i.idx.Names[mappedFirstLevel][i.secondLevel*i.idx.idSize():])
		entry.Offset = i.idx.getOffset(mappedFirstLevel, i.secondLevel)
		entry.CRC32 = i.idx.getCRC32(mappedFirstLevel, i.secondLevel)

		i.secondLevel++
		i.total++

		return entry, nil
	}
}

func (i *idxfileEntryIter) Close() error {
	i.firstLevel = fanout
	return nil
}

// Entry is the in memory representation of an object entry in the idx file.
type Entry struct {
	Hash   plumbing.Hash
	CRC32  uint32
	Offset uint64
}

type idxfileEntryOffsetIter struct {
	entries entriesByOffset
	pos     int
}

func (i *idxfileEntryOffsetIter) Next() (*Entry, error) {
	if i.pos >= len(i.entries) {
		return nil, io.EOF
	}

	entry := i.entries[i.pos]
	i.pos++

	return entry, nil
}

func (i *idxfileEntryOffsetIter) Close() error {
	i.pos = len(i.entries) + 1
	return nil
}

type entriesByOffset []*Entry

func (o entriesByOffset) Len() int {
	return len(o)
}

func (o entriesByOffset) Less(i, j int) bool {
	return o[i].Offset < o[j].Offset
}

func (o entriesByOffset) Swap(i, j int) {
	o[i], o[j] = o[j], o[i]
}
