package store

import (
	"encoding/binary"

	"github.com/cockroachdb/pebble"
	"github.com/pkg/errors"
	"go.uber.org/zap"
	"golang.org/x/crypto/sha3"
	"google.golang.org/protobuf/proto"
	"source.quilibrium.com/quilibrium/monorepo/node/protobufs"
)

type DataProofStore interface {
	NewTransaction() (Transaction, error)
	GetAggregateProof(
		filter []byte,
		commitment []byte,
		frameNumber uint64,
		inclusionReassembler func(typeUrl string, data [][]byte) ([]byte, error),
	) (
		*protobufs.InclusionAggregateProof,
		error,
	)
	PutAggregateProof(
		txn Transaction,
		aggregateProof *protobufs.InclusionAggregateProof,
		commitment []byte,
		inclusionSplitter func(typeUrl string, data []byte) ([][]byte, error),
	) error
}

type PebbleDataProofStore struct {
	db     KVDB
	logger *zap.Logger
}

func NewPebbleDataProofStore(
	db KVDB,
	logger *zap.Logger,
) *PebbleDataProofStore {
	return &PebbleDataProofStore{
		db,
		logger,
	}
}

const (
	DATA_PROOF           = 0x04
	DATA_PROOF_METADATA  = 0x00
	DATA_PROOF_INCLUSION = 0x01
	DATA_PROOF_SEGMENT   = 0x02
)

func dataProofMetadataKey(filter []byte, commitment []byte) []byte {
	key := []byte{DATA_PROOF, DATA_PROOF_METADATA}
	key = append(key, commitment...)
	key = append(key, filter...)
	return key
}

func dataProofInclusionKey(
	filter []byte,
	commitment []byte,
	seqNo uint64,
) []byte {
	key := []byte{DATA_PROOF, DATA_PROOF_INCLUSION}
	key = append(key, commitment...)
	key = binary.BigEndian.AppendUint64(key, seqNo)
	key = append(key, filter...)
	return key
}

func dataProofSegmentKey(
	filter []byte,
	hash []byte,
) []byte {
	key := []byte{DATA_PROOF, DATA_PROOF_SEGMENT}
	key = append(key, hash...)
	key = append(key, filter...)
	return key
}

func (p *PebbleDataProofStore) NewTransaction() (Transaction, error) {
	return p.db.NewBatch(), nil
}

func internalGetAggregateProof(
	db KVDB,
	filter []byte,
	commitment []byte,
	frameNumber uint64,
) (*protobufs.InclusionAggregateProof, error) {
	value, closer, err := db.Get(dataProofMetadataKey(filter, commitment))
	if err != nil {
		if errors.Is(err, pebble.ErrNotFound) {
			return nil, ErrNotFound
		}

		return nil, errors.Wrap(err, "get aggregate proof")
	}

	defer closer.Close()
	copied := make([]byte, len(value[8:]))
	limit := binary.BigEndian.Uint64(value[0:8])
	copy(copied, value[8:])

	aggregate := &protobufs.InclusionAggregateProof{
		Filter:               filter,
		FrameNumber:          frameNumber,
		InclusionCommitments: []*protobufs.InclusionCommitment{},
		Proof:                copied,
	}

	iter, err := db.NewIter(
		dataProofInclusionKey(filter, commitment, 0),
		dataProofInclusionKey(filter, commitment, limit+1),
	)
	if err != nil {
		return nil, errors.Wrap(err, "get aggregate proof")
	}

	i := uint32(0)

	for iter.First(); iter.Valid(); iter.Next() {
		incCommit := iter.Value()

		urlLength := binary.BigEndian.Uint16(incCommit[:2])
		commitLength := binary.BigEndian.Uint16(incCommit[2:4])

		url := make([]byte, urlLength)
		copy(url, incCommit[4:urlLength+4])

		commit := make([]byte, commitLength)
		copy(commit, incCommit[urlLength+4:urlLength+4+commitLength])

		remainder := int(urlLength + 4 + commitLength)

		inclusionCommitment := &protobufs.InclusionCommitment{
			Filter:      filter,
			FrameNumber: frameNumber,
			Position:    i,
			TypeUrl:     string(url),
			Commitment:  commit,
		}

		chunks := [][]byte{}
		for j := 0; j < (len(incCommit)-remainder)/32; j++ {
			start := remainder + (j * 32)
			end := remainder + ((j + 1) * 32)
			segValue, dataCloser, err := db.Get(
				dataProofSegmentKey(filter, incCommit[start:end]),
			)
			if err != nil {
				if errors.Is(err, pebble.ErrNotFound) {
					// If we've lost this key it means we're in a corrupted state
					return nil, ErrInvalidData
				}

				return nil, errors.Wrap(err, "get aggregate proof")
			}

			segCopy := make([]byte, len(segValue))
			copy(segCopy, segValue)
			chunks = append(chunks, segCopy)

			if err = dataCloser.Close(); err != nil {
				return nil, errors.Wrap(err, "get aggregate proof")
			}
		}

		if string(url) == protobufs.IntrinsicExecutionOutputType {
			o := &protobufs.IntrinsicExecutionOutput{}
			copiedLeft := make([]byte, len(chunks[0]))
			copiedRight := make([]byte, len(chunks[1]))
			copy(copiedLeft, chunks[0])
			copy(copiedRight, chunks[1])

			o.Address = copiedLeft[:32]
			o.Output = copiedLeft[32:]
			o.Proof = copiedRight
			inclusionCommitment.Data, err = proto.Marshal(o)
			if err != nil {
				return nil, errors.Wrap(err, "get aggregate proof")
			}
		} else {
			copied := make([]byte, len(chunks[0]))
			copy(copied, chunks[0])
			inclusionCommitment.Data = copied
		}

		aggregate.InclusionCommitments = append(
			aggregate.InclusionCommitments,
			inclusionCommitment,
		)
		i++
	}

	if err = iter.Close(); err != nil {
		return nil, errors.Wrap(err, "get aggregate proof")
	}

	return aggregate, nil
}

