package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/binary"
	"errors"
	"fmt"
	"sync"
)

// Frame layout on the wire:
//
//	[8 bytes counter (big endian)][ciphertext + 16 byte auth tag]
//
// The counter travels in the clear because the receiver needs it to rebuild
// the nonce; it carries no secret, and the authentication tag covers it.
const (
	counterSize = 8
	nonceSize   = 12 // AES-GCM standard nonce
	tagSize     = 16
	// FrameOverhead is how much a frame grows over its plaintext.
	FrameOverhead = counterSize + tagSize
)

// Direction distinguishes the two halves of a tunnel.
//
// Both ends derive the same tunnel key, so without this they would encrypt
// different plaintexts under the same (key, nonce) pair — the one failure that
// breaks AES-GCM completely, revealing the XOR of both messages and allowing
// forgery.
type Direction uint8

const (
	// DirShareToUse marks traffic from the sharing side to the using side.
	DirShareToUse Direction = 1
	// DirUseToShare marks traffic in the opposite direction.
	DirUseToShare Direction = 2
)

// ErrCounterExhausted is returned once a sealer has used every counter value.
var ErrCounterExhausted = errors.New("tunnel counter exhausted, reconnect required")

// ErrReplay is returned for a frame whose counter was already seen.
var ErrReplay = errors.New("replayed or out-of-order tunnel frame")

// Sealer encrypts outgoing tunnel frames.
type Sealer struct {
	mu      sync.Mutex
	aead    cipher.AEAD
	dir     Direction
	counter uint64
}

// Opener decrypts incoming tunnel frames.
type Opener struct {
	mu       sync.Mutex
	aead     cipher.AEAD
	dir      Direction
	lastSeen uint64
	started  bool
}

// NewSealer creates a sealer for one direction of a tunnel.
func NewSealer(key []byte, dir Direction) (*Sealer, error) {
	aead, err := newAEAD(key)
	if err != nil {
		return nil, err
	}
	return &Sealer{aead: aead, dir: dir}, nil
}

// NewOpener creates an opener for one direction of a tunnel.
// The direction must be the one the *sender* used.
func NewOpener(key []byte, dir Direction) (*Opener, error) {
	aead, err := newAEAD(key)
	if err != nil {
		return nil, err
	}
	return &Opener{aead: aead, dir: dir}, nil
}

func newAEAD(key []byte) (cipher.AEAD, error) {
	if len(key) != keySize {
		return nil, fmt.Errorf("key is %d bytes, want %d", len(key), keySize)
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("creating cipher: %w", err)
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("creating GCM: %w", err)
	}
	return aead, nil
}

// nonceFor builds the 12-byte nonce: direction, four zero bytes, counter.
// Distinct directions therefore never share a nonce under the same key.
func nonceFor(dir Direction, counter uint64) [nonceSize]byte {
	var nonce [nonceSize]byte
	nonce[0] = byte(dir)
	binary.BigEndian.PutUint64(nonce[4:], counter)
	return nonce
}

// Seal encrypts one frame and returns it ready for transmission.
func (s *Sealer) Seal(plaintext []byte) ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.counter == ^uint64(0) {
		return nil, ErrCounterExhausted
	}
	counter := s.counter
	s.counter++

	nonce := nonceFor(s.dir, counter)

	out := make([]byte, counterSize, counterSize+len(plaintext)+tagSize)
	binary.BigEndian.PutUint64(out, counter)

	// The counter prefix is authenticated as additional data, so it cannot be
	// altered to make a frame decrypt under a different nonce.
	return s.aead.Seal(out, nonce[:], plaintext, out[:counterSize]), nil
}

// Open decrypts one frame.
//
// Frames must arrive in order, which holds for both transports in use: a
// direct TCP connection and a relayed WebSocket both preserve ordering. A
// counter that does not advance means duplication or tampering.
func (o *Opener) Open(frame []byte) ([]byte, error) {
	if len(frame) < FrameOverhead {
		return nil, fmt.Errorf("frame is %d bytes, minimum is %d", len(frame), FrameOverhead)
	}

	counter := binary.BigEndian.Uint64(frame[:counterSize])

	o.mu.Lock()
	if o.started && counter <= o.lastSeen {
		o.mu.Unlock()
		return nil, ErrReplay
	}
	o.mu.Unlock()

	nonce := nonceFor(o.dir, counter)
	plaintext, err := o.aead.Open(nil, nonce[:], frame[counterSize:], frame[:counterSize])
	if err != nil {
		return nil, fmt.Errorf("authentication failed: %w", err)
	}

	// Only advance after the frame proves authentic, so a forged frame with a
	// high counter cannot make us reject the genuine ones that follow.
	o.mu.Lock()
	if counter > o.lastSeen || !o.started {
		o.lastSeen = counter
		o.started = true
	}
	o.mu.Unlock()

	return plaintext, nil
}