// +build linux

package main

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"syscall"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/oklog/run"
	"github.com/spf13/cobra"
	"github.com/vishvananda/netlink"
	"golang.org/x/sys/unix"
	"golang.zx2c4.com/wireguard/wgctrl"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/squat/kilo/pkg/iproute"
	"github.com/squat/kilo/pkg/k8s"
	"github.com/squat/kilo/pkg/k8s/apis/kilo/v1alpha1"
	kiloclient "github.com/squat/kilo/pkg/k8s/clientset/versioned"
	"github.com/squat/kilo/pkg/mesh"
	"github.com/squat/kilo/pkg/route"
	"github.com/squat/kilo/pkg/wireguard"
)

func takeIPNet(_ net.IP, i *net.IPNet, _ error) *net.IPNet {
	return i
}

func connect() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "connect",
		RunE:  connectAsPeer,
		Short: "connect to a Kilo cluster as a peer over WireGuard",
	}
	cmd.Flags().IPNetP("allowed-ip", "a", *takeIPNet(net.ParseCIDR("10.10.10.10/32")), "Allowed IP of the peer")
	cmd.Flags().IPNetP("service-cidr", "c", *takeIPNet(net.ParseCIDR("10.43.0.0/16")), "service CIDR of the cluster")
	cmd.Flags().String("log-level", logLevelInfo, fmt.Sprintf("Log level to use. Possible values: %s", availableLogLevels))
	cmd.Flags().String("config-path", "/tmp/wg.ini", "path to WireGuard configuation file")
	cmd.Flags().Bool("clean-up", true, "clean up routes and interface")
	cmd.Flags().Uint("mtu", uint(1420), "clean up routes and interface")
	cmd.Flags().Duration("resync-period", 30*time.Second, "How often should Kilo reconcile?")

	availableLogLevels = strings.Join([]string{
		logLevelAll,
		logLevelDebug,
		logLevelInfo,
		logLevelWarn,
		logLevelError,
		logLevelNone,
	}, ", ")

	return cmd
}

