/**
Author: Mark George <mark.george@otago.ac.nz
License: WTFPL v2 <http://www.wtfpl.net/txt/copying/>
Original Source: https://gist.github.com/yowu/f7dc34bd4736a65ff28d
*/

package main

import (
	// "flag"
	"fmt"
	// "io"
	// "log"
	"net/http"
	"net/http/httputil"
	"net/url"

	// "strings"
	"time"

	"github.com/fatih/color"
)

type ProxyCommand struct {
	port   string `short:"p" long:"port" description:"The port that the proxy listens on."`
	target string `short:"t" long:"target" descrtiption:"The target web server that the proxy forwards requests to."`
}

func (c *ProxyCommand) Execute(args []string) error {
	fmt.Println("Start proxy")
	return nil
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

func CreateProxy(target string) (*httputil.ReverseProxy, error) {

	url, err := url.Parse(target)

	// bad target URL
	if err != nil {
		return nil, err
	}

	proxy := httputil.NewSingleHostReverseProxy(url)

	OriginalDirector := proxy.Director
	proxy.Director = func(req *http.Request) {
		FixHeaders(url, req)
		LogRequest(req)

		// send request to target
		OriginalDirector(req)
	}

	proxy.ModifyResponse = func(rsp *http.Response) error {
		LogResponse(rsp)
		return nil
	}

	proxy.ErrorHandler = func(http.ResponseWriter, *http.Request, error) {
		fmt.Println("Could not connect to target server\n")
	}

	return proxy, nil
}

func FixHeaders(url *url.URL, req *http.Request) {
	// this may be necessary for proxying HTTPS (need to test)
	req.URL.Host = url.Host
	req.URL.Scheme = url.Scheme
	req.Host = url.Host
}

func RequestHandler(proxy *httputil.ReverseProxy) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		proxy.ServeHTTP(w, r)
	}
}

func LogRequest(req *http.Request) error {
	color.Set(color.FgWhite)
	fmt.Println(time.Now().Format(time.StampMilli))

	dumpReq, _ := httputil.DumpRequest(req, true)

	color.Set(color.FgBlue)
	fmt.Println(string(dumpReq[:]))
	color.Unset()
	return nil
}

func LogResponse(rsp *http.Response) {
	color.Set(color.FgWhite)
	fmt.Println(time.Now().Format(time.StampMilli))
	dumpRsp, _ := httputil.DumpResponse(rsp, true)
	color.Set(color.FgRed)
	fmt.Println(string(dumpRsp[:]))
	color.Unset()
}

func ymain() {
	// flag.StringVar(&listen, "listen", "9090", "The port that the proxy listens on.")
	// flag.StringVar(&listen, "l", "9090", "The port that the proxy listens on.")
	// flag.StringVar(&target, "target", "8080", "The port that the proxy forwards requests to.")
	// flag.StringVar(&target, "t", "8080", "The port that the proxy forwards requests to.")

	// flag.Parse()

	// handler := &proxy{}

	// log.Println("Forwarding port", listen, "to port", target)
	// if err := http.ListenAndServe("localhost:"+listen, handler); err != nil {
	// 	log.Fatal("ListenAndServe:", err)
	// }

	target := "http://localhost:8080"
	listenPort := "9090"
	proxy, err := CreateProxy(target)

	if err != nil {
		panic(err)
	}

	http.HandleFunc("/", RequestHandler(proxy))
	if err := http.ListenAndServe(":"+listenPort, nil); err != nil {
		panic(err)
	}
}