// Copyright 2018 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.

// +build linux

// This sample creates a stack with TCP and IPv4 protocols on top of a TUN
// device, and listens on a port. Data received by the server in the accepted
// connections is echoed back to the clients.
package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"os"
	"runtime"
	"strconv"
	"time"

	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
	"gvisor.dev/gvisor/pkg/tcpip/config"
	"gvisor.dev/gvisor/pkg/tcpip/network/arp"
	"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/icmp"
	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
	"gvisor.dev/gvisor/pkg/waiter"
 

	"gopkg.in/yaml.v3"

	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/insomniacslk/dhcp/dhcpv6"
	"github.com/insomniacslk/dhcp/dhcpv6/nclient6"
	"github.com/insomniacslk/dhcp/iana"
)

func echo(wq *waiter.Queue, ep tcpip.Endpoint) {
	defer ep.Close()

	// Create wait queue entry that notifies a channel.
	waitEntry, notifyCh := waiter.NewChannelEntry(nil)

	wq.EventRegister(&waitEntry, waiter.EventIn)
	defer wq.EventUnregister(&waitEntry)

	for {
		v, _, err := ep.Read(nil)
		if err != nil {
			if err == tcpip.ErrWouldBlock {
				<-notifyCh
				continue
			}

			return
		}

		ep.Write(tcpip.SlicePayload(v), tcpip.WriteOptions{})
	}
}

func TestFinialize(conn *net.UDPConn) {
	event := make(chan string)
	
	runtime.SetFinalizer(conn, func (obj *net.UDPConn) {
		fmt.Println("Finalize: ", obj.LocalAddr(), obj.RemoteAddr())
		event <- "Finalize"
	})

	runtime.GC()

	select {
	case res := <-event:
		fmt.Println(res)
	case <-time.After(2 * time.Second):
		fmt.Println("No finalize")
	}
}

// func (routes *Routes) OnDuplicateAddressDetectionStatus(nicID tcpip.NICID, addr tcpip.Address, resolved bool, err *tcpip.Error) {
// 	fmt.Println("OnDuplicateAddressDetectionStatus ", addr)
// }

// func (routes *Routes) OnDefaultRouterDiscovered(nicID tcpip.NICID, addr tcpip.Address) bool {
// 	fmt.Println("OnDefaultRouterDiscovered")
// 	return true
// }

// func (routes *Routes) OnDefaultRouterInvalidated(nicID tcpip.NICID, addr tcpip.Address) {
// }

// func (routes *Routes) OnOnLinkPrefixDiscovered(nicID tcpip.NICID, prefix tcpip.Subnet) bool {
// 	fmt.Println("OnOnLinkPrefixDiscovered ", prefix)
// 	return true
// }

// func (routes *Routes) OnOnLinkPrefixInvalidated(nicID tcpip.NICID, prefix tcpip.Subnet) {
// }

// func (routes *Routes) OnAutoGenAddress(tcpip.NICID, tcpip.AddressWithPrefix) bool {
// 	return true
// }

// func (routes *Routes) OnAutoGenAddressDeprecated(tcpip.NICID, tcpip.AddressWithPrefix) {
// }

// func (routes *Routes) OnAutoGenAddressInvalidated(tcpip.NICID, tcpip.AddressWithPrefix) {
// }

// func (routes *Routes) OnRecursiveDNSServerOption(nicID tcpip.NICID, addrs []tcpip.Address, lifetime time.Duration) {
// 	fmt.Println("OnRecursiveDNSServerOption ", addrs)
// }

// func (routes *Routes) OnDHCPv6Configuration(nicID tcpip.NICID, config stack.DHCPv6ConfigurationFromNDPRA) {
// 	str := "undefined"
// 	switch config {
// 	case stack.DHCPv6NoConfiguration: str = "no config"
// 	case stack.DHCPv6ManagedAddress: str = "managed"
// 	case stack.DHCPv6OtherConfigurations: str = "other"
// 	}
	
// 	fmt.Println("OnDHCPv6Configuration ", str)
// }


func protocolToString(proto tcpip.NetworkProtocolNumber) string {
	switch proto {
	case ipv4.ProtocolNumber: return "ipv4"
	case ipv6.ProtocolNumber: return "ipv6"
	case arp.ProtocolNumber: return "arp "
	default: return "?"
	}
}

func dumpAddress(addr tcpip.ProtocolAddress) {
	fmt.Println("Addr: ", protocolToString(addr.Protocol),
		addr.AddressWithPrefix)
}

func dumpAddresses(s *stack.Stack) {
	for nic, addrs := range s.AllAddresses() {
		fmt.Println("NIC:", nic)
		for _, addr := range addrs {
			dumpAddress(addr)
		}
	}
}

