// Copyright 2016 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package grpc

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	_ "runtime/debug"
	"sync"
	"time"

	"golang.org/x/crypto/nacl/box"
)

// Implements net.Conn. Encrypts messages using keys negotioated via NaCl.
// TODO: blessings
// TODO: discharges and caveats
// TODO: ensuer that all 5 of these are set. Constructor?
// TODO: do I need to check the message size as in box_coipher.Seal?
// TODO: make all these keys not pointers?
// TODO: should I have locks here?
type conn struct {
	rawConn   net.Conn
	publicKey *[32]byte
	secretKey *[32]byte
	sharedKey *[32]byte
	binding   []byte
	nonce     [24]byte // TODO: do I need a second nonce? Does this need to be random?
	counter   uint64   // Why is this so weirdly handled in advanceNonce?
	mu        sync.Mutex
}

func (c *conn) Read(b []byte) (n int, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	log.Printf("Beginning to Read.\n")
	resBuf := make([]byte, 4096*4096) // TODO better (dynamic) size or way of reading?
	bytesRead, err := c.rawConn.Read(resBuf)
	if err != nil {
		log.Printf("Failed to read.\n")
		log.Fatal(err)
		return bytesRead, err
	}

	// tmp := make([]byte, 0, bytesRead-box.Overhead) // TODO: is this enough? Also, why do we need both of tmp and out?
	out, ok := box.OpenAfterPrecomputation(nil, resBuf[:bytesRead], c.currentNonce(), c.sharedKey)
	c.advanceNonce() // TODO: defer? Is there harm in advancing often than necessary (i.e. on error)?
	if !ok {
		log.Printf("Failed to decrypt.\n")
		//log.Fatal(err)
		return 0, err
	}
	copy(b, out)
	// should we return bytesRead or len(b)?
	log.Printf("Succeeded in reading.\n")
	return bytesRead, nil
}

// TODO: all this casting is gross
func (c *conn) Write(b []byte) (n int, err error) {
	// log.Printf("Write called:\n%v\n", string(debug.Stack()))
	c.mu.Lock()
	defer c.mu.Unlock()
	log.Printf("Beginning to write.\n")
	// tmp := make([]byte, 0, len(b)+box.Overhead) // TODO: is this enough? Also, why do we need both of tmp and out?
	out := box.SealAfterPrecomputation(nil, b, c.currentNonce(), c.sharedKey)
	c.advanceNonce()
	bytesCopied, err := io.Copy(c.rawConn, bytes.NewReader(out))
	if err != nil {
		log.Printf("Failed to copy to rawConn.\n")
		log.Fatal(err)
		return int(bytesCopied), err
	}
	if bytesCopied != int64(len(out)) {
		errMsg := fmt.Sprintf("Did not write entire message. Expected to write %d bytes but wrote %d.", len(out), bytesCopied)
		log.Printf("Did not write entire message. Expected to write %d bytes but wrote %d.", len(out), bytesCopied)
		log.Printf("Failed to write.\n")
		log.Fatal(err)
		return -1, errors.New(errMsg)
	}
	log.Printf("Succeeded in writing!\n")
	return int(bytesCopied), nil
}

// TODO: cover these up with an interface?
// TODO: remove all these useless mutex locks and unlocks
func (c *conn) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.rawConn.Close()
}

func (c *conn) LocalAddr() net.Addr {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.rawConn.LocalAddr()
}

func (c *conn) RemoteAddr() net.Addr {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.rawConn.RemoteAddr()
}

func (c *conn) SetDeadline(t time.Time) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.rawConn.SetDeadline(t)
}

func (c *conn) SetReadDeadline(t time.Time) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.rawConn.SetReadDeadline(t)
}

func (c *conn) SetWriteDeadline(t time.Time) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.rawConn.SetWriteDeadline(t)
}

// TODO: should I just collapse current and advance?
func (c *conn) currentNonce() *[24]byte {
	return &c.nonce
}

func (c *conn) advanceNonce() {
	c.counter++
	binary.LittleEndian.PutUint64(c.nonce[:], c.counter)
}