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)
	server := Server{
		limiter: rate.NewLimiter(1, 1),
		links:   links,
		proxy:   proxy,
	}
	if err := http.ListenAndServe(fmt.Sprintf(":%d", port), server); err != nil {
		panic(err)
	}
}

type Server struct {
	limiter *rate.Limiter
	proxy   bool
	links   [][]byte
}

func (server Server) ServeHTTP(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
	}

	links := make([]string, n)
	for i := range links {
		link := ""
		for len(link) == 0 || strings.Contains(fmt.Sprint(links), link) {
			link = string(server.links[rand.Intn(len(server.links))])
		}
		links[i] = link
	}

	for _, link := range links {
		if strings.HasPrefix(link, "http") {
			server.serveHTTP(w, r, link)
		} else if strings.HasPrefix(link, "literal://") {
			w.Write([]byte(strings.TrimPrefix(link, "literal://")))
		}
		w.Write([]byte("\n"))
	}
}

func (server Server) serveHTTP(w http.ResponseWriter, r *http.Request, link string) {
	if server.proxy {
		server.limiter.Wait(r.Context())
		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)
	}
}