func dumpRoutes(s *stack.Stack) {
	for _, r := range s.GetRouteTable() {
		fmt.Println("Route:", r)
	}
}

func KeepAliveTunnel(np *config.Netplan) {
	KeepAliveTunnelEx(np, true)
}

func KeepAliveTunnelEx(np *config.Netplan, debug bool) {
	for _, tun := range np.Network.Tunnels {
		if debug {
			fmt.Println("Tunnel ", tun.Mode, tun.Local, tun.Remote, tun.Conn, tun.Sd)
		}
		runtime.KeepAlive(tun.Conn)
	}
}

func withIAPD(iaid [4]byte, prefixLength int) dhcpv6.Modifier {
	return func(d dhcpv6.DHCPv6) {
		opt := d.GetOneOption(dhcpv6.OptionIAPD)
		if opt == nil {
			opt = &dhcpv6.OptIAPD{}
		}
		iaPd := opt.(*dhcpv6.OptIAPD)

		if prefixLength > 0 {
			iaPrefix := &dhcpv6.OptIAPrefix{}
			iaPrefix.Prefix.Mask = net.CIDRMask(prefixLength,
				128-prefixLength)
			iaPrefix.Prefix.IP = net.ParseIP("::")
			iaPd.Options.Add(iaPrefix)
		}
		copy(iaPd.IaId[:], iaid[:])
		d.UpdateOption(iaPd)
	}
}

func withIAPDFromAdvertise(adv *dhcpv6.Message) dhcpv6.Modifier {
	return func(d dhcpv6.DHCPv6) {
		opt := adv.GetOneOption(dhcpv6.OptionIAPD)
		if opt != nil {
			d.AddOption(opt)
		}
	}
}

func withDHCP4oDHCP6Server(addrs ...net.IP) dhcpv6.Modifier {
	return func(d dhcpv6.DHCPv6) {
		opt := dhcpv6.OptDHCP4oDHCP6Server{
			DHCP4oDHCP6Servers: addrs,
		}
		d.UpdateOption(&opt)
	}
}

func withDHCPv4Msg(msg *dhcpv4.DHCPv4) dhcpv6.Modifier {
	return func(d dhcpv6.DHCPv6) {
		opt := dhcpv6.OptDHCPv4Msg{
			Msg: msg,
		}
		d.UpdateOption(&opt)
	}
}

func NewDHCPv4Query(flags uint32, modifiers ...dhcpv6.Modifier) (*dhcpv6.Message, error) {
	msg, err := dhcpv6.NewMessage()
	if err != nil {
		return nil, err
	}
	//msg.MessageType = dhcpv6.MessageTypeDHCPv4Query FIXME
	msg.MessageType = 20
	msg.TransactionID = dhcpv6.TransactionID{byte(flags >> 16), byte(flags >> 8), byte(flags)}
	//msg.AddOption(dhcpv6.OptElapsedTime(0))
	//modifiers = append([]dhcpv6.Modifier{dhcpv6.WithRequestedOptions(dhcpv6.OptionDHCP4oDHCP6Server)}, modifiers...)

	for _, mod := range modifiers {
		mod(msg)
	}
	return msg, nil
}

