package wire

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

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

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

// WriteTimeout bounds a single frame write. A peer that stops reading must not
// be able to park this connection inside WriteFrame forever: the write mutex
// is held for the whole socket write, so one stalled write would otherwise
// wedge liveness (PING/PONG) and WND/FIN on this tunnel.
const WriteTimeout = 30 * time.Second

var (
	errFrameTooBig = errors.New("wire: frame exceeds max size")
	// ErrBroken is returned once a write has failed. The ChaCha20 keystream has
	// already advanced (and the socket may hold a partial frame), so the
	// connection can never be resynchronized and is closed for good.
	ErrBroken = errors.New("wire: connection is broken")
)

// 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
	broken bool
}

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.
//
// The write is bounded by WriteTimeout. On any write error the connection is
// marked broken and closed, which unblocks the reader so the owner can tear the
// session down instead of leaving every stream parked on the write mutex.
func (f *FramedConn) WriteFrame(payload []byte) error {
	f.wmu.Lock()
	defer f.wmu.Unlock()
	if f.broken {
		return ErrBroken
	}
	ct := make([]byte, len(payload))
	f.out.XORKeyStream(ct, payload)
	out := AppendVarInt(make([]byte, 0, VarIntMaxBytes+len(ct)), len(ct))
	out = append(out, ct...)
	_ = f.conn.SetWriteDeadline(time.Now().Add(WriteTimeout))
	_, err := f.conn.Write(out)
	if err != nil {
		f.broken = true
		_ = f.conn.Close()
		return err
	}
	_ = f.conn.SetWriteDeadline(time.Time{})
	return nil
}

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

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