// Copyright 2021 The gVisor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tunnel

import (
	"context"
	"errors"
	"log"

	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/buffer"
	"gvisor.dev/gvisor/pkg/tcpip/header"
	"gvisor.dev/gvisor/pkg/tcpip/link/channel"
	"gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
	"gvisor.dev/gvisor/pkg/tcpip/network/ipv6"
	"gvisor.dev/gvisor/pkg/tcpip/stack"
	"gvisor.dev/gvisor/pkg/tcpip/transport/gre"
	"gvisor.dev/gvisor/pkg/tcpip/transport/raw"
	"gvisor.dev/gvisor/pkg/waiter"
)

var SIZE = 16

type Endpoint struct {
	channel.Endpoint
	LocalAddr tcpip.Address
	RemoteAddr tcpip.Address
}

type writer Endpoint

func New(mtu uint32) *Endpoint {
	var linkAddress tcpip.LinkAddress

	ch := channel.New(SIZE, mtu, linkAddress)
	return &Endpoint{
		Endpoint: *ch,
	}
}

func (e *Endpoint) GetChannel() *channel.Endpoint {
	return &e.Endpoint
}

func (e *Endpoint) GetLocalAddress() tcpip.Address {
	return e.LocalAddr
}

func (e *Endpoint) GetRemoteAddress() tcpip.Address {
	return e.RemoteAddr
}

// Attach saves the stack network-layer dispatcher for use later when packets
// are injected.
func (e *Endpoint) Attach(dispatcher stack.NetworkDispatcher) {
	log.Println("GRE: Attach")

	e.Endpoint.Attach(dispatcher)
}

type GrePacketInfo struct {
	pi *channel.PacketInfo
}

func (g GrePacketInfo) Len() int {
	return g.pi.Pkt.Size()
}

func (g GrePacketInfo) Read(p []byte) (n int, err error){
	pkt := g.pi.Pkt
	size := pkt.Size()

	if size > len(p) {
		return 0, errors.New("Buffer too small")
	}
	vv := buffer.NewVectorisedView(size, pkt.Views())
	pos := 0
	for {
		copied, err := vv.Read(p[pos:])
		log.Printf("VectorisedView Read: %d %d %d %v", size, pos, copied, err)
		if err != nil {
			return 0, errors.New("&tcpip.ErrBadBuffer{}")
		}
		pos = pos + copied
		if pos == size {
			break
		}
	}

	return pos, nil
}

func (g GrePacketInfo) Payload(size int) ([]byte, tcpip.Error){
	log.Println("Payload")
	return nil, &tcpip.ErrNotSupported{}
}

type GreHandlerInfo struct {
	ChEp *channel.Endpoint
	Raw tcpip.Endpoint
	Raddr tcpip.FullAddress
	Stack *stack.Stack
	NIC tcpip.NICID
}

func (info *GreHandlerInfo) greHandler(req *gre.ForwarderRequest) {
	pkt := req.Pkt
	log.Println("greHandler: ", req.Pkt.Size(), req.Pkt.Views())
	greHdr := header.GRE(pkt.TransportHeader().View())
	proto := greHdr.ProtocolType()
	pktData := pkt.Data()
	views := pktData.Views()
	size := pktData.Size()
	data := buffer.NewVectorisedView(size, views)
	newPkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
		Data: data,
	})
	log.Printf("greHandler proto: %d cloned: %v", proto, newPkt.Views())

	info.ChEp.InjectInbound(proto, newPkt)
}

func ipAny(protocol tcpip.NetworkProtocolNumber) (tcpip.Address, tcpip.Error) {
	switch protocol {
	case header.IPv4ProtocolNumber:
		return header.IPv4Any, nil
	case header.IPv6ProtocolNumber:
		return header.IPv6Any, nil
	default:
		log.Printf("Unsupported protocol %v", protocol)
		return tcpip.Address([]byte{}), &tcpip.ErrNotSupported{}
	}
}

func (info *GreHandlerInfo) greRead(ep *channel.Endpoint) {
	for {
		pi, err := ep.ReadContext(context.Background());
		linkHdr := pi.Pkt.LinkHeader()
		greHdr := header.GRE(linkHdr.Push(header.GREHeaderSize))
		greHdr.SetProtocolType(pi.Proto)
		log.Printf("greRead %d %v %v %v", pi.Proto, pi, err, greHdr)
		opts := tcpip.WriteOptions{
			//To: &info.Raddr
		}
		if info.Raddr.Addr == header.IPv4Any {
			protocol := pi.Pkt.NetworkProtocolNumber
			nw := pi.Pkt.Network()
			addr := nw.DestinationAddress()
			laddr, err := ipAny(protocol); if err != nil {
				continue
			}

			err = info.Stack.GetLinkAddress(info.NIC, addr, laddr,  protocol, func (res stack.LinkResolutionResult) {

				if res.Err != nil {
					return
				}

				opts.To = &tcpip.FullAddress{
					NIC: 0,
					Addr: tcpip.Address(res.LinkAddress),
				}
			})
			if err != nil {
				continue
			}
		}
		info.Raw.Write(GrePacketInfo{&pi}, opts)
	}
}

func networkProtocolNumber(addr *tcpip.Address) tcpip.NetworkProtocolNumber {
	if addr.To4() != "" {
		return ipv4.ProtocolNumber
	} else {
		return ipv6.ProtocolNumber
	}
}

// ARPHardwareType implements stack.LinkEndpoint.ARPHardwareType.
func (*Endpoint) ARPHardwareType() header.ARPHardwareType {
	return header.ARPHardwareIPGRE
}

func (e *Endpoint) Start(s *stack.Stack, nic tcpip.NICID, laddr, raddr *tcpip.Address) tcpip.Error {
	proto := networkProtocolNumber(laddr)
	if proto != networkProtocolNumber(raddr) {
		return &tcpip.ErrBadAddress{}
	}

	// Create TCP endpoint.
	var rawWq waiter.Queue
	rawEp, tcperr := raw.NewEndpoint(s, proto, header.GREProtocolNumber, &rawWq)
	if tcperr != nil {
		return tcperr
	}
	log.Println("EP: %s", rawEp)

	e.LocalAddr = *laddr
	e.RemoteAddr = *raddr

	fraddr := tcpip.FullAddress{NIC: 0, Addr: *raddr}
	flAddr := tcpip.FullAddress{NIC: 0, Addr: *laddr}
	tcperr = rawEp.Bind(flAddr)
	if tcperr != nil {
		return tcperr
	}
	log.Printf("Remote: %v %v", raddr, fraddr)

	// Create GRE
	// greEP := grelink.New(mtu - 24)
	chEP := e.GetChannel()
	// TODO detect IPv4/IPv6
	e.NMaxHeaderLength = header.IPv6FixedHeaderSize + header.GREHeaderSize

	greInfo := GreHandlerInfo{
//		Ep: loEP,
		ChEp: chEP,
		Raw: rawEp,
		Raddr: fraddr,
		Stack: s,
		NIC: nic,
	}
	greFwd := gre.NewForwarder(s, greInfo.greHandler)
	s.SetTransportProtocolHandler(header.GREProtocolNumber, greFwd.HandlePacket)

	go greInfo.greRead(chEP)

	any, _ := ipAny(proto)
	if fraddr.Addr != any {
		tcperr = rawEp.Connect(fraddr)
		if tcperr != nil {
			return tcperr
		}
	}

	return nil
}