func doClient(s *stack.Stack, cfg *config.Config, nic tcpip.NICID) {
	fmt.Println("doClient start")

	// TODO use link local address
	fsm := FSM{}
	fsm.Init(DHCPv6{})
	time.Sleep(500 * time.Millisecond)
//	log.Fatal("FSM")
//	state = state(dhcp, nil)
	
	//src := tcpip.Address(net.ParseIP("2001:470:dfae:6300::1:111").To16())
	//dst := tcpip.Address(net.ParseIP("2001:470:dfae:6300::3").To16())
	srcIp := net.ParseIP("fe80::111")
	srcAddr := tcpip.Address(srcIp.To16())
	//srcUdp := net.UDPAddr{IP: srcIp, Port: 546}
	src := tcpip.FullAddress{NIC: nic, Addr: srcAddr, Port: 546}
	//dstIp := net.ParseIP("fe80::3")
	//dstUdp := net.UDPAddr{IP: dstIp, Port: 547}
	//dstAddr := tcpip.Address(dstIp.To16())
	//dst := tcpip.FullAddress{NIC: nic, Addr: dstAddr, Port: 547}
	
	conn, err := gonet.DialUDP(s,
		&src,
		//&dst,
		nil,
		ipv6.ProtocolNumber); if err != nil {
		log.Fatal(err)
	}

	hwaddr := []byte("ABCDEF")
	
	client, err := nclient6.NewWithConn(conn, hwaddr, nclient6.WithDebugLogger()); if err != nil {
		log.Fatal(err)
	}

	//nclient6.WithBroadcastAddr(&dstUdp)(client)

	duid := dhcpv6.Duid{
	 	Type:          dhcpv6.DUID_LL,
	 	HwType:        iana.HWTypeEthernet,
	 	LinkLayerAddr: hwaddr,
	}

	hostName := "gvisor"
	fqdn := hostName + ".m7n.se"
	fqdnOpt := dhcpv6.WithFQDN(0x1, fqdn)

	iaPrefix := dhcpv6.OptIAPrefix{}
	iaPrefix.Prefix.Mask = net.CIDRMask(64, 128-64)
	//iaPrefix.SetIPv6Prefix(net.ParseIP("::"))

	iaid := []byte{0, 0, 0, 3}
	ident := []byte{255}	   		 // Type IAID+DUID
	ident = append(ident, iaid...)		 // IAID
	ident = append(ident, duid.ToBytes()...) // DUID
	clientIDOpt := dhcpv4.OptClientIdentifier(ident)
	
	adv, err := client.Solicit(context.Background(),
		dhcpv6.WithIAPD([4]byte{0, 0, 0, 1}, &iaPrefix),
		dhcpv6.WithIAID([4]byte{0, 0, 0, 2}),
		fqdnOpt,
		dhcpv6.WithClientID(duid),
		dhcpv6.WithRequestedOptions(dhcpv6.OptionDHCP4oDHCP6Server),
		withDHCP4oDHCP6Server(net.ParseIP("fe80::1"),
			net.ParseIP("fe80::2"))); if err != nil {
		log.Fatal(err)
	}

	// Unicast require OPTION_UNICAST from server
	//nclient6.WithBroadcastAddr(&dstUdp)(client)

	//withIAPDFromAdvertise(adv)
	msg, err := client.Request(context.Background(), adv, fqdnOpt); if err != nil {
		log.Fatal(err)
	}

	modHostName := dhcpv4.WithGeneric(dhcpv4.OptionHostName, []byte(hostName))
	disc, err := dhcpv4.NewDiscovery(hwaddr, dhcpv4.WithOption(clientIDOpt), modHostName); if err != nil {
		log.Fatal(err)
	}
	disc_query, err := NewDHCPv4Query(0x800000, withDHCPv4Msg(disc)); if err != nil {
		log.Fatal(err)
	}

	serverAddr := nclient6.AllDHCPRelayAgentsAndServers
	disc_resp, err := client.SendAndRead(context.Background(), serverAddr, disc_query, nil); if err != nil {
		log.Fatal(err)
	}

	msgOpt := disc_resp.GetOneOption(dhcpv6.OptionDHCPv4Msg); if msgOpt == nil {
		log.Fatal("Missing DHCPv4Msg option in discovery")
	}

	offer := msgOpt.(*dhcpv6.OptDHCPv4Msg).Msg
	req, err := dhcpv4.NewRequestFromOffer(offer, modHostName); if err != nil {
		log.Fatal(err)
	}

	req_query, err := NewDHCPv4Query(0x800000, withDHCPv4Msg(req)); if err != nil {
		log.Fatal(err)
	}
	req_resp, err := client.SendAndRead(context.Background(), serverAddr, req_query, nil); if err != nil {
		log.Fatal(err)
	}

	msgOpt = req_resp.GetOneOption(dhcpv6.OptionDHCPv4Msg); if msgOpt == nil {
		log.Fatal("Missing DHCPv4Msg option in ack")
	}
	ack := msgOpt.(*dhcpv6.OptDHCPv4Msg).Msg

	//	client.Close()
	fmt.Println("doClient end", ack.YourIPAddr, ack.SubnetMask())
	ip := net.IPNet{IP: ack.YourIPAddr, Mask:ack.SubnetMask()}
	cfg.AddAddress(s, nic, ip.String())

	iana := msg.GetOneOption(dhcpv6.OptionIANA).(*dhcpv6.OptIANA)
	for _, addr := range iana.Options.Get(dhcpv6.OptionIAAddr) {
		str := addr.(*dhcpv6.OptIAAddress).IPv6Addr.String() + "/128"
		cfg.AddAddress(s, nic, str)
		cfg.AddRoute(nic, addr.(*dhcpv6.OptIAAddress).IPv6Addr.String() + "/64")
	}

	var loNic tcpip.NICID = 0
	iapd := msg.GetOneOption(dhcpv6.OptionIAPD).(*dhcpv6.OptIAPD)
	for _, opt := range iapd.Options.Get(dhcpv6.OptionIAPrefix) {
		prefix := opt.(*dhcpv6.OptIAPrefix)
		str := prefix.Prefix.IP.String()+"1"+"/128"
		cfg.AddAddress(s, loNic, str)
	}

	dumpAddresses(s)
	dumpRoutes(s)

	// if false {
	// 	// FIXME can't send non icmpv6 echo request
	// 	var wq waiter.Queue
	// 	ep, e := s.NewEndpoint(icmp.ProtocolNumber6, ipv6.ProtocolNumber, &wq)
	// 	if err != nil {
	// 		log.Fatal("NewEndpoint", e)
	// 	}

	// 	v := []byte{0, 1, 2, 3}
	// 	raSrc := tcpip.FullAddress{NIC: nic, Addr: src, Port: 0}
	// 	raDst := tcpip.FullAddress{NIC: nic, Addr: dst, Port: 0}
	// 	if err := ep.Bind(raSrc); err != nil {
	// 		log.Fatal("Bind failed: ", err)
	// 	}
	// 	fmt.Println("Before write", raSrc.NIC, raSrc.Addr, raSrc.Port, raDst.NIC, raDst.Addr, raDst.Port)
	// 	ep.Write(tcpip.SlicePayload(v),
	// 		tcpip.WriteOptions{To:&raDst})

	// 	defer ep.Close()
	// 	fmt.Println("After write")
	// }
	// Exchange runs a Solicit-Advertise-Request-Reply transaction on the
	// specified network interface, and returns a list of DHCPv6 packets
	// (a "conversation") and an error if any. Notice that Exchange may
	// return a non-empty packet list even if there is an error. This is
	// intended, because the transaction may fail at any point, and we
	// still want to know what packets were exchanged until then.
	// A default Solicit packet will be used during the "conversation",
	// which can be manipulated by using modifiers.
//	conversation, err := client.Exchange(iface)

	// Summary() prints a verbose representation of the exchanged packets.
//	for _, packet := range conversation {
//		log.Print(packet.Summary())
//	}
	// error handling is done *after* printing, so we still print the
	// exchanged packets if any, as explained above.
//	if err != nil {
//		log.Fatal(err)
//	}
}

