package frame

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"image"
	"io"
	"sync"

	"gocv.io/x/gocv"
	"golang.org/x/crypto/pbkdf2"
)

const (
	saltSize   = 16
	keySize    = 32
	nonceSize  = 12

	// MaxPBKDF2Iter bounds the iteration count a Clip may claim on the wire,
	// so a malicious sender cannot force a receiver into an arbitrarily long
	// key-derivation CPU burn. Kept generous above every default and
	// configured value, but low enough that a rejected clip costs at most a
	// few hundred milliseconds of KDF work.
	MaxPBKDF2Iter = 1_000_000

	// encryptedOverhead is how many bytes Frame.Encrypt prepends to a
	// plaintext frame: 16 salt + 12 nonce + 16 GCM tag.
	encryptedOverhead = saltSize + nonceSize + 16

	// clipEncryptedOverhead is the per-frame overhead added by Clip.Encrypt:
	// the 16-byte salt lives on the Clip, so each frame gains only the
	// 12-byte nonce and 16-byte GCM tag.
	clipEncryptedOverhead = nonceSize + 16
)

// DefaultPBKDF2Iter is the PBKDF2 iteration count used for new clip and frame
// encryption. Encrypt stores the count actually used on the Clip, so Decrypt
// honors whatever the origin matched regardless of a node's local setting
// (override per node with frame.SetDefaultPBKDF2Iter).
var DefaultPBKDF2Iter = 600_000

// SetDefaultPBKDF2Iter overrides the default PBKDF2 iteration count used to
// encrypt new clips and decrypt foreign clips that do not carry their own
// count.
func SetDefaultPBKDF2Iter(n int) {
	if n < 1 {
		n = 1
	}
	if n > MaxPBKDF2Iter {
		n = MaxPBKDF2Iter
	}
	DefaultPBKDF2Iter = n
}

// isCV8U reports whether a MatType is an 8-bit unsigned depth, the only depth
// the pipeline understands: PixelBytes is counted as one byte per element, so
// a 16-bit or float type with a matching byte count would be decoded as a
// different size of matrix than the dimensions claim.
func isCV8U(t gocv.MatType) bool {
	return int(t)&0x07 == 0
}

// matTypeChannels recovers the channel count encoded in a MatType, so a clip
// declaring Channels can be cross-checked against its declared type.
func matTypeChannels(t gocv.MatType) int {
	return (int(t)>>3)&0x07 + 1
}

type DetectionCollection map[string][]Detection

type Detection struct {
	DetectionTitle        string // Car, human, human torso, bear, cat, animal, etc...
	DetectionMajorVersion uint
	DetectionMinorVersion uint
	DetectionPostfix      string // 2.3.4a4444248f
	DetectionRegion       image.Rectangle
	DetectionCertainty    float32
}

// Stores single frame data
type Frame struct {
	PixelBytes    []byte
	Width, Height uint
	GocvImageType gocv.MatType
	Detections    map[string][]Detection
	Channels      int
	Guid          []byte
	SourceData    string // sourcedata is <username>.<cameraID>
	Timestamp     uint64
}

func (f *Frame) ToMat() (gocv.Mat, error) {
	if f.Width == 0 || f.Height == 0 || f.Channels < 1 || f.Channels > 4 {
		return gocv.NewMat(), fmt.Errorf("frame: implausible dimensions %dx%d (%d channels)", f.Width, f.Height, f.Channels)
	}
	want := int(f.Width) * int(f.Height) * f.Channels
	if len(f.PixelBytes) != want {
		return gocv.NewMat(), fmt.Errorf("frame: %d pixel bytes, want %d for %dx%d x%d",
			len(f.PixelBytes), want, f.Width, f.Height, f.Channels)
	}
	mat, err := gocv.NewMatFromBytes(int(f.Height), int(f.Width), f.GocvImageType, f.PixelBytes)
	if err != nil {
		return gocv.NewMat(), err
	}
	return mat, nil
}

