// KEKS/CM -- KEKS-encoded cryptographic messages
// Copyright (C) 2024-2025 Sergey Matveev <stargrave@stargrave.org>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser 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 Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this program.  If not, see <http://www.gnu.org/licenses/>.

package chapoly

import (
	"bytes"
	"crypto/hkdf"
	"crypto/subtle"
	"errors"
	"hash"
	"io"

	"github.com/aead/chacha20"
	"go.cypherpunks.su/keks"
	"golang.org/x/crypto/blake2b"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/poly1305"
)

const (
	ChunkLen       = 128 * 1024
	CommitmentLen  = 32
	CEKLen         = blake2b.Size
	DEMAlgo        = "xchapoly-krkc"
	Poly1305KeyLen = 32
)

type keymat struct {
	key []byte
	mac [Poly1305KeyLen]byte
	iv  []byte
}

type job struct {
	bufReady  chan struct{}
	processed chan error
	keymat    keymat
	buf       []byte
}

func blake2bHash() hash.Hash {
	h, err := blake2b.New512(nil)
	if err != nil {
		panic(err)
	}
	return h
}

func do(
	w io.Writer,
	blob, seal bool,
	r io.Reader,
	cek []byte,
	procs int,
) (total int64, err error) {
	if len(cek) != CEKLen {
		return 0, errors.New("wrong CEK len")
	}
	ready := make(chan *job, procs)
	dones := make(chan *job, procs)
	keymats := make(chan keymat, procs)
	go func() {
		ck := cek
		var key []byte
		var mac []byte
		var iv []byte
		var errHKDF error
		for {
			ck, errHKDF = hkdf.Extract(blake2bHash, nil, ck)
			if errHKDF != nil {
				panic(errHKDF)
			}
			ck, errHKDF = hkdf.Expand(
				blake2bHash, ck, "cm/encrypted/xchapoly-krkc/kr", CEKLen)
			if errHKDF != nil {
				panic(errHKDF)
			}
			key, errHKDF = hkdf.Expand(
				blake2bHash, ck, "cm/encrypted/xchapoly-krkc/key",
				chacha20poly1305.KeySize)
			if errHKDF != nil {
				panic(errHKDF)
			}
			mac, errHKDF = hkdf.Expand(
				blake2bHash, ck, "cm/encrypted/xchapoly-krkc/mac",
				Poly1305KeyLen)
			if errHKDF != nil {
				panic(errHKDF)
			}
			iv, errHKDF = hkdf.Expand(
				blake2bHash, ck, "cm/encrypted/xchapoly-krkc/iv",
				chacha20poly1305.NonceSizeX)
			if errHKDF != nil {
				panic(errHKDF)
			}
			k := keymat{key: key, iv: iv}
			copy(k.mac[:], mac)
			keymats <- k
		}
	}()
	blobChunkLen := ChunkLen + +CommitmentLen
	var blobDecoder *keks.BlobDecoder
	if seal {
		if blob {
			total, err = keks.BlobAtomEncode(w, int64(blobChunkLen))
		}
	} else {
		if blob {
			blobDecoder, err = keks.NewBlobDecoder(r, int64(blobChunkLen))
		}
	}
	if err != nil {
		return
	}
	for range procs {
		go func() {
			j := job{
				buf:       make([]byte, blobChunkLen),
				bufReady:  make(chan struct{}),
				processed: make(chan error),
			}
			ready <- &j
			var errJob error
			var com hash.Hash
			var tag [16]byte
			ourCom := make([]byte, CommitmentLen)
			for {
				<-j.bufReady
				com, errJob = blake2b.New(CommitmentLen, nil)
				if errJob != nil {
					panic(errJob)
				}
				if seal {
					chacha20.XORKeyStream(
						j.buf[:len(j.buf)-CommitmentLen],
						j.buf[:len(j.buf)-CommitmentLen],
						j.keymat.iv,
						j.keymat.key,
					)
					poly1305.Sum(
						&tag,
						j.buf[:len(j.buf)-CommitmentLen],
						&j.keymat.mac,
					)
					com.Write(j.keymat.key)
					com.Write(j.keymat.iv)
					com.Write(j.keymat.mac[:])
					com.Write(tag[:])
					com.Sum(j.buf[:len(j.buf)-CommitmentLen])
				} else {
					poly1305.Sum(
						&tag,
						j.buf[:len(j.buf)-CommitmentLen],
						&j.keymat.mac,
					)
					com.Write(j.keymat.key)
					com.Write(j.keymat.iv)
					com.Write(j.keymat.mac[:])
					com.Write(tag[:])
					ourCom = com.Sum(ourCom[:0])
					if subtle.ConstantTimeCompare(
						ourCom,
						j.buf[len(j.buf)-CommitmentLen:],
					) == 1 {
						chacha20.XORKeyStream(
							j.buf[:len(j.buf)-CommitmentLen],
							j.buf[:len(j.buf)-CommitmentLen],
							j.keymat.iv,
							j.keymat.key,
						)
						j.buf = j.buf[:len(j.buf)-CommitmentLen]
					} else {
						errJob = errors.New("commitment differs")
					}
				}
				j.processed <- errJob
			}
		}()
	}
	finished := make(chan struct{})
	var errW error
	go func() {
		var n int64
		for j := range dones {
			errW = <-j.processed
			if errW != nil {
				j.buf = j.buf[:blobChunkLen]
				ready <- j
				break
			}
			if seal {
				if blob {
					n, errW = keks.BinEncode(w, j.buf)
				} else {
					n, errW = io.Copy(w, bytes.NewReader(j.buf))
				}
			} else {
				if len(j.buf) == 0 {
					n = 0
				} else {
					n, errW = io.Copy(w, bytes.NewReader(j.buf))
				}
			}
			total += n
			if errW != nil || n == 0 {
				break
			}
			j.buf = j.buf[:blobChunkLen]
			ready <- j
		}
		close(finished)
	}()
	var chunk []byte
	var errR error
	var eof bool
	var n int
	var j *job
	var tailMet bool
	for !eof {
		select {
		case <-finished:
			err = errW
			return
		default:
		}
		j = <-ready
		if seal {
			n, errR = io.ReadFull(r, j.buf[:ChunkLen])
		} else {
			if blobDecoder == nil {
				n, errR = io.ReadFull(r, j.buf)
				j.buf = j.buf[:n]
			} else {
				blobDecoder.SetReadBuf(j.buf)
				chunk, errR = blobDecoder.Next()
				n = len(chunk)
				j.buf = j.buf[:n]
			}
		}
		total += int64(n)
		if errR != nil {
			if !(errR == io.ErrUnexpectedEOF || errR == io.EOF) {
				err = errR
				close(dones)
				<-finished
				return
			}
			errR = nil
			eof = true
			if seal {
				j.buf = j.buf[:n+CommitmentLen]
			}
		}
		if !seal && n == 0 {
			break
		}
		j.keymat = <-keymats
		if (seal && n < ChunkLen) || (!seal && n < blobChunkLen) {
			j.keymat.iv[chacha20poly1305.NonceSizeX-1] |= 0x01
			tailMet = true
		} else {
			j.keymat.iv[chacha20poly1305.NonceSizeX-1] &= 0xFE
		}
		j.bufReady <- struct{}{}
		dones <- j
	}
	close(dones)
	<-finished
	err = errW
	if err != nil {
		return
	}
	if !tailMet {
		err = errors.New("no tail met")
	}
	return
}

func Seal(w io.Writer, r io.Reader, key []byte, procs int) (total int64, err error) {
	return do(w, false, true, r, key, procs)
}

func Open(w io.Writer, r io.Reader, key []byte, procs int) (total int64, err error) {
	return do(w, false, false, r, key, procs)
}

func SealBlob(w io.Writer, r io.Reader, key []byte, procs int) (total int64, err error) {
	return do(w, true, true, r, key, procs)
}

func OpenBlob(w io.Writer, r io.Reader, key []byte, procs int) (total int64, err error) {
	return do(w, true, false, r, key, procs)
}
