package lb

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	coreapi "k8s.io/api/core/v1"
	discovery "k8s.io/api/discovery/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"

	"github.com/rajch/weave/common"
	"github.com/rajch/weave/lb/admit"
)

// Controller watches Services, EndpointSlices, Nodes, EgressServices and this node's Pods, and keeps this
// node's part of the picture true. It also answers the CNI plugin's egress admissions.
type Controller struct {
	self   string
	client kubernetes.Interface
	egress rest.Interface
	dp     *Datapath
	// readLocal is the node's own interfaces (swappable in tests)
	readLocal func() (Local, error)

	services, slices, nodes, pods, egresses cache.Store
	synced                                  []cache.InformerSynced
	egressSynced                            cache.InformerSynced // nil while the EgressService API is absent
	kick                                    chan struct{}
	ready                                   chan struct{} // closed once the first pass ran

	mu           sync.Mutex // one pass at a time (the loop, or an admission)
	lastProblems map[string]bool
	// Services whose status this node wrote since it started (so a restart rewrites ipMode once)
	written map[types.UID]string
	// admitted: IP -> pod the CNI plugin gave it; fetched: pods admitted before the pod informer had them
	admitted map[string]Admission
	fetched  map[string]*coreapi.Pod
	at       map[string]time.Time
	// lastEgress is what was last applied, reused while the EgressService caches are not synced (never wipe
	// members' rules on a partial view)
	lastEgress EgressDesired
}

// Resync is how often the node re-reads its interfaces and re-applies everything even without API events.
const Resync = 30 * time.Second

// NewController runs nothing yet; Run does. egress is the EgressService API client.
func NewController(self string, client kubernetes.Interface, egress rest.Interface, dp *Datapath) *Controller {
	return &Controller{self: self, client: client, egress: egress, dp: dp, readLocal: ReadLocal,
		kick: make(chan struct{}, 1), ready: make(chan struct{}), lastProblems: map[string]bool{},
		written: map[types.UID]string{}, admitted: map[string]Admission{}, fetched: map[string]*coreapi.Pod{},
		at: map[string]time.Time{}}
}

func (c *Controller) poke() {
	select {
	case c.kick <- struct{}{}:
	default:
	}
}

func newInformer(c *Controller, lw cache.ListerWatcher, obj runtime.Object) (cache.Store, cache.Controller) {
	return cache.NewInformer(lw, obj, 0, cache.ResourceEventHandlerFuncs{
		AddFunc:    func(interface{}) { c.poke() },
		UpdateFunc: func(interface{}, interface{}) { c.poke() },
		DeleteFunc: func(interface{}) { c.poke() },
	})
}

// egressAPI reports whether the EgressService API is served (the CRD is installed).
func (c *Controller) egressAPI() (bool, error) {
	_, err := c.client.Discovery().ServerResourcesForGroupVersion(GroupVersion.String())
	if err == nil {
		return true, nil
	}
	if apierrors.IsNotFound(err) {
		return false, nil
	}
	return false, err
}

func (c *Controller) startEgress(ctx context.Context) {
	var ctl cache.Controller
	c.egresses, ctl = newInformer(c, cache.NewListWatchFromClient(c.egress, EgressServices, "", fields.Everything()), &EgressService{})
	go ctl.Run(ctx.Done())
	c.egressSynced = ctl.HasSynced
}

