package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"strconv"
	"strings"
	"time"

	"golang.org/x/time/rate"
)

func main() {
	rand.Seed(time.Now().UnixNano())
	var port int
	var conf string
	var proxy bool
	flag.IntVar(&port, "p", 8080, "port to listen on")
	flag.StringVar(&conf, "c", "/dev/null", "line delimited file to read urls from")
	flag.BoolVar(&proxy, "proxy", false, "proxy content rather than redirect, so refresh works")
	flag.Parse()
	b, err := ioutil.ReadFile(conf)
	if err != nil {
		panic(err)
	}
	links := bytes.Split(b, []byte{'\n'})
	log.Print(port)
	limiter := rate.NewLimiter(1, 1)
	if err := http.ListenAndServe(fmt.Sprintf(":%d", port), ServeHTTP(limiter, links, proxy)); err != nil {
		panic(err)
	}
}

func ServeHTTP(limiter *rate.Limiter, links [][]byte, proxy bool) http.HandlerFunc {
	handler := serveHTTP(limiter, links, proxy)
	return func(w http.ResponseWriter, r *http.Request) {
		s := r.URL.Query().Get("n")
		if s == "" {
			s = "1"
		}
		n, err := strconv.Atoi(s)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		for i := 0; i < n; i++ {
			handler(w, r)
		}
	}
}

func serveHTTP(limiter *rate.Limiter, links [][]byte, proxy bool) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		limiter.Wait(r.Context())
		link := ""
		for len(link) == 0 {
			link = string(links[rand.Intn(len(links))])
		}
		if strings.HasPrefix(link, "http") {
			_serveHTTP(w, r, proxy, link)
		} else if strings.HasPrefix(link, "literal://") {
			w.Write([]byte(strings.TrimPrefix(link, "literal://") + "\n"))
		}
	}
}

func _serveHTTP(w http.ResponseWriter, r *http.Request, proxy bool, link string) {
	if proxy {
		resp, err := http.Get(link)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadGateway)
		}
		for k, v := range resp.Header {
			for _, v2 := range v {
				w.Header().Add(k, v2)
			}
		}
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
	} else {
		http.Redirect(w, r, link, http.StatusTemporaryRedirect)
	}
}