// Series of frames in a sequence from the same camera.
// Most metadata should be identical
// Frames should be in order as stored, so Timestamps[5] should be taken directly from PIxelMats[5] and Guids[5]
type Clip struct {
	PixelMats     [][]byte // First order is different frames, second order is pixel bytes from frames
	Width, Height uint
	Types         gocv.MatType
	Detections    []map[string][]Detection
	Comparisons   []Comparison
	Channels      int
	Guids         [][]byte
	SourceData    string
	Timestamps    []uint64
	Salt          []byte
	PBKDF2Iter    int // iteration count used by Encrypt, so a decrypting node can match it (0 = unknown)
}

func (clp *Clip) Sublimate() []Frame {
	frames := make([]Frame, len(clp.PixelMats))
	for i := range clp.PixelMats {
		f := Frame{
			PixelBytes:    clp.PixelMats[i],
			Width:         clp.Width,
			Height:        clp.Height,
			GocvImageType: clp.Types,
			Detections:    nil,
			Channels:      clp.Channels,
			SourceData:    clp.SourceData,
		}
		if i < len(clp.Detections) {
			f.Detections = clp.Detections[i]
		}
		if i < len(clp.Guids) {
			f.Guid = clp.Guids[i]
		}
		if i < len(clp.Timestamps) {
			f.Timestamp = clp.Timestamps[i]
		}
		frames[i] = f
	}
	return frames
}

// PixelMats, Guids, and Timestamps should all be the same len.
func (clp *Clip) CheckLenCorrelations() (bool, error) {
	l1 := len(clp.PixelMats)
	l2 := len(clp.Guids)
	l3 := len(clp.Timestamps)

	switch {
	case l1 != l2:
		return false, fmt.Errorf("clip has %d pixel frames but %d guids", l1, l2)
	case l2 != l3:
		return false, fmt.Errorf("clip has %d guids but %d timestamps", l2, l3)
	case l1 != l3:
		return false, fmt.Errorf("clip has %d pixel frames but %d timestamps", l1, l3)
	default:
		return true, nil
	}
}

func deriveKey(passphrase string, salt []byte, iterations int) []byte {
	return pbkdf2.Key([]byte(passphrase), salt, iterations, keySize, sha256.New)
}

// deriverCache memoizes the last handful of PBKDF2 derivations keyed by
// (salt, iterations), so a flood of repeated clips (replays or forged copies
// that share a salt) cannot repeatedly re-run an expensive KDF to bog the
// node down. It is bounded and safe for concurrent use.
var deriverCache = struct {
	sync.Mutex
	entries map[string][]byte
}{}

func cachedDeriveKey(passphrase string, salt []byte, iterations int) []byte {
	cacheKey := fmt.Sprintf("%x/%d", salt, iterations)

	deriverCache.Lock()
	got, ok := deriverCache.entries[cacheKey]
	deriverCache.Unlock()
	if ok {
		return got
	}

	derived := deriveKey(passphrase, salt, iterations)

	deriverCache.Lock()
	if deriverCache.entries == nil {
		deriverCache.entries = make(map[string][]byte)
	}
	if len(deriverCache.entries) >= 16 {
		deriverCache.entries = make(map[string][]byte)
	}
	deriverCache.entries[cacheKey] = derived
	deriverCache.Unlock()
	return derived
}

// gcmEncrypt seals plaintext with fresh random nonce and returns
// nonce||ciphertext (ciphertext includes the GCM tag).
func gcmEncrypt(plaintext, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}
	return append(nonce, gcm.Seal(nil, nonce, plaintext, nil)...), nil
}

// gcmDecrypt opens data produced by gcmEncrypt with the same key.
func gcmDecrypt(data, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	if len(data) < gcm.NonceSize() {
		return nil, errors.New("encrypted data too short")
	}
	nonce := data[:gcm.NonceSize()]
	return gcm.Open(nil, nonce, data[gcm.NonceSize():], nil)
}

