// License: GPLv3 Copyright: 2023, Kovid Goyal, <kovid at kovidgoyal.net>

// First create a patcher with:
// p = NewPatcher()
// Create a signature for the file you want to update using
// p.CreateSignatureIterator(file_to_update)
// Now create a Differ with the created signature
// d = NewDiffer()
// d.AddSignatureData(signature_data_from_previous_step)
// Now create a delta based on the signature and the reference file
// d.CreateDelta(reference_file)
// Finally, apply this delta using the patcher to produce a file identical to reference_file
// based ont he delta data and file_to_update
// p.StartDelta(output_file, file_to_update)
// p.UpdateDelta(...)
// p.FinishDelta()
package rsync

import (
	"fmt"
	"io"
	"math"

	"kitty/tools/utils"
)

var _ = fmt.Print

const MaxBlockSize int = 1024 * 1024 // sqrt of 1TB

type StrongHashType uint16
type WeakHashType uint16
type ChecksumType uint16

const (
	XXH3 StrongHashType = iota
)
const (
	XXH3128Sum ChecksumType = iota
)
const (
	Rsync WeakHashType = iota
)

type GrowBufferFunction = func(slice []byte, sz int) []byte

type Api struct {
	rsync     rsync
	signature []BlockHash

	Checksum_type    ChecksumType
	Strong_hash_type StrongHashType
	Weak_hash_type   WeakHashType
	Grow_buffer_func GrowBufferFunction
}

type Differ struct {
	Api
	unconsumed_signature_data []byte
}

type Patcher struct {
	Api
	unconsumed_delta_data                        []byte
	expected_input_size_for_signature_generation int64
	delta_output                                 io.Writer
	delta_input                                  io.ReadSeeker
	total_data_in_delta                          int
}

// internal implementation {{{
func (self *Api) read_signature_header(data []byte) (consumed int, err error) {
	if len(data) < 12 {
		return -1, io.ErrShortBuffer
	}
	if version := bin.Uint16(data); version != 0 {
		return consumed, fmt.Errorf("Invalid version in signature header: %d", version)
	}
	switch csum := ChecksumType(bin.Uint16(data[2:])); csum {
	case XXH3128Sum:
		self.Checksum_type = XXH3128Sum
		self.rsync.SetChecksummer(new_xxh3_128)
	default:
		return consumed, fmt.Errorf("Invalid checksum_type in signature header: %d", csum)
	}
	switch strong_hash := StrongHashType(bin.Uint16(data[4:])); strong_hash {
	case XXH3:
		self.Strong_hash_type = strong_hash
		self.rsync.SetHasher(new_xxh3_64)
	default:
		return consumed, fmt.Errorf("Invalid strong_hash in signature header: %d", strong_hash)
	}
	switch weak_hash := WeakHashType(bin.Uint16(data[6:])); weak_hash {
	case Rsync:
		self.Weak_hash_type = weak_hash
	default:
		return consumed, fmt.Errorf("Invalid weak_hash in signature header: %d", weak_hash)
	}
	block_size := int(bin.Uint32(data[8:]))
	consumed = 12
	if block_size == 0 {
		return consumed, fmt.Errorf("rsync signature header has zero block size")
	}
	if block_size > MaxBlockSize {
		return consumed, fmt.Errorf("rsync signature header has too large block size %d > %d", block_size, MaxBlockSize)
	}
	self.rsync.BlockSize = block_size
	self.signature = make([]BlockHash, 0, 1024)
	return
}

func (self *Api) read_signature_blocks(data []byte) (consumed int) {
	block_hash_size := self.rsync.HashSize() + 12
	for ; len(data) >= block_hash_size; data = data[block_hash_size:] {
		bl := BlockHash{}
		bl.Unserialize(data[:block_hash_size])
		self.signature = append(self.signature, bl)
		consumed += block_hash_size
	}
	return
}

func (self *Differ) finish_signature_data() (err error) {
	if len(self.unconsumed_signature_data) > 0 {
		return fmt.Errorf("There were %d leftover bytes in the signature data", len(self.unconsumed_signature_data))
	}
	self.unconsumed_signature_data = nil
	if !self.rsync.HasHasher() {
		return fmt.Errorf("No header was found in the signature data")
	}
	return
}

func (self *Patcher) update_delta(data []byte) (consumed int, err error) {
	op := Operation{}
	for len(data) > 0 {
		n, uerr := op.Unserialize(data)
		if uerr == nil {
			consumed += n
			data = data[n:]
			if err = self.rsync.ApplyDelta(self.delta_output, self.delta_input, op); err != nil {
				return
			}
			if op.Type == OpData {
				self.total_data_in_delta += len(op.Data)
			}
		} else {
			if n < 0 {
				return consumed, nil
			}
			return consumed, uerr
		}
	}
	return
}

// }}}

// Start applying serialized delta
func (self *Patcher) StartDelta(delta_output io.Writer, delta_input io.ReadSeeker) {
	self.delta_output = delta_output
	self.delta_input = delta_input
	self.total_data_in_delta = 0
	self.unconsumed_delta_data = nil
}

