// KEKS -- Go KEKS codec implementation
// 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 keks

import (
	"bytes"
	"errors"
	"io"
	"math/big"

	"go.cypherpunks.su/keks/be"
	"go.cypherpunks.su/tai64n/v4"
)

// Just write single byte. Can be convenient for writing LIST/MAP atoms.
func ByteEncode(w io.Writer, b byte) (written int64, err error) {
	_, err = w.Write([]byte{b})
	if err != nil {
		return
	}
	written = 1
	return
}

// Write an encoded TRUE/FALSE atom.
func BoolEncode(w io.Writer, v bool) (written int64, err error) {
	data := []byte{byte(AtomFalse)}
	if v {
		data[0] = byte(AtomTrue)
	}
	_, err = w.Write(data)
	if err != nil {
		return
	}
	written = 1
	return
}

// Write an encoded Hexlet atom.
func HexletEncode(w io.Writer, v *Hexlet) (written int64, err error) {
	return io.Copy(w, bytes.NewReader(append([]byte{byte(AtomHexlet)}, v[:]...)))
}

// Write an encoded Magic atom.
func MagicEncode(w io.Writer, m Magic) (written int64, err error) {
	if len(m) > 12 {
		return 0, errors.New("too long magic")
	}
	var v [16]byte
	v[0] = 'K'
	v[1] = 'E'
	v[2] = 'K'
	v[3] = 'S'
	copy(v[4:], []byte(m))
	return io.Copy(w, bytes.NewReader(v[:]))
}

func atomUintEncode(w io.Writer, v uint64) (written int64, err error) {
	if v == 0 {
		return BinEncode(w, []byte{})
	}
	l := 0
	for ; l < 7; l++ {
		if v < (1 << ((l + 1) * 8)) {
			break
		}
	}
	buf := make([]byte, l+1)
	be.Put(buf, v)
	return BinEncode(w, buf)
}

// Write an encoded +INT atom.
func UIntEncode(w io.Writer, v uint64) (written int64, err error) {
	_, err = w.Write([]byte{byte(AtomPInt)})
	if err != nil {
		return
	}
	written, err = atomUintEncode(w, v)
	written++
	return
}

// Write an encoded -INT atom if v is negative.
// +INT atom otherwise, same as UIntEncode.
func IntEncode(w io.Writer, v int64) (written int64, err error) {
	if v >= 0 {
		return UIntEncode(w, uint64(v))
	}
	_, err = w.Write([]byte{byte(AtomNInt)})
	if err != nil {
		return
	}
	written, err = atomUintEncode(w, uint64(-(v + 1)))
	written++
	return
}

// Write an encoded ±INT atom.
func BigIntEncode(w io.Writer, v *big.Int) (written int64, err error) {
	if v.Sign() >= 0 {
		_, err = w.Write([]byte{byte(AtomPInt)})
	} else {
		n := new(big.Int).Abs(v)
		v = n.Sub(n, big.NewInt(1))
		_, err = w.Write([]byte{byte(AtomNInt)})
	}
	if err != nil {
		return
	}
	written, err = BinEncode(w, v.Bytes())
	written++
	return
}

// Write an encoded FLOAT atom.
func FloatEncode(w io.Writer, v *Float) (written int64, err error) {
	if v.NaN {
		return ByteEncode(w, byte(AtomFloatNaN))
	}
	if v.Pinf {
		return ByteEncode(w, byte(AtomFloatPinf))
	}
	if v.Ninf {
		return ByteEncode(w, byte(AtomFloatNinf))
	}
	_, err = w.Write([]byte{byte(AtomFloat)})
	if err != nil {
		return
	}
	written++
	var wr int64
	wr, err = BigIntEncode(w, v.M)
	written += wr
	if err != nil {
		return
	}
	wr, err = BigIntEncode(w, v.E)
	written += wr
	return
}

// Write an encoded BLOB atom.
func BlobAtomEncode(w io.Writer, chunkLen int64) (written int64, err error) {
	l := make([]byte, 9)
	l[0] = byte(AtomBLOB)
	be.Put(l[1:], uint64(chunkLen-1))
	return io.Copy(w, bytes.NewReader(l))
}

// Write an encoded BLOB.
func BlobEncode(w io.Writer, chunkLen int64, r io.Reader) (written int64, err error) {
	if chunkLen == 0 {
		return 0, errors.New("zero chunkLen specified")
	}
	written, err = BlobAtomEncode(w, chunkLen)
	if err != nil {
		return
	}
	chunk := make([]byte, chunkLen)
	var n int
	var n64 int64
	for {
		n, err = io.ReadFull(r, chunk)
		if err != nil {
			if err == io.ErrUnexpectedEOF || err == io.EOF {
				chunk = chunk[:n]
				break
			}
			return
		}
		n64, err = BinEncode(w, chunk)
		if err != nil {
			return
		}
		written += n64
	}
	n64, err = BinEncode(w, chunk)
	written += n64
	return
}

func atomStrEncode(w io.Writer, data []byte, utf8 bool) (written int64, err error) {
	tag := byte(AtomStrings)
	if utf8 {
		tag |= AtomIsUTF8
	}
	var hdr []byte
	if len(data) >= 63+((1<<8)-1)+((1<<16)-1) {
		hdr = make([]byte, 8+1)
		hdr[0] = tag | 63
		be.Put(hdr[1:], uint64(len(data)-(63+((1<<8)-1)+((1<<16)-1))))
	} else if len(data) >= 62+255 {
		hdr = make([]byte, 2+1)
		hdr[0] = tag | 62
		be.Put(hdr[1:], uint64(len(data)-(62+((1<<8)-1))))
	} else if len(data) >= 61 {
		hdr = []byte{tag | 61, byte(len(data) - 61)}
	} else {
		hdr = []byte{tag | byte(len(data))}
	}
	var hdrLen int64
	hdrLen, err = io.Copy(w, bytes.NewReader(hdr))
	if err != nil {
		written = hdrLen
		return
	}
	written, err = io.Copy(w, bytes.NewReader(data))
	written += hdrLen
	return
}

// Write an encoded STR atom.
func StrEncode(w io.Writer, str string) (written int64, err error) {
	return atomStrEncode(w, []byte(str), true)
}

// Write an encoded BIN atom.
func BinEncode(w io.Writer, bin []byte) (written int64, err error) {
	return atomStrEncode(w, bin, false)
}

// Write an encoded TAI64 atom.
func TAI64Encode(w io.Writer, tai *tai64n.TAI64) (written int64, err error) {
	_, err = w.Write([]byte{byte(AtomTAI64)})
	if err != nil {
		return
	}
	written, err = io.Copy(w, bytes.NewReader(tai[:]))
	written++
	return
}

// Write an encoded TAI64N atom.
func TAI64NEncode(w io.Writer, tai *tai64n.TAI64N) (written int64, err error) {
	_, err = w.Write([]byte{byte(AtomTAI64N)})
	if err != nil {
		return
	}
	written, err = io.Copy(w, bytes.NewReader(tai[:]))
	written++
	return
}

// Write an encoded TAI64NA atom.
func TAI64NAEncode(w io.Writer, tai *tai64n.TAI64NA) (written int64, err error) {
	_, err = w.Write([]byte{byte(AtomTAI64NA)})
	if err != nil {
		return
	}
	written, err = io.Copy(w, bytes.NewReader(tai[:]))
	written++
	return
}