func connectAsPeer(cmd *cobra.Command, args []string) error {
	var cancel context.CancelFunc
	resyncPersiod, err := cmd.Flags().GetDuration("resync-period")
	if err != nil {
		return err
	}
	mtu, err := cmd.Flags().GetUint("mtu")
	if err != nil {
		return err
	}
	configPath, err := cmd.Flags().GetString("config-path")
	if err != nil {
		return err
	}
	serviceCIDR, err := cmd.Flags().GetIPNet("service-cidr")
	if err != nil {
		return err
	}
	allowedIP, err := cmd.Flags().GetIPNet("allowed-ip")
	if err != nil {
		return err
	}
	logger := log.NewJSONLogger(log.NewSyncWriter(os.Stdout))
	logLevel, err := cmd.Flags().GetString("log-level")
	if err != nil {
		return err
	}
	switch logLevel {
	case logLevelAll:
		logger = level.NewFilter(logger, level.AllowAll())
	case logLevelDebug:
		logger = level.NewFilter(logger, level.AllowDebug())
	case logLevelInfo:
		logger = level.NewFilter(logger, level.AllowInfo())
	case logLevelWarn:
		logger = level.NewFilter(logger, level.AllowWarn())
	case logLevelError:
		logger = level.NewFilter(logger, level.AllowError())
	case logLevelNone:
		logger = level.NewFilter(logger, level.AllowNone())
	default:
		return fmt.Errorf("log level %s unknown; possible values are: %s", logLevel, availableLogLevels)
	}
	logger = log.With(logger, "ts", log.DefaultTimestampUTC)
	logger = log.With(logger, "caller", log.DefaultCaller)
	peername := "random"
	if len(args) > 0 {
		peername = args[0]
	}

	var kiloClient *kiloclient.Clientset
	switch backend {
	case k8s.Backend:
		config, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
		if err != nil {
			return fmt.Errorf("failed to create Kubernetes config: %v", err)
		}
		kiloClient = kiloclient.NewForConfigOrDie(config)
	default:
		return fmt.Errorf("backend %v unknown; posible values are: %s", backend, availableBackends)
	}
	privateKey, err := wgtypes.GeneratePrivateKey()
	if err != nil {
		return fmt.Errorf("failed to generate private key: %w", err)
	}
	publicKey := privateKey.PublicKey()
	level.Info(logger).Log("msg", "generated public key", "key", publicKey)

	peer := &v1alpha1.Peer{
		ObjectMeta: metav1.ObjectMeta{
			Name: peername,
		},
		Spec: v1alpha1.PeerSpec{
			AllowedIPs:          []string{allowedIP.String()},
			PersistentKeepalive: 10,
			PublicKey:           publicKey.String(),
		},
	}
	if p, err := kiloClient.KiloV1alpha1().Peers().Get(context.TODO(), peername, metav1.GetOptions{}); err != nil || p == nil {
		peer, err = kiloClient.KiloV1alpha1().Peers().Create(context.TODO(), peer, metav1.CreateOptions{})
		if err != nil {
			return fmt.Errorf("failed to create peer: %w", err)
		}
	}

	kiloIfaceName := "kilo0"

	iface, _, err := wireguard.New(kiloIfaceName, mtu)
	if err != nil {
		return fmt.Errorf("failed to create wg interface: %w", err)
	} else {
		level.Info(logger).Log("msg", "successfully created wg interface", "name", kiloIfaceName, "no", iface)
	}
	if err := iproute.Set(iface, false); err != nil {
		return err
	}

	if err := iproute.SetAddress(iface, &allowedIP); err != nil {
		return err
	} else {
		level.Info(logger).Log("mag", "successfully set IP address of wg interface", "IP", allowedIP.String())
	}

	if err := iproute.Set(iface, true); err != nil {
		return err
	}

	var g run.Group
	ctx, cancel := context.WithCancel(context.Background())
	g.Add(run.SignalHandler(ctx, syscall.SIGINT, syscall.SIGTERM))

	table := route.NewTable()
	stop := make(chan struct{}, 1)
	errCh := make(<-chan error, 1)
	{
		ch := make(chan struct{})
		g.Add(
			func() error {
				for {
					select {
					case err, ok := <-errCh:
						if ok {
							level.Error(logger).Log("err", err.Error())
						} else {
							return nil
						}
					case <-ch:
						return nil
					}
				}
			},
			func(err error) {
				ch <- struct{}{}
				close(ch)
				stop <- struct{}{}
				close(stop)
				level.Error(logger).Log("msg", "stopped ip routes table", "err", err.Error())
			},
		)
	}
	{
		ch := make(chan struct{})
		g.Add(
			func() error {
				for {
					ns, err := opts.backend.Nodes().List()
					if err != nil {
						return fmt.Errorf("failed to list nodes: %v", err)
					}
					ps, err := opts.backend.Peers().List()
					if err != nil {
						return fmt.Errorf("failed to list peers: %v", err)
					}
					// Obtain the Granularity by looking at the annotation of the first node.
					if opts.granularity, err = optainGranularity(opts.granularity, ns); err != nil {
						return fmt.Errorf("failed to obtain granularity: %w", err)
					}
					var hostname string
					subnet := mesh.DefaultKiloSubnet
					nodes := make(map[string]*mesh.Node)
					for _, n := range ns {
						if n.Ready() {
							nodes[n.Name] = n
							hostname = n.Name
						}
						if n.WireGuardIP != nil {
							subnet = n.WireGuardIP
						}
					}
					subnet.IP = subnet.IP.Mask(subnet.Mask)
					if len(nodes) == 0 {
						return errors.New("did not find any valid Kilo nodes in the cluster")
					}
					peers := make(map[string]*mesh.Peer)
					for _, p := range ps {
						if p.Ready() {
							peers[p.Name] = p
						}
					}
					if _, ok := peers[peername]; !ok {
						return fmt.Errorf("did not find any peer named %q in the cluster", peername)
					}

					t, err := mesh.NewTopology(nodes, peers, opts.granularity, hostname, opts.port, []byte{}, subnet, peers[peername].PersistentKeepalive, logger)
					if err != nil {
						return fmt.Errorf("failed to create topology: %v", err)
					}
					conf := t.PeerConf(peername)
					conf.Interface = &wireguard.Interface{
						PrivateKey: []byte(privateKey.String()),
						ListenPort: uint32(55555),
					}
					buf, err := conf.Bytes()
					if err != nil {
						//level.Error(m.logger).Log("error", err)
						//m.errorCounter.WithLabelValues("apply").Inc()
						return err
					}
					if err := ioutil.WriteFile("/tmp/wg.ini", buf, 0600); err != nil {
						//level.Error(m.logger).Log("error", err)
						//m.errorCounter.WithLabelValues("apply").Inc()
						return err
					}
					if err := wireguard.SetConf(kiloIfaceName, "/tmp/wg.ini"); err != nil {
						return err
					}
					wgClient, err := wgctrl.New()
					if err != nil {
						return fmt.Errorf("failed to initialize wg Client: %w", err)
					}
					defer func() {
						wgClient.Close()
					}()
					wgConf := wgtypes.Config{
						PrivateKey: &privateKey,
						//	Peers:      []wgtypes.PeerConfig{},
					}
					if err := wgClient.ConfigureDevice(kiloIfaceName, wgConf); err != nil {
						return fmt.Errorf("failed to configure wg interface: %w", err)
					}

					var routes []*netlink.Route
					for _, segment := range t.Segments() {
						for i := range segment.CIDRS() {
							// Add routes to the Pod CIDRs of nodes in other segments.
							routes = append(routes, &netlink.Route{
								Dst: segment.CIDRS()[i],
								//Flags:     int(netlink.FLAG_ONLINK),
								LinkIndex: iface,
								Protocol:  unix.RTPROT_STATIC,
							})
						}
						for i := range segment.PrivateIPs() {
							// Add routes to the private IPs of nodes in other segments.
							routes = append(routes, &netlink.Route{
								Dst: mesh.OneAddressCIDR(segment.PrivateIPs()[i]),
								//Flags:     int(netlink.FLAG_ONLINK),
								LinkIndex: iface,
								Protocol:  unix.RTPROT_STATIC,
							})
						}
						// For segments / locations other than the location of this instance of kg,
						// we need to set routes for allowed location IPs over the leader in the current location.
						for i := range segment.AllowedLocationIPs() {
							routes = append(routes, &netlink.Route{
								Dst: segment.AllowedLocationIPs()[i],
								//Flags:     int(netlink.FLAG_ONLINK),
								LinkIndex: iface,
								Protocol:  unix.RTPROT_STATIC,
							})
						}
						routes = append(routes, &netlink.Route{
							Dst: mesh.OneAddressCIDR(segment.WireGuardIP()),
							//Flags:     int(netlink.FLAG_ONLINK),
							LinkIndex: iface,
							Protocol:  unix.RTPROT_STATIC,
						})
					}
					// Add routes for the allowed IPs of peers.
					for _, peer := range t.Peers() {
						for i := range peer.AllowedIPs {
							routes = append(routes, &netlink.Route{
								Dst: peer.AllowedIPs[i],
								//Flags:     int(netlink.FLAG_ONLINK),
								LinkIndex: iface,
								Protocol:  unix.RTPROT_STATIC,
							})
						}
					}
					routes = append(routes, &netlink.Route{
						Dst: &serviceCIDR,
						//Flags:     int(netlink.FLAG_ONLINK),
						Gw:        nil,
						LinkIndex: iface,
						Protocol:  unix.RTPROT_STATIC,
					})
					for _, r := range routes {
						fmt.Println(r)
					}

					level.Debug(logger).Log("routes", routes)
					if err := table.Set(routes, []*netlink.Rule{}); err != nil {
						return fmt.Errorf("failed to set ip routes table: %w", err)
					}
					errCh, err = table.Run(stop)
					if err != nil {
						return fmt.Errorf("failed to start ip routes tables: %w", err)
					}
					select {
					case <-time.After(resyncPersiod):
					case <-ch:
						return nil
					}
				}
			}, func(err error) {
				// Cancel the root context in the very end.
				defer func() {
					cancel()
					level.Debug(logger).Log("msg", "canceled parent context")
				}()
				ch <- struct{}{}
				var serr run.SignalError
				if ok := errors.As(err, &serr); ok {
					level.Info(logger).Log("msg", "received signal", "signal", serr.Signal.String(), "err", err.Error())
				} else {
					level.Error(logger).Log("msg", "received error", "err", err.Error())
				}
				level.Debug(logger).Log("msg", "stoped ip routes table")
				ctxWithTimeOut, cancelWithTimeOut := context.WithTimeout(ctx, 10*time.Second)
				defer func() {
					cancelWithTimeOut()
					level.Debug(logger).Log("msg", "canceled timed context")
				}()
				if err := kiloClient.KiloV1alpha1().Peers().Delete(ctxWithTimeOut, peername, metav1.DeleteOptions{}); err != nil {
					level.Error(logger).Log("failed to delete peer: %w", err)
				} else {
					level.Info(logger).Log("msg", "deleted peer", "peer", peername)
				}
				if ok, err := cmd.Flags().GetBool("clean-up"); err != nil {
					level.Error(logger).Log("err", err.Error(), "msg", "failed to get value from clean-up flag")
				} else if ok {
					cleanUp(iface, table, configPath, logger)
				}
			})
	}
	err = g.Run()
	var serr run.SignalError
	if ok := errors.As(err, &serr); ok {
		return nil
	}
	return err
}

