/**
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"
)

var (
	listen string
	target string
)

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

func ServeProxy(target string, rsp http.ResponseWriter, req *http.Request) {

	url, _ := url.Parse(target)

	proxy := httputil.NewSingleHostReverseProxy(url)

	req.URL.Host = url.Host
	req.URL.Scheme = url.Scheme
	req.Header.Set("X-Forwarded-Host", req.Header.Get("Host"))

	proxy.ServeHTTP(rsp, req)

// 	fmt.Println()

// 	log.Println("\n")

// 	// log the original request
	dumpReq, _ := httputil.DumpRequest(req, true)

	color.Set(color.FgBlue)
	fmt.Println(string(dumpReq[:]))

	fmt.Println()

// 	// log the response
// 	dumpRsp, _ := httputil.DumpResponse(resp, true)
// 	color.Set(color.FgRed)
// 	fmt.Println(string(dumpRsp[:]))

// 	copyHeader(wr.Header(), resp.Header)
// 	wr.WriteHeader(resp.StatusCode)
// 	io.Copy(wr, resp.Body)

// 	color.Unset()

// 	defer resp.Body.Close()
}

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) {
		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 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 main() {
	// 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)
	}
}