func main() {
	flag.Parse()
	if len(flag.Args()) != 1 {
		log.Fatal("Usage: ", os.Args[0], " <local-port>")
	}

	data, err := ioutil.ReadFile("config.yaml")
	if err != nil {
		log.Fatalf("File reading error", err)
	}

	var np config.Netplan
	err = yaml.Unmarshal(data, &np)
	fmt.Println("err", err)
	fmt.Println("res", np)

	portName := flag.Arg(0)

	rand.Seed(time.Now().UnixNano())

	localPort, err := strconv.Atoi(portName)
	if err != nil {
		log.Fatalf("Unable to convert port %v: %v", portName, err)
	}

	cfg := config.Config{}

	// Create the stack with ip and tcp protocols, then add a tun-based
	// NIC and address.
	s := stack.New(stack.Options{
		NetworkProtocols:   []stack.NetworkProtocol{ipv4.NewProtocol(), ipv6.NewProtocol(), arp.NewProtocol()},
		TransportProtocols: []stack.TransportProtocol{
			tcp.NewProtocol(),
			udp.NewProtocol(),
			//icmp.NewProtocol6(),
		},
		//NDPConfigs: stack.DefaultNDPConfigurations(),
		//NDPDisp: &routes,
	})

	// FIXME enable
	cfg.Setup(s, &np)
	
	KeepAliveTunnel(&np)

	// FIXME disabled for now, to test startSolicitingRouters
	if false {
		// FIXME
		var wg2Nic tcpip.NICID = -1
		doClient(s, &cfg, wg2Nic)
	}

	cfg.SetRouteTable(s)
	
	dumpAddresses(s)
	dumpRoutes(s)

	KeepAliveTunnel(&np)

	runtime.GC()

	KeepAliveTunnel(&np)

	// Create TCP endpoint, bind it, then start listening.
	if true {
		proto := ipv6.ProtocolNumber // Dummy
		var wq waiter.Queue
		ep, e := s.NewEndpoint(tcp.ProtocolNumber, proto, &wq)
		if err != nil {
			log.Fatal("NewEndpoint", e)
		}

		defer ep.Close()

		if err := ep.Bind(tcpip.FullAddress{0, "", uint16(localPort)}); err != nil {
			log.Fatal("Bind failed: ", err)
		
}
		if err := ep.Listen(10); err != nil {
			log.Fatal("Listen failed: ", err)
		}

		// Wait for connections to appear.
		waitEntry, notifyCh := waiter.NewChannelEntry(nil)
		wq.EventRegister(&waitEntry, waiter.EventIn)
		defer wq.EventUnregister(&waitEntry)

		fmt.Println("Echo server\n")

		for {
			n, wq, err := ep.Accept()
			if err != nil {
				if err == tcpip.ErrWouldBlock {
					<-notifyCh
					continue
				}

				log.Fatal("Accept() failed:", err)
			}

			KeepAliveTunnelEx(&np, false)
			go echo(wq, n)
		}
	}

}