package cache

import (
	"bytes"
	"compress/gzip"
	"compress/zlib"
	"encoding/json"
	jsonv2 "encoding/json/v2"
	"fmt"
	"io"
	"math"
	"sync"

	jsoniter "github.com/json-iterator/go"
	"github.com/klauspost/compress/s2"
	"github.com/vmihailenco/msgpack/v5"
)

// ManifestCompressionType defines the compression algorithm for cached manifests.
type ManifestCompressionType string

// ManifestStorageType defines the serialization format for cached manifests.
type ManifestStorageType string

const (
	ManifestCompressionGZip           ManifestCompressionType = "gzip"
	ManifestCompressionGZipDefault    ManifestCompressionType = "gzip-default"
	ManifestCompressionGZipBestSpeed  ManifestCompressionType = "gzip-bestspeed"
	ManifestCompressionS2Encode       ManifestCompressionType = "s2-encode"
	ManifestCompressionS2EncodeBetter ManifestCompressionType = "s2-encodebetter"
	ManifestCompressionZLib           ManifestCompressionType = "zlib"
	ManifestCompressionNone           ManifestCompressionType = "none"

	ManifestStorageJSON     ManifestStorageType = "json"
	ManifestStorageJSONIter ManifestStorageType = "jsoniter"
	ManifestStorageMsgPack  ManifestStorageType = "msgpack"
)

// manifestSerializer handles serialization and deserialization of manifest objects.
type manifestSerializer interface {
	marshal(obj map[string]any) ([]byte, error)
	unmarshal(data []byte) (map[string]any, error)
}

// manifestCompressor handles compression and decompression of manifest bytes.
type manifestCompressor interface {
	compress(data []byte) ([]byte, error)
	decompress(data []byte) ([]byte, error)
}

// serializers maps each ManifestStorageType to its codec implementation.
var serializers = map[ManifestStorageType]manifestSerializer{
	ManifestStorageJSON:     jsonSerializer{},
	ManifestStorageJSONIter: jsoniterSerializer{},
	ManifestStorageMsgPack:  msgpackSerializer{},
}

// compressors maps each ManifestCompressionType to its codec implementation.
var compressors = map[ManifestCompressionType]manifestCompressor{
	ManifestCompressionGZipDefault:    &gzipCompressor{level: gzip.DefaultCompression},
	ManifestCompressionGZipBestSpeed:  &gzipCompressor{level: gzip.BestSpeed},
	ManifestCompressionS2Encode:       s2Compressor{better: false},
	ManifestCompressionS2EncodeBetter: s2Compressor{better: true},
	ManifestCompressionZLib:           zlibCompressor{},
	ManifestCompressionNone:           noneCompressor{},
}

// normalizeManifestCompressionType normalizes a compression type to a known value.
// "gzip" is an alias for "gzip-bestspeed"; unknown values default to "gzip-bestspeed".
func normalizeManifestCompressionType(ct ManifestCompressionType) ManifestCompressionType {
	switch ct {
	case ManifestCompressionGZip:
		return ManifestCompressionGZipBestSpeed
	case ManifestCompressionGZipDefault, ManifestCompressionGZipBestSpeed,
		ManifestCompressionS2Encode, ManifestCompressionS2EncodeBetter,
		ManifestCompressionZLib, ManifestCompressionNone:
		return ct
	default:
		return ManifestCompressionGZipBestSpeed
	}
}

// normalizeManifestStorageType normalizes a storage type to a known value.
// Unknown values default to "json".
func normalizeManifestStorageType(st ManifestStorageType) ManifestStorageType {
	switch st {
	case ManifestStorageJSON, ManifestStorageJSONIter, ManifestStorageMsgPack:
		return st
	default:
		return ManifestStorageJSON
	}
}

// serializeManifestObject serializes a manifest object using the registered serializer.
func serializeManifestObject(obj map[string]any, storageType ManifestStorageType) ([]byte, error) {
	s, ok := serializers[storageType]
	if !ok {
		return nil, fmt.Errorf("unknown storage type: %s", storageType)
	}
	return s.marshal(obj)
}

