package nwsync

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"fmt"

	"github.com/klauspost/compress/zstd"
)

// NWCompressedBuffer framing, as upstream's neverwinter/compressedbuf.nim
// writes it for NWSync blobs. All fields are little-endian uint32:
//
//	magic "NSYC", version 3, algorithm 2 (zstd), uncompressed size,
//	zstd header version 1, dictionary 0, then the raw zstd frame.
const (
	blobMagic       = 0x4359534E // "NSYC" little-endian
	blobVersion     = 3
	algorithmZstd   = 2
	zstdHeaderVer   = 1
	zstdDictionary  = 0
	blobHeaderBytes = 24
)

// EncodeAll/DecodeAll are single-threaded per call, so the default pool of one
// encoder per CPU only buys idle memory: each holds a window-sized history, so
// on a 24-core runner that is ~200 MB of live heap doing nothing. Concurrency 1
// produces byte-identical output.
var (
	blobEncoder, _ = zstd.NewWriter(nil, zstd.WithEncoderConcurrency(1))
	blobDecoder, _ = zstd.NewReader(nil, zstd.WithDecoderConcurrency(1))
)

// zstd frame header bits we care about. A frame starts with the magic, then a
// one-byte Frame_Header_Descriptor: bits 7-6 size the Frame_Content_Size field,
// bit 5 is Single_Segment_flag, bits 1-0 size the Dictionary_ID field.
const (
	zstdFrameMagic      = "\x28\xb5\x2f\xfd"
	frameSingleSegment  = 1 << 5
	frameDictionaryMask = 0x03
	// The size below which klauspost/compress writes no Frame_Content_Size,
	// matching the field's own encoding threshold.
	frameSizeThreshold = 256
)

// compressBlob wraps data in NWCompressedBuffer framing.
func compressBlob(data []byte) []byte {
	var out bytes.Buffer
	header := []uint32{blobMagic, blobVersion, algorithmZstd, uint32(len(data)), zstdHeaderVer, zstdDictionary}
	for _, field := range header {
		_ = binary.Write(&out, binary.LittleEndian, field)
	}
	out.Write(declareFrameContentSize(blobEncoder.EncodeAll(data, nil), len(data)))
	return out.Bytes()
}

// inspectBlob unwraps a stored blob the way the game client reads it, and is the
// only reader that should be trusted to judge a published blob.
//
// It asserts the frame property on top of the round trip. Go's decoder — like
// the zstd CLI — streams a frame that declares no content size, so a check that
// only decompresses and hashes is a *more* capable decoder than the client's: it
// certifies exactly the blobs the client rejects, which is how #86 reached
// production and survived an audit.
func inspectBlob(blob []byte) ([]byte, error) {
	data, err := decompressBlob(blob)
	if err != nil {
		return nil, fmt.Errorf("malformed framing: %w", err)
	}
	if len(data) > 0 && !frameDeclaresContentSize(blob[blobHeaderBytes:]) {
		return nil, fmt.Errorf("malformed framing: the zstd frame declares no content size, which the game client cannot decode")
	}
	return data, nil
}

// blobMatchesName holds a stored blob to its own file name: a blob is named
// after the sha1 of its uncompressed bytes, so the name is a complete statement
// about the contents and nothing else is needed to check it.
func blobMatchesName(blob []byte, sha1Hex string) error {
	data, err := inspectBlob(blob)
	if err != nil {
		return err
	}
	if got := hex.EncodeToString(sha1Sum(data)); got != sha1Hex {
		return fmt.Errorf("blob %s holds the contents of %s", sha1Hex, got)
	}
	return nil
}

// frameDeclaresContentSize reports whether a zstd frame states how many bytes it
// decompresses to. A frame with a zero-sized Frame_Content_Size field declares
// one only when Single_Segment_flag is set; otherwise the size is unknown.
func frameDeclaresContentSize(frame []byte) bool {
	if len(frame) < 5 || string(frame[:4]) != zstdFrameMagic {
		return false
	}
	descriptor := frame[4]
	return descriptor>>6 != 0 || descriptor&frameSingleSegment != 0
}

// declareFrameContentSize rewrites a frame that does not declare its
// Frame_Content_Size so that it does, and returns any other frame unchanged.
//
// klauspost/compress omits the field for inputs under 256 bytes, which the spec
// permits. Reference libzstd never does, so the NWN client — which sizes its
// output buffer from ZSTD_getFrameContentSize and has therefore never met a
// frame without one — rejects the blob outright with an empty "potential
// compression error" (#86). No encoder option changes this, so the frame is
// re-headered here.
//
// The result is the shape libzstd itself emits for the same input: setting
// Single_Segment_flag drops the Window_Descriptor byte, and the freed byte pays
// for a one-byte Frame_Content_Size. Window_Size then equals the content size,
// which is sound because the content is under 256 bytes and every match in it
// therefore falls inside that window. Same length in, same length out.
func declareFrameContentSize(frame []byte, size int) []byte {
	if size <= 0 || size >= frameSizeThreshold || len(frame) < 6 || string(frame[:4]) != zstdFrameMagic {
		return frame
	}
	descriptor := frame[4]
	// Rewrite only the exact shape a small input produces: no declared size, no
	// single segment, no dictionary. Anything else either declares a size
	// already or is not a frame this reinterpretation is safe on.
	if descriptor>>6 != 0 || descriptor&frameSingleSegment != 0 || descriptor&frameDictionaryMask != 0 {
		return frame
	}
	reframed := make([]byte, len(frame))
	copy(reframed, frame)
	reframed[4] = descriptor | frameSingleSegment
	reframed[5] = byte(size) // replaces Window_Descriptor
	return reframed
}

// decompressBlob unwraps NWCompressedBuffer framing. It exists so a blob this
// package wrote — or one upstream wrote — can be compared by its uncompressed
// bytes, which is the only comparison that is meaningful across zstd
// implementations.
func decompressBlob(blob []byte) ([]byte, error) {
	if len(blob) < blobHeaderBytes {
		return nil, fmt.Errorf("blob too small: %d bytes", len(blob))
	}
	header := make([]uint32, 6)
	if err := binary.Read(bytes.NewReader(blob[:blobHeaderBytes]), binary.LittleEndian, header); err != nil {
		return nil, fmt.Errorf("decode blob header: %w", err)
	}
	switch {
	case header[0] != blobMagic:
		return nil, fmt.Errorf("invalid blob magic: %#x", header[0])
	case header[1] != blobVersion:
		return nil, fmt.Errorf("unsupported blob version: %d", header[1])
	case header[2] != algorithmZstd:
		return nil, fmt.Errorf("unsupported compression algorithm: %d", header[2])
	case header[4] != zstdHeaderVer:
		return nil, fmt.Errorf("unsupported zstd header version: %d", header[4])
	case header[5] != zstdDictionary:
		return nil, fmt.Errorf("zstd dictionaries are not supported")
	}
	if header[3] == 0 {
		return nil, nil
	}
	data, err := blobDecoder.DecodeAll(blob[blobHeaderBytes:], nil)
	if err != nil {
		return nil, fmt.Errorf("decompress blob: %w", err)
	}
	if uint32(len(data)) != header[3] {
		return nil, fmt.Errorf("blob size mismatch: header says %d, got %d", header[3], len(data))
	}
	return data, nil
}