package lafsession

import (
	"crypto/cipher"
	"encoding/binary"
	"sync"
)

// Config configures a Session before handshake.
type Config struct {
	// Identity is this side's long-term key. Required.
	Identity *Identity

	// Role is RoleServer or RoleClient.
	Role HandshakeRole

	// AcceptPeer is called with the peer's 32-byte identity public key during
	// handshake. Return false to reject (ErrIdentityRejected).
	// If nil, any peer identity is accepted (not recommended for production).
	AcceptPeer func(peerPub []byte) bool

	// ReplayWindowSize is the anti-replay window in packets (default 1024).
	ReplayWindowSize uint64
}

// Session is a mutually authenticated AEAD channel that carries opaque LAF payloads.
//
// Typical server flow:
//
//	s, _ := New(Config{Identity: serverID, Role: RoleServer, AcceptPeer: pinCheck})
//	// read clientHello from network
//	serverHello, _ := s.HandleHandshake(clientHello)
//	// write serverHello
//	// read clientFinished
//	_, _ = s.HandleHandshake(clientFinished)
//	// optional: write serverFinished
//	sf, _ := s.SendFinished()
//	// write sf
//	// now Established()
//	frame, _ := s.Seal(lafBytes)
//	laf, _ := s.Open(frame)
type Session struct {
	mu sync.Mutex

	cfg Config
	hs  *handshakeState

	keys     *trafficKeys
	sendAEAD cipher.AEAD
	recvAEAD cipher.AEAD

	sendSeq uint64
	recvWin *replayWindow

	sealBuf []byte
	openBuf []byte
	nonce   [nonceLen]byte
	aad     [recordHeaderLen]byte

	established           bool
	closed                bool
	pendingServerFinished bool
}

// New creates a Session in pre-handshake state.
func New(cfg Config) (*Session, error) {
	if cfg.Identity == nil {
		return nil, ErrIdentityRejected
	}
	if cfg.ReplayWindowSize == 0 {
		cfg.ReplayWindowSize = 1024
	}
	hs, err := newHandshake(cfg.Role, cfg.Identity, cfg.AcceptPeer)
	if err != nil {
		return nil, err
	}
	return &Session{
		cfg:     cfg,
		hs:      hs,
		recvWin: newReplayWindow(cfg.ReplayWindowSize),
	}, nil
}

// StartHandshake is called by the client to produce ClientHello.
func (s *Session) StartHandshake() ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, ErrClosed
	}
	if s.established {
		return nil, ErrAlreadyEstablished
	}
	if s.cfg.Role != RoleClient {
		return nil, ErrHandshakeState
	}
	return s.hs.buildClientHello()
}

// HandleHandshake processes one inbound handshake message and returns an
// optional outbound response (may be nil).
//
// Server: ClientHello → ServerHello; ClientFinished → (nil, establishes keys).
// Client: ServerHello → ClientFinished; ServerFinished → (nil, establishes keys).
func (s *Session) HandleHandshake(msg []byte) (response []byte, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, ErrClosed
	}
	if s.established {
		return nil, ErrAlreadyEstablished
	}
	if len(msg) < 1 {
		return nil, ErrShortMessage
	}

	switch s.cfg.Role {
	case RoleServer:
		return s.handleServer(msg)
	case RoleClient:
		return s.handleClient(msg)
	default:
		return nil, ErrHandshakeState
	}
}

func (s *Session) handleServer(msg []byte) ([]byte, error) {
	switch msg[0] {
	case msgClientHello:
		if err := s.hs.processClientHello(msg); err != nil {
			return nil, err
		}
		return s.hs.buildServerHello()
	case msgClientFinished:
		if err := s.hs.processFinished(msg, msgClientFinished, s.hs.peerIDSlice()); err != nil {
			return nil, err
		}
		keys, err := s.hs.deriveKeys()
		if err != nil {
			return nil, err
		}
		if err := s.installKeys(keys); err != nil {
			return nil, err
		}
		s.pendingServerFinished = true
		s.established = true
		return nil, nil
	default:
		return nil, ErrHandshakeState
	}
}

func (s *Session) handleClient(msg []byte) ([]byte, error) {
	switch msg[0] {
	case msgServerHello:
		if err := s.hs.processServerHello(msg); err != nil {
			return nil, err
		}
		cf, err := s.hs.buildFinished(msgClientFinished)
		if err != nil {
			return nil, err
		}
		// Keys from transcript through ClientFinished — same point the server uses.
		keys, err := s.hs.deriveKeys()
		if err != nil {
			return nil, err
		}
		if err := s.installKeys(keys); err != nil {
			return nil, err
		}
		// Established only after ServerFinished confirms the peer finished.
		return cf, nil
	case msgServerFinished:
		if err := s.hs.processFinished(msg, msgServerFinished, s.hs.peerIDSlice()); err != nil {
			return nil, err
		}
		if s.keys == nil {
			return nil, ErrHandshakeState
		}
		s.established = true
		return nil, nil
	default:
		return nil, ErrHandshakeState
	}
}