// Run blocks until ctx ends.
func (c *Controller) Run(ctx context.Context) {
	var sc, ec, nc, pc cache.Controller
	c.services, sc = newInformer(c, cache.NewListWatchFromClient(c.client.CoreV1().RESTClient(), "services", "", fields.Everything()), &coreapi.Service{})
	c.slices, ec = newInformer(c, cache.NewListWatchFromClient(c.client.DiscoveryV1().RESTClient(), "endpointslices", "", fields.Everything()), &discovery.EndpointSlice{})
	c.nodes, nc = newInformer(c, cache.NewListWatchFromClient(c.client.CoreV1().RESTClient(), "nodes", "", fields.Everything()), &coreapi.Node{})
	c.pods, pc = newInformer(c, cache.NewListWatchFromClient(c.client.CoreV1().RESTClient(), "pods", "",
		fields.OneTermEqualSelector("spec.nodeName", c.self)), &coreapi.Pod{})
	for _, ctl := range []cache.Controller{sc, ec, nc, pc} {
		go ctl.Run(ctx.Done())
		c.synced = append(c.synced, ctl.HasSynced)
	}
	// with the API present, nothing is applied before EgressServices are known: an empty view would wipe
	// members' rules (their traffic would leave from the node); without it, there can be no members
	for {
		present, err := c.egressAPI()
		if err == nil {
			if present {
				c.startEgress(ctx)
				c.synced = append(c.synced, c.egressSynced)
			}
			break
		}
		common.Log.Warnf("lb: cannot tell whether the EgressService API is served (retrying): %v", err)
		select {
		case <-ctx.Done():
			return
		case <-time.After(5 * time.Second):
		}
	}
	go c.serveAdmission(ctx)
	if !cache.WaitForCacheSync(ctx.Done(), c.synced...) {
		return
	}
	tick := time.NewTicker(Resync)
	defer tick.Stop()
	c.poke()
	first := true
	for {
		select {
		case <-ctx.Done():
			return
		case <-tick.C:
			if c.egressSynced == nil {
				// the API may have been installed since
				if present, err := c.egressAPI(); err == nil && present {
					c.mu.Lock()
					c.startEgress(ctx)
					c.mu.Unlock()
				}
			}
		case <-c.kick:
			// let a burst of events settle into one pass
			time.Sleep(200 * time.Millisecond)
			select {
			case <-c.kick:
			default:
			}
		}
		c.mu.Lock()
		err := c.sync(ctx)
		c.mu.Unlock()
		if err != nil {
			common.Log.Errorf("lb: %v", err)
		}
		if first {
			first = false
			close(c.ready)
		}
	}
}

// egressViewSynced: the EgressService view can be trusted (absent API: trivially empty; else synced).
func (c *Controller) egressViewSynced() bool {
	return c.egressSynced == nil || c.egressSynced()
}

// sync is one pass: publish this node, compute, apply, write statuses. Callers hold c.mu.
func (c *Controller) sync(ctx context.Context) error {
	local, err := c.readLocal()
	if err != nil {
		return fmt.Errorf("reading interfaces: %v", err)
	}
	if err := c.publish(ctx, local); err != nil {
		return err
	}
	if local.PodNet == nil {
		return fmt.Errorf("weave bridge has no address yet; not forwarding")
	}

	var nodes []Node
	var problems []string
	for _, o := range c.nodes.List() {
		n, err := ParseNode(o.(*coreapi.Node))
		if err != nil {
			problems = append(problems, err.Error())
		}
		nodes = append(nodes, n)
	}
	var svcs []*coreapi.Service
	for _, o := range c.services.List() {
		svcs = append(svcs, o.(*coreapi.Service))
	}
	var slices []*discovery.EndpointSlice
	for _, o := range c.slices.List() {
		slices = append(slices, o.(*discovery.EndpointSlice))
	}
	d := Compute(c.self, nodes, svcs, slices)
	// rewrite once after start what an earlier run may have written without ipMode
	for _, f := range d.Forwards {
		if s := c.service(f.Namespace, f.Name); s != nil && c.written[s.UID] != f.Address.String() && !hasStatusFor(d, f) {
			d.Status = append(d.Status, Status{Namespace: f.Namespace, Name: f.Name, Address: f.Address})
		}
	}

	e := c.lastEgress
	e.Status = nil // reusing the datapath, never the writes of an earlier view
	if c.egressViewSynced() {
		e = c.computeEgress(nodes, svcs)
	} else {
		problems = append(problems, "EgressServices not synced yet: keeping the egress rules last applied")
	}

	current := map[string]bool{}
	for _, p := range append(append(problems, d.Problems...), e.Problems...) {
		current[p] = true
		if !c.lastProblems[p] {
			common.Log.Warnf("lb: %s", p)
		}
	}
	c.lastProblems = current

	if err := c.dp.Apply(d, e, c.self, local); err != nil {
		return fmt.Errorf("applying datapath: %v", err)
	}
	c.lastEgress = e
	for _, st := range d.Status {
		if err := c.writeStatus(ctx, st); err != nil {
			return err
		}
	}
	for _, st := range e.Status {
		if err := PatchEgressStatus(ctx, c.egress, st.Namespace, st.Name, st.Status); err != nil {
			return fmt.Errorf("status of egress %s/%s: %v", st.Namespace, st.Name, err)
		}
		common.Log.Infof("lb: egress %s/%s: host %q, address %q, %d port blocks", st.Namespace, st.Name, st.Status.Host, st.Status.Address, len(st.Status.Blocks))
	}
	return nil
}