// deserializeManifestObject deserializes manifest bytes using the registered serializer.
func deserializeManifestObject(data []byte, storageType ManifestStorageType) (map[string]any, error) {
	s, ok := serializers[storageType]
	if !ok {
		return nil, fmt.Errorf("unknown storage type: %s", storageType)
	}
	return s.unmarshal(data)
}

// compressManifestData compresses manifest bytes using the registered compressor.
func compressManifestData(data []byte, compressionType ManifestCompressionType) ([]byte, error) {
	c, ok := compressors[compressionType]
	if !ok {
		return nil, fmt.Errorf("unknown compression type: %s", compressionType)
	}
	compressed, err := c.compress(data)
	if err != nil {
		return nil, err
	}
	return trimExcessCapacity(compressed), nil
}

// trimExcessCapacity copies b into a right-sized allocation when its spare capacity is large.
// The result of compressManifestData is kept in the cluster cache for the lifetime of the
// resource, so spare capacity left by the encoder (bytes.Buffer growth, s2.MaxEncodedLen
// preallocation, serializer output buffers) would otherwise stay allocated alongside every
// cached manifest.
//
// Small slack is left alone, because a fresh copy would be rounded up to the same allocation size
// and the copy would buy nothing. Up to 32 KiB, Go rounds allocations to a size class (at most
// ~12.5% waste, hence len/8). Above that it rounds to whole 8 KiB pages, so slack beyond one page
// can always be reclaimed.
func trimExcessCapacity(b []byte) []byte {
	if cap(b)-len(b) <= min(len(b)/8, 8<<10) {
		return b
	}
	trimmed := make([]byte, len(b))
	copy(trimmed, b)
	return trimmed
}

// decompressManifestData decompresses manifest bytes using the registered compressor.
func decompressManifestData(data []byte, compressionType ManifestCompressionType) ([]byte, error) {
	c, ok := compressors[compressionType]
	if !ok {
		return nil, fmt.Errorf("unknown compression type: %s", compressionType)
	}
	return c.decompress(data)
}

type jsonSerializer struct{}

func (jsonSerializer) marshal(obj map[string]any) ([]byte, error) {
	data, err := json.Marshal(obj)
	if err != nil {
		return nil, fmt.Errorf("json marshal failed: %w", err)
	}
	return data, nil
}

func (jsonSerializer) unmarshal(data []byte) (map[string]any, error) {
	var obj map[string]any
	// v2 API is ~2x faster: v1's AllowDuplicateNames disables v2's non-reflective path for any.
	// Its stricter defaults (no duplicate names, valid UTF-8) always hold because data is our
	// own json.Marshal output.
	if err := jsonv2.Unmarshal(data, &obj); err != nil {
		return nil, fmt.Errorf("json unmarshal failed: %w", err)
	}
	return obj, nil
}

type jsoniterSerializer struct{}

var jsoniterAPI = jsoniter.ConfigCompatibleWithStandardLibrary

func (jsoniterSerializer) marshal(obj map[string]any) ([]byte, error) {
	data, err := jsoniterAPI.Marshal(obj)
	if err != nil {
		return nil, fmt.Errorf("jsoniter marshal failed: %w", err)
	}
	return data, nil
}

func (jsoniterSerializer) unmarshal(data []byte) (map[string]any, error) {
	var obj map[string]any
	if err := jsoniterAPI.Unmarshal(data, &obj); err != nil {
		return nil, fmt.Errorf("jsoniter unmarshal failed: %w", err)
	}
	return obj, nil
}

type msgpackSerializer struct{}

func (msgpackSerializer) marshal(obj map[string]any) ([]byte, error) {
	data, err := msgpack.Marshal(obj)
	if err != nil {
		return nil, fmt.Errorf("msgpack marshal failed: %w", err)
	}
	return data, nil
}

func (msgpackSerializer) unmarshal(data []byte) (map[string]any, error) {
	var obj map[string]any
	if err := msgpack.Unmarshal(data, &obj); err != nil {
		return nil, fmt.Errorf("msgpack unmarshal failed: %w", err)
	}
	normalizeManifestValue(obj)
	return obj, nil
}

