package utils

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"git.openprivacy.ca/openprivacy/log"
	"golang.org/x/crypto/nacl/secretbox"
	"io"
	"sync"
)

const (
	// InvalidPacketLengthError is returned whenever ricochet receives a packet too small or too large to conform to the spec.
	InvalidPacketLengthError = Error("InvalidPacketLengthError")

	// InvalidChannelIDError channels must be between 0 and 65535
	InvalidChannelIDError = Error("InvalidChannelIDError")
)

// RicochetData is a structure containing the raw data and the channel it the
// message originated on.
type RicochetData struct {
	Channel int32
	Data    []byte
}

// Equals compares a RicochetData object to another and returns true if contain
// the same data.
func (rd RicochetData) Equals(other RicochetData) bool {
	return rd.Channel == other.Channel && bytes.Equal(rd.Data, other.Data)
}

// RicochetNetworkInterface abstract operations that interact with ricochet's
// packet layer.
type RicochetNetworkInterface interface {
	SendRicochetPacket(dst io.Writer, channel int32, data []byte) error
	RecvRicochetPacket(reader io.Reader) (RicochetData, error)
}

// RicochetNetwork is a concrete implementation of the RicochetNetworkInterface
type RicochetNetwork struct {
	// Derived ephemeral session key for connection
	key     [32]byte
	encrypt bool
	lock    sync.Mutex
}

// SetEncryptionKey sets the ephemeral encryption key for this session.
func (rn *RicochetNetwork) SetEncryptionKey(key [32]byte) {
	rn.lock.Lock()
	defer rn.lock.Unlock()
	log.Debugf("turning on ephemeral session encryption for connection")
	copy(rn.key[:], key[:])

	rn.encrypt = true
}

// SendRicochetPacket places the data into a structure needed for the client to
// decode the packet and writes the packet to the network.
func (rn *RicochetNetwork) SendRicochetPacket(dst io.Writer, channel int32, data []byte) error {
	packet := make([]byte, 4+len(data))
	if len(packet) > 65535 {
		return InvalidPacketLengthError
	}
	binary.BigEndian.PutUint16(packet[0:2], uint16(len(packet)))
	if channel < 0 || channel > 65535 {
		return InvalidChannelIDError
	}
	binary.BigEndian.PutUint16(packet[2:4], uint16(channel))
	copy(packet[4:], data[:])

	rn.lock.Lock()
	if rn.encrypt {
		var nonce [24]byte
		if _, err := io.ReadFull(rand.Reader, nonce[:]); err != nil {
			panic(err)

		}
		encrypted := secretbox.Seal(nonce[:], packet[2:], &nonce, &rn.key)
		binary.BigEndian.PutUint16(packet[0:2], uint16(len(encrypted)+2))
		packet = append(packet[0:2], encrypted...)
	}
	rn.lock.Unlock()

	for pos := 0; pos < len(packet); {
		n, err := dst.Write(packet[pos:])
		if err != nil {
			return err
		}
		pos += n
	}

	return nil
}

// RecvRicochetPacket returns the next packet from reader as a RicochetData
// structure, or an error.
func (rn *RicochetNetwork) RecvRicochetPacket(reader io.Reader) (RicochetData, error) {

	packet := RicochetData{}

	// Read the four-byte header to get packet length
	header := make([]byte, 2)
	if _, err := io.ReadAtLeast(reader, header, len(header)); err != nil {
		return packet, err
	}

	size := int(binary.BigEndian.Uint16(header[0:2]))
	if size < 4 {
		return packet, InvalidPacketLengthError
	}

	packetBytes := make([]byte, size-2)
	_, err := io.ReadAtLeast(reader, packetBytes, size-2)
	if err != nil {
		return packet, err
	}

	rn.lock.Lock()
	if rn.encrypt {
		var decryptNonce [24]byte
		if len(packetBytes) > 24 {
			copy(decryptNonce[:], packetBytes[:24])
			decrypted, ok := secretbox.Open(nil, packetBytes[24:], &decryptNonce, &rn.key)
			if ok {
				packetBytes = decrypted
			} else {
				return packet, errors.New("failed to decrypt encrypted ricochet packet")
			}

		} else {
			return packet, errors.New("ciphertext length was too short")
		}
	}
	rn.lock.Unlock()

	packet.Channel = int32(binary.BigEndian.Uint16(packetBytes[0:2]))
	packet.Data = make([]byte, len(packetBytes)-2)
	copy(packet.Data[:], packetBytes[2:])
	return packet, nil
}