// computeEgress gathers this node's pods (informer, plus pods admitted before the informer had them) and computes.
func (c *Controller) computeEgress(nodes []Node, svcs []*coreapi.Service) EgressDesired {
	var egresses []*EgressService
	if c.egresses != nil {
		for _, o := range c.egresses.List() {
			egresses = append(egresses, o.(*EgressService))
		}
	}
	seen := map[string]bool{}
	var pods []*coreapi.Pod
	for _, o := range c.pods.List() {
		p := o.(*coreapi.Pod)
		seen[p.Namespace+"/"+p.Name] = true
		pods = append(pods, p)
	}
	// forget admissions the API has caught up with or that are long gone
	for ip, a := range c.admitted {
		key := a.Namespace + "/" + a.Name
		o, ok, _ := c.pods.GetByKey(key)
		caughtUp := ok && o.(*coreapi.Pod).Status.PodIP == ip
		if caughtUp || (!ok && time.Since(c.at[ip]) > 2*time.Minute) {
			delete(c.admitted, ip)
			delete(c.at, ip)
		}
	}
	for key, p := range c.fetched {
		stillAdmitted := false
		for _, a := range c.admitted {
			if a.Namespace+"/"+a.Name == key {
				stillAdmitted = true
			}
		}
		if !stillAdmitted {
			delete(c.fetched, key)
			continue
		}
		if !seen[key] {
			pods = append(pods, p)
		}
	}
	return ComputeEgress(c.self, nodes, svcs, egresses, pods, c.admitted)
}

func hasStatusFor(d Desired, f Forward) bool {
	for _, s := range d.Status {
		if s.Namespace == f.Namespace && s.Name == f.Name {
			return true
		}
	}
	return false
}

func (c *Controller) service(ns, name string) *coreapi.Service {
	o, ok, err := c.services.GetByKey(ns + "/" + name)
	if err != nil || !ok {
		return nil
	}
	return o.(*coreapi.Service)
}

// serveAdmission answers the CNI plugin on admit.SocketPath: record the pod's IP, run a pass, answer once the
// pod's rules are in place. A pod is refused (the plugin fails the ADD, kubelet retries) while this node cannot
// yet tell whether it is a member.
func (c *Controller) serveAdmission(ctx context.Context) {
	if err := os.MkdirAll(filepath.Dir(admit.SocketPath), 0o700); err != nil {
		common.Log.Errorf("lb: admission socket: %v", err)
		return
	}
	_ = os.Remove(admit.SocketPath)
	l, err := net.Listen("unix", admit.SocketPath)
	if err != nil {
		common.Log.Errorf("lb: admission socket: %v", err)
		return
	}
	_ = os.Chmod(admit.SocketPath, 0o600)
	mux := http.NewServeMux()
	mux.HandleFunc("/admit", func(w http.ResponseWriter, r *http.Request) {
		var req admit.Request
		if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<16)).Decode(&req); err != nil {
			http.Error(w, "bad request: "+err.Error(), http.StatusBadRequest)
			return
		}
		member, code, err := c.admit(r.Context(), req)
		if err != nil {
			common.Log.Warnf("lb: admission of %s/%s (%s): %v", req.Namespace, req.Name, req.IP, err)
			http.Error(w, err.Error(), code)
			return
		}
		_ = json.NewEncoder(w).Encode(admit.Response{Member: member})
	})
	srv := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second}
	go func() {
		<-ctx.Done()
		_ = srv.Close()
	}()
	if err := srv.Serve(l); err != nil && err != http.ErrServerClosed {
		common.Log.Errorf("lb: admission server: %v", err)
	}
}