// Apply a chunk of delta data
func (self *Patcher) UpdateDelta(data []byte) (err error) {
	self.unconsumed_delta_data = append(self.unconsumed_delta_data, data...)
	consumed, err := self.update_delta(self.unconsumed_delta_data)
	if err != nil {
		return err
	}
	self.unconsumed_delta_data = utils.ShiftLeft(self.unconsumed_delta_data, consumed)
	return
}

// Finish applying delta data
func (self *Patcher) FinishDelta() (err error) {
	if err = self.UpdateDelta([]byte{}); err != nil {
		return err
	}
	if len(self.unconsumed_delta_data) > 0 {
		return fmt.Errorf("There are %d leftover bytes in the delta", len(self.unconsumed_delta_data))
	}
	self.delta_input = nil
	self.delta_output = nil
	self.unconsumed_delta_data = nil
	return
}

func double_size(output []byte, count int) []byte {
	ncap := utils.Max(8192, cap(output)*2, cap(output)+count)
	t := make([]byte, len(output), ncap)
	copy(t, output)
	return t
}

func ensure_size(output []byte, count int, esf GrowBufferFunction) ([]byte, []byte) {
	if cap(output)-len(output) < count {
		output = esf(output, count)
	}
	return output[:len(output)+count], output[len(output) : len(output)+count]
}

func write_block_hash(output []byte, bl BlockHash, f GrowBufferFunction) []byte {
	output, b := ensure_size(output, BlockHashSize, f)
	bl.Serialize(b)
	return output
}

// Append the next item to the provided slice and return the resulting slice. Data is written to the slice iff err == nil.
// When no more data is available, err = io.EOF and the unmodified slice is returned.
type OutputIterator = func([]byte) ([]byte, error)

// Create a signature for the data source in src.
func (self *Patcher) CreateSignatureIterator(src io.Reader) OutputIterator {
	var it func() (BlockHash, error)
	finished := false
	gbf := self.Grow_buffer_func
	if gbf == nil {
		gbf = double_size
	}
	return func(output []byte) ([]byte, error) {
		if finished {
			return output, io.EOF
		}
		if it == nil { // write signature header
			var b []byte
			it = self.rsync.CreateSignatureIterator(src)
			output, b = ensure_size(output, 12, gbf)
			bin.PutUint16(b, 0)
			bin.PutUint16(b[2:], uint16(self.Checksum_type))
			bin.PutUint16(b[4:], uint16(self.Strong_hash_type))
			bin.PutUint16(b[6:], uint16(self.Weak_hash_type))
			bin.PutUint32(b[8:], uint32(self.rsync.BlockSize))
		}
		bl, err := it()
		switch err {
		case io.EOF:
			finished = true
			return output, io.EOF
		case nil:
			return write_block_hash(output, bl, gbf), nil
		default:
			return nil, err
		}
	}
}

// Create a serialized delta based on the previously loaded signature
func (self *Differ) CreateDelta(src io.Reader, output io.Writer) func() error {
	if err := self.finish_signature_data(); err != nil {
		return func() error { return err }
	}
	if self.signature == nil {
		return func() error {
			return fmt.Errorf("Cannot call CreateDelta() before loading a signature")
		}
	}
	return self.rsync.CreateDiff(src, self.signature, output)
}

// Add more external signature data
func (self *Differ) AddSignatureData(data []byte) (err error) {
	self.unconsumed_signature_data = append(self.unconsumed_signature_data, data...)
	if !self.rsync.HasHasher() {
		consumed, err := self.read_signature_header(self.unconsumed_signature_data)
		if err != nil {
			if consumed < 0 {
				return nil
			}
			return err
		}
		self.unconsumed_signature_data = utils.ShiftLeft(self.unconsumed_signature_data, consumed)
	}
	consumed := self.read_signature_blocks(self.unconsumed_signature_data)
	self.unconsumed_signature_data = utils.ShiftLeft(self.unconsumed_signature_data, consumed)
	return nil
}

// Use to calculate a delta based on a supplied signature, via AddSignatureData
func NewDiffer() *Differ {
	return &Differ{}
}

// Use to create a signature and possibly apply a delta
func NewPatcher(expected_input_size int64) (ans *Patcher) {
	bs := DefaultBlockSize
	sz := utils.Max(0, expected_input_size)
	if sz > 0 {
		bs = int(math.Round(math.Sqrt(float64(sz))))
	}
	ans = &Patcher{}
	ans.rsync.BlockSize = utils.Min(bs, MaxBlockSize)
	ans.rsync.SetHasher(new_xxh3_64)
	ans.rsync.SetChecksummer(new_xxh3_128)

	if ans.rsync.HashBlockSize() > 0 && ans.rsync.HashBlockSize() < ans.rsync.BlockSize {
		ans.rsync.BlockSize = (ans.rsync.BlockSize / ans.rsync.HashBlockSize()) * ans.rsync.HashBlockSize()
	}

	ans.expected_input_size_for_signature_generation = sz
	return
}