package main

import (
	"flag"
	"net/http"
	"net/url"
	"io"
	"os"
	"os/signal"
	"runtime"
	"time"
	"fmt"
	"strings"
	"sync"

	"github.com/elazarl/goproxy"
	"golang.org/x/sys/unix"
	"github.com/dustin/go-humanize"
	"golang.org/x/net/proxy"

)

var wg sync.WaitGroup

func main() {
	flag.Parse()
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, unix.SIGUSR1, unix.SIGTERM, unix.SIGINT, unix.SIGABRT)
	go func() {
		// Wait for signal
		for sig := range sigs {
			switch sig {
			case unix.SIGUSR1:
				bt := make([]byte, 250*1024)
				size := runtime.Stack(bt, true)

				fmt.Printf("Stack dump requested\n%s", bt[:size])
			case unix.SIGTERM:
				os.Exit(0)
			case unix.SIGINT:
				os.Exit(0)
			}
		}
	}()

	go startGoProxy(`127.0.0.1:9090`)
	time.Sleep(2 * time.Second)
	go DownloadFile("/tmp/tmp.x","https://testfileorg.netwet.net/testfile.org-1GB- Corrupt.zip")
	wg.Add(1)
	wg.Wait()

}

func startGoProxy(aInAddr string) {
	addr := aInAddr

	proxy := goproxy.NewProxyHttpServer()
	proxyServ := &http.Server{
		Addr:           addr,
		Handler:        proxy,
		MaxHeaderBytes: 1 << 20,
	}
	for retryCount := 0; retryCount < 10; retryCount++ {
		fmt.Println("Starting GoProxy ")
		err := proxyServ.ListenAndServe()
		if err != nil {
			if err == http.ErrServerClosed {
				fmt.Println("Proxy has closed connections")
				break
			} else if retryCount == 10-1 {
				fmt.Println("Proxy has quit " + err.Error())
			}
			time.Sleep(1 * time.Minute)
		}
	}
}

// WriteCounter counts the number of bytes written to it. It implements to the io.Writer interface
// and we can pass this into io.TeeReader() which will report progress on each write cycle.
type WriteCounter struct {
	Total uint64
}

func (wc *WriteCounter) Write(p []byte) (int, error) {
	n := len(p)
	wc.Total += uint64(n)
	wc.PrintProgress()
	return n, nil
}

func (wc WriteCounter) PrintProgress() {
	// Clear the line by using a character return to go back to the start and remove
	// the remaining characters by filling it with spaces
	fmt.Printf("\r%s", strings.Repeat(" ", 35))

	// Return again and print current status of download
	// We use the humanize package to print the bytes in a meaningful way (e.g. 10 MB)
	fmt.Printf("\rDownloading... %s complete", humanize.Bytes(wc.Total))
}

func proxyAwareHttpClient() *http.Client {
	var dialer proxy.Dialer
	dialer = proxy.Direct
	proxyUrl, err := url.Parse("http://127.0.0.1:9090")
	if err != nil {
		fmt.Println("Unable to parse URL")
		return nil
	}
	fmt.Printf("Proxy URL %v, %v", proxyUrl, proxy.Direct)

	// setup a http client
	httpTransport := &http.Transport{Proxy:http.ProxyURL(proxyUrl)}
	httpClient := &http.Client{Transport: httpTransport}
	httpTransport.Dial = dialer.Dial
	return httpClient
}

// DownloadFile will download a url to a local file. It's efficient because it will
// write as it downloads and not load the whole file into memory. We pass an io.TeeReader
// into Copy() to report progress on the download.
func DownloadFile(filepath string, url string) error {

	defer wg.Done()
	// Create the file, but give it a tmp file extension, this means we won't overwrite a
	// file until it's downloaded, but we'll remove the tmp extension once downloaded.
	out, err := os.CreateTemp("", "test")
	if err != nil {
		return err
	}

	defer os.Remove(out.Name())
	fmt.Println("Download Started")
	req, err := http.NewRequest("GET", url, nil)
	client := proxyAwareHttpClient()
	if client == nil {
		return nil	
	}
	resp, err := client.Do(req)
	if err != nil {
		out.Close()
		panic(err)
	}
	defer resp.Body.Close()

	// Create our progress reporter and pass it to be used alongside our writer
	counter := &WriteCounter{}
	if _, err = io.Copy(out, io.TeeReader(resp.Body, counter)); err != nil {
		out.Close()
		return err
	}

	// The progress uses the same line so print a new line once it's finished downloading
	fmt.Print("\n")

	// Close the file without defer so it can happen before Rename()
	out.Close()

	fmt.Println("Download Finished")
	return nil
}