// SendFinished returns ServerFinished (server only, after ClientFinished processed).
func (s *Session) SendFinished() ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, ErrClosed
	}
	if s.cfg.Role != RoleServer || !s.pendingServerFinished {
		return nil, ErrHandshakeState
	}
	msg, err := s.hs.buildFinished(msgServerFinished)
	if err != nil {
		return nil, err
	}
	s.pendingServerFinished = false
	return msg, nil
}

func (s *Session) installKeys(keys *trafficKeys) error {
	var sendKey, recvKey []byte
	if s.cfg.Role == RoleClient {
		sendKey, recvKey = keys.c2s[:], keys.s2c[:]
	} else {
		sendKey, recvKey = keys.s2c[:], keys.c2s[:]
	}
	sa, err := newAEAD(sendKey)
	if err != nil {
		return err
	}
	ra, err := newAEAD(recvKey)
	if err != nil {
		return err
	}
	if s.keys != nil {
		s.keys.wipe()
	}
	s.keys = keys
	s.sendAEAD = sa
	s.recvAEAD = ra
	s.sendSeq = 0
	s.recvWin = newReplayWindow(s.cfg.ReplayWindowSize)
	return nil
}

// Established reports whether application Seal/Open are allowed.
func (s *Session) Established() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.established && !s.closed
}

// PeerPublicKey returns the peer identity public key after it has been received.
func (s *Session) PeerPublicKey() []byte {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.hs == nil || !s.hs.hasPeer {
		return nil
	}
	out := make([]byte, idPubLen)
	copy(out, s.hs.peerIDPub[:])
	return out
}

func (s *Session) Seal(lafPayload []byte) ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, ErrClosed
	}
	if !s.established || s.sendAEAD == nil {
		return nil, ErrNotEstablished
	}
	seq := s.sendSeq
	s.sendSeq++
	s.sealBuf = sealRecord(s.sendAEAD, seq, recordTypeData, lafPayload, s.sealBuf, &s.nonce, &s.aad)
	return s.sealBuf, nil
}

func (s *Session) SealTo(dst, lafPayload []byte) ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, ErrClosed
	}
	if !s.established || s.sendAEAD == nil {
		return nil, ErrNotEstablished
	}
	seq := s.sendSeq
	s.sendSeq++
	return sealRecord(s.sendAEAD, seq, recordTypeData, lafPayload, dst, &s.nonce, &s.aad), nil
}

// Open decrypts and verifies a record frame.
func (s *Session) Open(frame []byte) ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, ErrClosed
	}
	if !s.established || s.recvAEAD == nil {
		return nil, ErrNotEstablished
	}
	if len(frame) < recordHeaderLen+tagLen {
		return nil, ErrShortMessage
	}
	seq := binary.BigEndian.Uint64(frame[1:9])
	if !s.recvWin.accept(seq) {
		return nil, ErrReplay
	}
	recType, pt, err := openRecord(s.recvAEAD, frame, s.openBuf, &s.nonce, &s.aad)
	if err != nil {
		return nil, err
	}
	s.openBuf = pt
	switch recType {
	case recordTypeData:
		return pt, nil
	case recordTypeClose:
		s.closed = true
		s.wipeKeys()
		return nil, ErrClosed
	default:
		return nil, ErrBadTag
	}
}

func (s *Session) OpenTo(dst, frame []byte) ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, ErrClosed
	}
	if !s.established || s.recvAEAD == nil {
		return nil, ErrNotEstablished
	}
	if len(frame) < recordHeaderLen+tagLen {
		return nil, ErrShortMessage
	}
	seq := binary.BigEndian.Uint64(frame[1:9])
	if !s.recvWin.accept(seq) {
		return nil, ErrReplay
	}
	recType, pt, err := openRecord(s.recvAEAD, frame, dst, &s.nonce, &s.aad)
	if err != nil {
		return nil, err
	}
	switch recType {
	case recordTypeData:
		return pt, nil
	case recordTypeClose:
		s.closed = true
		s.wipeKeys()
		return nil, ErrClosed
	default:
		return nil, ErrBadTag
	}
}

// Close sends an authenticated close record (if established) and wipes keys.
// Write the returned frame to the peer when non-nil.
func (s *Session) Close() (frame []byte, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil, nil
	}
	if s.established && s.sendAEAD != nil {
		seq := s.sendSeq
		s.sendSeq++
		s.sealBuf = sealRecord(s.sendAEAD, seq, recordTypeClose, nil, s.sealBuf, &s.nonce, &s.aad)
		frame = s.sealBuf
	}
	s.closed = true
	s.established = false
	s.wipeKeys()
	return frame, err
}

func (s *Session) wipeKeys() {
	if s.keys != nil {
		s.keys.wipe()
		s.keys = nil
	}
	s.sendAEAD = nil
	s.recvAEAD = nil
}