func (c *Controller) admit(ctx context.Context, req admit.Request) (bool, int, error) {
	ip := net.ParseIP(req.IP).To4()
	if req.Namespace == "" || req.Name == "" || ip == nil {
		return false, http.StatusBadRequest, fmt.Errorf("need namespace, name and an IPv4 address")
	}
	select {
	case <-c.ready:
	case <-ctx.Done():
		return false, http.StatusServiceUnavailable, ctx.Err()
	case <-time.After(admit.Timeout - 5*time.Second):
		return false, http.StatusServiceUnavailable, fmt.Errorf("not ready yet")
	}
	key := req.Namespace + "/" + req.Name
	var pod *coreapi.Pod
	if o, ok, _ := c.pods.GetByKey(key); ok {
		pod = o.(*coreapi.Pod)
	} else {
		gctx, cancel := context.WithTimeout(ctx, 10*time.Second)
		p, err := c.client.CoreV1().Pods(req.Namespace).Get(gctx, req.Name, metav1.GetOptions{})
		cancel()
		if err != nil {
			return false, http.StatusServiceUnavailable, fmt.Errorf("fetching pod: %v", err)
		}
		pod = p
	}
	if pod.Spec.NodeName != c.self {
		return false, http.StatusConflict, fmt.Errorf("pod is bound to %q, not this node", pod.Spec.NodeName)
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if !c.egressViewSynced() {
		return false, http.StatusServiceUnavailable, fmt.Errorf("EgressServices not synced yet")
	}
	c.admitted[ip.String()] = Admission{Namespace: req.Namespace, Name: req.Name}
	c.at[ip.String()] = time.Now()
	if _, ok, _ := c.pods.GetByKey(key); !ok {
		c.fetched[key] = pod
	}
	if err := c.sync(ctx); err != nil {
		// with no EgressService and no member rules in place there is nothing to protect or to clean up for
		// this IP: a broken pass must not stop every pod on the node from starting
		if len(c.lastEgress.Members) == 0 && (c.egresses == nil || len(c.egresses.List()) == 0) {
			common.Log.Warnf("lb: admitting %s/%s despite a failed pass (no EgressServices): %v", req.Namespace, req.Name, err)
			return false, http.StatusOK, nil
		}
		return false, http.StatusServiceUnavailable, err
	}
	for _, m := range c.lastEgress.Members {
		if m.IP.Equal(ip) {
			return true, http.StatusOK, nil
		}
	}
	return false, http.StatusOK, nil
}

// publish keeps this node's annotations equal to what its interfaces say.
func (c *Controller) publish(ctx context.Context, local Local) error {
	want := local.Annotations()
	o, ok, err := c.nodes.GetByKey(c.self)
	if err != nil {
		return err
	}
	if ok {
		have := o.(*coreapi.Node).Annotations
		same := true
		for k, v := range want {
			if have[k] != v {
				same = false
			}
		}
		if same {
			return nil
		}
	}
	patch, err := json.Marshal(map[string]interface{}{"metadata": map[string]interface{}{"annotations": want}})
	if err != nil {
		return err
	}
	if _, err := c.client.CoreV1().Nodes().Patch(ctx, c.self, types.MergePatchType, patch, metav1.PatchOptions{}); err != nil {
		return fmt.Errorf("publishing node annotations: %v", err)
	}
	common.Log.Infof("lb: node %s holds [%s]", c.self, want[AnnAddresses])
	return nil
}

// writeStatus sets (ipMode Proxy: kube-proxy leaves the address to us) or clears a Service's load-balancer status.
func (c *Controller) writeStatus(ctx context.Context, st Status) error {
	ingress := []map[string]string{}
	if st.Address != nil {
		ingress = append(ingress, map[string]string{"ip": st.Address.String(), "ipMode": "Proxy"})
	}
	patch, err := json.Marshal(map[string]interface{}{"status": map[string]interface{}{"loadBalancer": map[string]interface{}{"ingress": ingress}}})
	if err != nil {
		return err
	}
	s, err := c.client.CoreV1().Services(st.Namespace).Patch(ctx, st.Name, types.MergePatchType, patch, metav1.PatchOptions{}, "status")
	if err != nil {
		return fmt.Errorf("status of %s/%s: %v", st.Namespace, st.Name, err)
	}
	if st.Address != nil {
		c.written[s.UID] = st.Address.String()
		common.Log.Infof("lb: %s/%s at %s", st.Namespace, st.Name, st.Address)
	} else {
		delete(c.written, s.UID)
		common.Log.Infof("lb: %s/%s has no address", st.Namespace, st.Name)
	}
	return nil
}