// 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

// TODO(krakauer): better name

import (
	_ "bufio"
	_ "bytes"
	"crypto/rand"
	"encoding/binary"
	_ "errors"
	"fmt"
	_ "io"
	"log"
	"net"
	"time"

	"github.com/golang/protobuf/proto"
	"golang.org/x/crypto/nacl/box"
	"google.golang.org/grpc/credentials"

	"v.io/x/ref/runtime/internal/flow/conn/grpc/handshake"
	"v.io/x/ref/runtime/internal/rpc/version"
)

const (
	// Just for the sake of stubbing. When succeed is true, we pass back the same net.Conn we got in.
	// It is not encrypted or authenticated.
	succeed    = false
	defaultMtu = 1 << 16
)

// This package provides a set of methods that can be exposed to non-Vanadium RPC systems (e.g.
// gRPC) so that they can utilize Vanadium's authentication (i.e. mutual authentication) and
// authorization (i.e. blessings) capabilities.
//
// This package is explicitly aimed at GRPC for now. See package
// "google.golang.org/grpc/credentials" for more information.

// TODO(krakauer):
// - We'll ignoring blessings at first. Once mutual auth works without it, we can add blessing support.
// - We'll ignore caveats and discharges at first.
// - Use timeouts. See net package docs.
// - When is it my responsibility to close a net.Conn?

type GRPCInfo struct{}

func (GRPCInfo) AuthType() string {
	return "vanadium"
}

// 1 - Send client's public key
// 2 - Get server's public key
// 3 - Compute shared key
// 4 - Compute channel bindings
// If everything checks out at this point, consider the the connection authenticated and return a
// net.Conn that encrypts messages as they are written.
//
// Later:
// 5 - Blessing support
// 6 - Caveat and discharge support.
func ClientHandshake(addr string, rawConn net.Conn, timeout time.Duration) (net.Conn, credentials.AuthInfo, error) {
	log.Printf("In ClientHandshake.\n")
	return setupSharedSecrets(rawConn, true)
}

// 1 - Send the server's public key
// 2 - Get the client's public key
// 3 - Compute shared key
// 4 - Compute channel bindings
// If everything checks out at this point, consider the the connection authenticated and return a
// net.Conn that encrypts messages as they are written.
//
// Later:
// 5 - Blessing support
// 6 - Caveat and discharge support
func ServerHandshake(rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
	log.Printf("In ServerHandshake.\n")
	return setupSharedSecrets(rawConn, true)
}

// Generates the shared key and channel bindings for both the server and client handshakes.
func setupSharedSecrets(rawConn net.Conn, doprint bool) (net.Conn, credentials.AuthInfo, error) {
	pk, sk, err := box.GenerateKey(rand.Reader)
	setupReq := &handshake.Setup{
		// Versions are typedefed uint32s.
		MinVersion:        uint32(version.Supported.Min),
		MaxVersion:        uint32(version.Supported.Max),
		PeerNaclPublicKey: binary.LittleEndian.Uint32(pk[:]),
		Mtu:               defaultMtu,
	}
	if doprint {
		log.Printf("setupReq: %v", setupReq)
		log.Printf("pk: %v", pk)
	}
	// Now we have to send this message to the server... do we just send it over rawConn while using box?

	// if proto.Size(setupReq) < box.Overhead {
	// 	return nil, nil, errors.New("Message too short") // TODO: what's this check for? It's in box_cipher.go.
	// }

	// There's nothing to seal with with 1st time. We just send our public key as plaintext.
	// First we have to serialize the setup message.
	setupReqData, err := proto.Marshal(setupReq)
	rawConn.Write(setupReqData)

	// Wait until we have our entire reponse.
	log.Printf("Going to wait on Copy.\n")
	// resBuf := bytes.Buffer{}
	// io.Copy(&resBuf, rawConn)
	resBuf := make([]byte, 4096) // TODO better (dynamic) size or way of reading?
	bytesRead, err := rawConn.Read(resBuf)
	if err != nil {
		return nil, nil, err
	}
	log.Printf("Done waiting on Copy.\n")
	log.Printf("Received %d bytes.\n", bytesRead)

	setupRes := &handshake.Setup{}
	// err = proto.Unmarshal(resBuf.Bytes(), setupRes)
	err = proto.Unmarshal(resBuf[:bytesRead], setupRes)
	if err != nil {
		if doprint {
			log.Printf("Failed to unmarshal proto.\n")
			log.Printf("resBuf: %v", resBuf)
		}
		return nil, nil, err
	}

	// TODO: negotiate stuff with the remote (version, mtu, etc.). For now, we are using identical
	// protos, so get this working after.
	var sharedKey [32]byte
	var peerPublicKey = make([]byte, 32) // TODO does this need both?
	binary.LittleEndian.PutUint32(peerPublicKey, setupRes.PeerNaclPublicKey)
	var peerPublicKeyArr [32]byte
	copy(peerPublicKeyArr[:], peerPublicKey)
	box.Precompute(&sharedKey, &peerPublicKeyArr, sk)
	if doprint {
		log.Printf("setupRes: %v", setupRes)
		log.Printf("peerPublicKeyArr: %v", peerPublicKeyArr)
		log.Printf("sharedKey: %v\n", sharedKey)
	}

	// TODO: maybe this should use a constructor. Also, we can just read the keys directly into this.
	// Declare this at the beginning.
	secureConn := &conn{
		rawConn:   rawConn,
		publicKey: pk,
		secretKey: sk,
		sharedKey: &sharedKey,
		// binding: /* TODO */,
	}

	log.Printf("secureConn: %#v\n", secureConn)

	return secureConn, GRPCInfo{}, nil
}

func SecurityProtocol() string {
	return "vanadium"
}

// Returns the latest version supported.
func SecurityVersion() string {
	return fmt.Sprint(version.Supported.Max)
}