package main

import (
	"bytes"
	"flag"
	"golang.org/x/net/websocket"
	"io"
	"k8s.io/client-go/1.4/rest"
	"k8s.io/client-go/1.4/tools/clientcmd"
	"log"
	"net/http"
)

var (
	kubeconfig = flag.String("kubeconfig", "/Users/myself/.kube/config", "absolute path to the kubeconfig file")
)

func main() {
	flag.Parse()
	// uses the current context in kubeconfig
	config, err := clientcmd.BuildConfigFromFlags("", *kubeconfig)
	if err != nil {
		panic(err.Error())
	}
	apiURL := "wss://192.168.64.17:8443/api/v1/pods?watch=true"
	wsc, err := websocket.NewConfig(apiURL, apiURL)
	if err != nil {
		log.Fatal(err)
	}
	wsc.TlsConfig, err = rest.TLSConfigFor(config)
	if err != nil {
		log.Fatal(err)
	}
	wsc.Header, err = headersForConfig(config)
	if err != nil {
		log.Fatal(err)
	}

	wsConn, err := websocket.DialConfig(wsc)
	if err != nil {
		log.Fatal(err)
	}
	defer wsConn.Close()

	buf := &bytes.Buffer{}
	for {
		var msg []byte
		if err := websocket.Message.Receive(wsConn, &msg); err != nil {
			if err == io.EOF {
				break
			}
			log.Fatalf("Failed to read completely from websocket %v", err)
		}
		if len(msg) == 0 {
			continue
		}
		log.Printf("Read %v %v", len(msg), string(msg))
		buf.Write(msg)
	}

	if buf.String() != "container is alive\n" {
		log.Fatalf("Unexpected websocket logs:\n%s", buf.String())
	}
}

type extractRT struct {
	http.Header
}

func (rt *extractRT) RoundTrip(req *http.Request) (*http.Response, error) {
	rt.Header = req.Header
	return nil, nil
}

func headersForConfig(c *rest.Config) (http.Header, error) {
	extract := &extractRT{}
	rt, err := rest.HTTPWrappersForConfig(c, extract)
	if err != nil {
		return nil, err
	}
	if _, err := rt.RoundTrip(&http.Request{}); err != nil {
		return nil, err
	}
	return extract.Header, nil
}