package wire

import (
	"bufio"
	"errors"
	"io"
	"net"
	"sync"

	"golang.org/x/crypto/chacha20"
)

// MaxFrame is the maximum decrypted frame payload size (1 MiB).
const MaxFrame = 1 << 20

var errFrameTooBig = errors.New("wire: frame exceeds max size")

// FramedConn is the encrypted, length-prefixed frame transport (PROTOCOL.md §3.1).
// The VarInt length prefix is plaintext; the payload is ChaCha20-encrypted with a
// continuous per-direction keystream. Writes are serialized; reads are expected
// from a single goroutine.
type FramedConn struct {
	conn net.Conn
	r    *bufio.Reader
	in   *chacha20.Cipher
	out  *chacha20.Cipher
	wmu  sync.Mutex
}

func NewFramedConn(conn net.Conn, in, out *chacha20.Cipher) *FramedConn {
	return &FramedConn{
		conn: conn,
		r:    bufio.NewReader(conn),
		in:   in,
		out:  out,
	}
}

// SwitchCiphers swaps both ciphers at a frame boundary (Phase A → Phase B).
// Only call this from the same goroutine sequence as reads/writes during the
// handshake, before concurrency begins.
func (f *FramedConn) SwitchCiphers(in, out *chacha20.Cipher) {
	f.in = in
	f.out = out
}

// ReadFrame reads and decrypts one frame payload.
func (f *FramedConn) ReadFrame() ([]byte, error) {
	n, err := ReadVarInt(f.r)
	if err != nil {
		return nil, err
	}
	if n < 0 || n > MaxFrame {
		return nil, errFrameTooBig
	}
	ct := make([]byte, n)
	if _, err := io.ReadFull(f.r, ct); err != nil {
		return nil, err
	}
	f.in.XORKeyStream(ct, ct) // decrypt in place
	return ct, nil
}

// WriteFrame encrypts and sends one frame payload. Safe for concurrent callers.
func (f *FramedConn) WriteFrame(payload []byte) error {
	f.wmu.Lock()
	defer f.wmu.Unlock()
	ct := make([]byte, len(payload))
	f.out.XORKeyStream(ct, payload)
	out := AppendVarInt(make([]byte, 0, VarIntMaxBytes+len(ct)), len(ct))
	out = append(out, ct...)
	_, err := f.conn.Write(out)
	return err
}

func (f *FramedConn) Close() error { return f.conn.Close() }

func (f *FramedConn) RemoteAddr() net.Addr { return f.conn.RemoteAddr() }