func (p *PebbleDataProofStore) GetAggregateProof(
	filter []byte,
	commitment []byte,
	frameNumber uint64,
) (*protobufs.InclusionAggregateProof, error) {
	return internalGetAggregateProof(
		p.db,
		filter,
		commitment,
		frameNumber,
	)
}

func internalPutAggregateProof(
	db KVDB,
	txn Transaction,
	aggregateProof *protobufs.InclusionAggregateProof,
	commitment []byte,
	inclusionSplitter func(typeUrl string, data []byte) ([][]byte, error),
) error {
	buf := binary.BigEndian.AppendUint64(
		nil,
		uint64(len(aggregateProof.InclusionCommitments)),
	)
	buf = append(buf, aggregateProof.Proof...)

	for i, inc := range aggregateProof.InclusionCommitments {
		segments, err := inclusionSplitter(inc.TypeUrl, inc.Data)
		if err != nil {
			return errors.Wrap(err, "get aggregate proof")
		}

		urlLength := len(inc.TypeUrl)
		commitLength := len(inc.Commitment)
		encoded := binary.BigEndian.AppendUint16(nil, uint16(urlLength))
		encoded = binary.BigEndian.AppendUint16(encoded, uint16(commitLength))

		encoded = append(encoded, []byte(inc.TypeUrl)...)
		encoded = append(encoded, inc.Commitment...)

		for _, segment := range segments {
			hash := sha3.Sum256(segment)
			if err = txn.Set(
				dataProofSegmentKey(aggregateProof.Filter, hash[:]),
				segment,
			); err != nil {
				return errors.Wrap(err, "put aggregate proof")
			}
			encoded = append(encoded, hash[:]...)
		}

		if err = txn.Set(
			dataProofInclusionKey(aggregateProof.Filter, commitment, uint64(i)),
			encoded,
		); err != nil {
			return errors.Wrap(err, "put aggregate proof")
		}
	}

	if err := txn.Set(
		dataProofMetadataKey(aggregateProof.Filter, commitment),
		buf,
	); err != nil {
		return errors.Wrap(err, "put aggregate proof")
	}

	return nil
}

func (p *PebbleDataProofStore) PutAggregateProof(
	txn Transaction,
	aggregateProof *protobufs.InclusionAggregateProof,
	commitment []byte,
	inclusionSplitter func(typeUrl string, data []byte) ([][]byte, error),
) error {
	return internalPutAggregateProof(
		p.db,
		txn,
		aggregateProof,
		commitment,
		inclusionSplitter,
	)
}