type gzipCompressor struct {
	level int
	pool  sync.Pool
}

func (c *gzipCompressor) compress(data []byte) ([]byte, error) {
	var buf bytes.Buffer

	var gz *gzip.Writer
	if v := c.pool.Get(); v != nil {
		gz = v.(*gzip.Writer)
		gz.Reset(&buf)
	} else {
		var err error
		gz, err = gzip.NewWriterLevel(&buf, c.level)
		if err != nil {
			return nil, fmt.Errorf("failed to create gzip writer: %w", err)
		}
	}

	if _, err := gz.Write(data); err != nil {
		return nil, fmt.Errorf("failed to write gzip data: %w", err)
	}
	if err := gz.Close(); err != nil {
		return nil, fmt.Errorf("failed to close gzip writer: %w", err)
	}
	// Detach the writer from buf before pooling it so the pool doesn't keep the untrimmed output buffer alive.
	gz.Reset(io.Discard)
	c.pool.Put(gz)
	return buf.Bytes(), nil
}

func (*gzipCompressor) decompress(data []byte) ([]byte, error) {
	gz, err := gzip.NewReader(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("failed to create gzip reader: %w", err)
	}
	defer gz.Close()
	result, err := io.ReadAll(gz)
	if err != nil {
		return nil, fmt.Errorf("failed to read gzip data: %w", err)
	}
	return result, nil
}

type s2Compressor struct{ better bool }

func (c s2Compressor) compress(data []byte) ([]byte, error) {
	if c.better {
		return s2.EncodeBetter(nil, data), nil
	}
	return s2.Encode(nil, data), nil
}

func (s2Compressor) decompress(data []byte) ([]byte, error) {
	result, err := s2.Decode(nil, data)
	if err != nil {
		return nil, fmt.Errorf("failed to decode s2 data: %w", err)
	}
	return result, nil
}

type zlibCompressor struct{}

func (zlibCompressor) compress(data []byte) ([]byte, error) {
	var buf bytes.Buffer
	w := zlib.NewWriter(&buf)
	if _, err := w.Write(data); err != nil {
		return nil, fmt.Errorf("failed to write zlib data: %w", err)
	}
	if err := w.Close(); err != nil {
		return nil, fmt.Errorf("failed to close zlib writer: %w", err)
	}
	return buf.Bytes(), nil
}

func (zlibCompressor) decompress(data []byte) ([]byte, error) {
	r, err := zlib.NewReader(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("failed to create zlib reader: %w", err)
	}
	defer r.Close()
	result, err := io.ReadAll(r)
	if err != nil {
		return nil, fmt.Errorf("failed to read zlib data: %w", err)
	}
	return result, nil
}

type noneCompressor struct{}

func (noneCompressor) compress(data []byte) ([]byte, error)   { return data, nil }
func (noneCompressor) decompress(data []byte) ([]byte, error) { return data, nil }

// normalizeManifestValue normalizes msgpack-decoded values to match
// unstructured.Unstructured expectations: integers become int64 where possible,
// and map[any]any becomes map[string]any.
func normalizeManifestValue(obj map[string]any) {
	for k, v := range obj {
		obj[k] = normalizeValue(v)
	}
}

func normalizeValue(v any) any {
	switch val := v.(type) {
	case map[string]any:
		normalizeManifestValue(val)
		return val
	case map[any]any:
		converted := make(map[string]any, len(val))
		for mk, mv := range val {
			converted[fmt.Sprintf("%v", mk)] = normalizeValue(mv)
		}
		return converted
	case []any:
		for i, item := range val {
			val[i] = normalizeValue(item)
		}
		return val
	case uint64:
		if val <= math.MaxInt64 {
			return int64(val)
		}
		return float64(val)
	case int64:
		return val
	case uint32:
		return int64(val)
	case int32:
		return int64(val)
	case int:
		return int64(val)
	case uint:
		if uint64(val) <= math.MaxInt64 {
			return int64(val)
		}
		return float64(val)
	default:
		return v
	}
}