// Encrypt encrypts PixelBytes using AES-256-GCM with a key derived from passphrase.
func (f *Frame) Encrypt(passphrase string) error {
	if len(f.PixelBytes) == 0 {
		return nil
	}

	salt := make([]byte, saltSize)
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return err
	}

	key := deriveKey(passphrase, salt, DefaultPBKDF2Iter)
	ciphertext, err := gcmEncrypt(f.PixelBytes, key)
	if err != nil {
		return err
	}

	f.PixelBytes = append(salt, ciphertext...)
	return nil
}

// Decrypt decrypts PixelBytes using AES-256-GCM with a key derived from passphrase.
func (f *Frame) Decrypt(passphrase string) error {
	return f.DecryptWithIter(passphrase, DefaultPBKDF2Iter)
}

// DecryptWithIter is Decrypt with an explicit iteration count, used by the
// legacy per-frame salted layout so a clip can honor the count its origin
// used.
func (f *Frame) DecryptWithIter(passphrase string, iterations int) error {
	if len(f.PixelBytes) == 0 {
		return nil
	}

	if len(f.PixelBytes) < saltSize+nonceSize+1 {
		return errors.New("encrypted data too short")
	}

	salt := f.PixelBytes[:saltSize]
	key := cachedDeriveKey(passphrase, salt, iterations)
	plaintext, err := gcmDecrypt(f.PixelBytes[saltSize:], key)
	if err != nil {
		return err
	}

	f.PixelBytes = plaintext
	return nil
}

// Encrypt seals every frame in the clip with a single PBKDF2 key derivation.
// The shared salt is stored on the clip; each frame gets its own random nonce.
// The iteration count used is recorded on the clip so any decrypting node can
// match it even when its local default differs. Frames are left untouched when
// their pixel buffer is empty.
func (c *Clip) Encrypt(passphrase string) error {
	if len(c.PixelMats) == 0 {
		return nil
	}

	c.PBKDF2Iter = DefaultPBKDF2Iter
	salt := make([]byte, saltSize)
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return err
	}
	key := deriveKey(passphrase, salt, c.PBKDF2Iter)

	var errs []error
	for i := range c.PixelMats {
		if len(c.PixelMats[i]) == 0 {
			continue
		}
		ct, err := gcmEncrypt(c.PixelMats[i], key)
		if err != nil {
			errs = append(errs, fmt.Errorf("frame %d: %w", i, err))
			continue
		}
		c.PixelMats[i] = ct
	}
	c.Salt = salt
	return errors.Join(errs...)
}

// Decrypt opens every frame previously sealed by Clip.Encrypt, deriving the
// key from the clip's stored salt. Clips without a clip-level salt
// (per-frame salted layout from before the single-KDF change) are handled via
// the per-frame fallback.
func (c *Clip) Decrypt(passphrase string) error {
	if len(c.PixelMats) == 0 {
		return nil
	}

	if len(c.Salt) == saltSize {
		iterations := c.PBKDF2Iter
		if iterations < 1 || iterations > MaxPBKDF2Iter {
			iterations = DefaultPBKDF2Iter
		}
		key := cachedDeriveKey(passphrase, c.Salt, iterations)
		var errs []error
		for i := range c.PixelMats {
			if len(c.PixelMats[i]) == 0 {
				continue
			}
			plaintext, err := gcmDecrypt(c.PixelMats[i], key)
			if err != nil {
				errs = append(errs, fmt.Errorf("frame %d: %w", i, err))
				continue
			}
			c.PixelMats[i] = plaintext
		}
		return errors.Join(errs...)
	}

	iterations := DefaultPBKDF2Iter
	if c.PBKDF2Iter >= 1 && c.PBKDF2Iter <= MaxPBKDF2Iter {
		iterations = c.PBKDF2Iter
	}
	var errs []error
	for i := range c.PixelMats {
		f := Frame{PixelBytes: c.PixelMats[i]}
		if err := f.DecryptWithIter(passphrase, iterations); err != nil {
			errs = append(errs, fmt.Errorf("frame %d: %w", i, err))
			continue
		}
		c.PixelMats[i] = f.PixelBytes
	}
	return errors.Join(errs...)
}