func cleanUp(iface int, t *route.Table, configPath string, logger log.Logger) {
	if err := iproute.Set(iface, false); err != nil {
		level.Error(logger).Log("err", err.Error(), "msg", "failed to set down wg interface")
	}
	if err := os.Remove(configPath); err != nil {
		level.Error(logger).Log("error", fmt.Sprintf("failed to delete configuration file: %v", err))
	}
	if err := iproute.RemoveInterface(iface); err != nil {
		level.Error(logger).Log("error", fmt.Sprintf("failed to remove WireGuard interface: %v", err))
	}
	if err := t.CleanUp(); err != nil {
		level.Error(logger).Log("failed to clean up routes: %w", err)
	}
	return
}

//func nodeToWGPeer(n *mesh.Node) (ret wgtypes.PeerConfig, err error) {
//	pubKey, err := wgtypes.ParseKey(string(n.Key))
//	if err != nil {
//		return ret, err
//	}
//	ret.PublicKey = pubKey
//	if err != nil {
//		return ret, err
//	}
//	aIPs := []net.IPNet{*n.WireGuardIP, *n.InternalIP, *n.Subnet}
//	for _, a := range n.AllowedLocationIPs {
//		aIPs = append(aIPs, *a)
//	}
//	ret.AllowedIPs = aIPs
//
//	dur := time.Second * time.Duration(n.PersistentKeepalive)
//	ret.PersistentKeepaliveInterval = &dur
//
//	udpEndpoint, err := net.ResolveUDPAddr("udp", n.Endpoint.String())
//	if err != nil {
//		return ret, err
//	} else if udpEndpoint.IP == nil {
//		udpEndpoint = nil
//	}
//	ret.Endpoint = udpEndpoint
//	return ret, nil
//}
//
//func toWGPeer(p wireguard.Peer) (ret wgtypes.PeerConfig, err error) {
//	pubKey, err := wgtypes.ParseKey(string(p.PublicKey))
//	if err != nil {
//		return ret, err
//	}
//	ret.PublicKey = pubKey
//	aIPs := make([]net.IPNet, len(p.AllowedIPs))
//	for i, a := range p.AllowedIPs {
//		aIPs[i] = *a
//	}
//	ret.AllowedIPs = aIPs
//
//	if preSharedKey, err := wgtypes.ParseKey(string(p.PresharedKey)); len(p.PresharedKey) > 0 && err != nil {
//		return ret, err
//	} else {
//		ret.PresharedKey = &preSharedKey
//	}
//	dur := time.Second * time.Duration(p.PersistentKeepalive)
//	ret.PersistentKeepaliveInterval = &dur
//
//	udpEndpoint, err := net.ResolveUDPAddr("udp", p.Endpoint.String())
//	if err != nil {
//		return ret, err
//	} else if udpEndpoint.IP == nil {
//		udpEndpoint = nil
//	}
//	fmt.Println("peer")
//	fmt.Println(aIPs)
//	ret.Endpoint = udpEndpoint
//	return ret, nil
//}