package main

import (
	"encoding/json"
	_ "expvar"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof"
	"sort"
	"strings"
	"sync"
	"time"

	"code.justin.tv/common/chitin"
	"golang.org/x/net/context"
)

func main() {
	listen := flag.String("http", "127.0.0.1:8080", "IP and port of HTTP listener")
	debug := flag.String("debug", "127.0.0.1:8081", "IP and port of debug HTTP listener")
	flag.Parse()

	log.SetFlags(log.LUTC | log.Ldate | log.Ltime | log.Lmicroseconds)

	s := newServer()

	l, err := net.Listen("tcp", *listen)
	if err != nil {
		log.Fatalf("listen err=%q", err)
	}
	dl, err := net.Listen("tcp", *debug)
	if err != nil {
		log.Fatalf("debug listen err=%q", err)
	}

	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)

	go s.serve(ctx, &http.Server{Handler: chitin.Handler(s)}, l)
	go s.serve(ctx, &http.Server{Handler: chitin.Handler(http.DefaultServeMux)}, dl)
	go s.run(ctx)

	err = <-s.errs
	if err != nil {
		cancel()
		log.Fatalf("serve err=%q", err)
	}
}

type server struct {
	errs chan error

	mu    sync.Mutex
	count map[cspReport]int
}

func newServer() *server {
	s := &server{
		errs:  make(chan error, 1),
		count: make(map[cspReport]int),
	}

	return s
}

func (s *server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Path {
	case "/csp-report":
		s.handleReport(w, r)
		return
	}

	w.Header().Set("Content-Security-Policy", strings.Join([]string{
		"default-src https:",
		"img-src https: *.justin.tv",
		"report-uri http://localhost:8080/csp-report",
	}, "; "))

	log.Printf("request-uri=%q", r.URL.Path)

	fmt.Fprintf(w,
		`<html><head></head><body><img src="http://graphite.internal.justin.tv/content/img/refresh.png" /></body></html>`)
}

func (s *server) handleReport(w http.ResponseWriter, r *http.Request) {
	var rep report
	err := json.NewDecoder(r.Body).Decode(&rep)
	if err != nil {
		log.Printf("body-read err=%q", err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	key := cspReport{
		DocumentURI:        rep.CSPReport.DocumentURI,
		BlockedURI:         rep.CSPReport.BlockedURI,
		SourceFile:         rep.CSPReport.SourceFile,
		ViolatedDirective:  rep.CSPReport.ViolatedDirective,
		EffectiveDirective: rep.CSPReport.EffectiveDirective,
	}

	s.mu.Lock()
	n := s.count[key]
	s.count[key] = n + 1
	s.mu.Unlock()

	if n == 0 {
		log.Printf("document-uri=%q blocked-uri=%q source-file=%q "+
			"violated-directive=%q effective-directive=%q",
			key.DocumentURI, key.BlockedURI, key.SourceFile,
			key.ViolatedDirective, key.EffectiveDirective)
	}

	w.WriteHeader(http.StatusNoContent)
}

func (s *server) serve(ctx context.Context, srv *http.Server, l net.Listener) {
	err := srv.Serve(l)
	if err != nil {
		select {
		case s.errs <- err:
		case <-ctx.Done():
		}
	}
}

func (s *server) run(ctx context.Context) {
	dump := time.NewTicker(10 * time.Second)
	defer dump.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-dump.C:
			s.mu.Lock()
			rc := make(reportCounts, 0, len(s.count))
			for report, count := range s.count {
				rc = append(rc, reportCount{report: report, count: count})
			}
			s.count = make(map[cspReport]int)
			s.mu.Unlock()

			sort.Sort(rc)
			for _, rep := range rc {
				log.Printf("count=%d document-uri=%q blocked-uri=%q source-file=%q "+
					"violated-directive=%q effective-directive=%q",
					rep.count, rep.report.DocumentURI, rep.report.BlockedURI, rep.report.SourceFile,
					rep.report.ViolatedDirective, rep.report.EffectiveDirective)
			}
		}
	}
}

type report struct {
	CSPReport cspReport `json:"csp-report"`
}

type cspReport struct {
	DocumentURI        string `json:"document-uri"`
	Referrer           string `json:"referrer"`
	ViolatedDirective  string `json:"violated-directive"`
	EffectiveDirective string `json:"effective-directive"`
	OriginalPolicy     string `json:"original-policy"`
	BlockedURI         string `json:"blocked-uri"`
	StatusCode         int    `json:"status-code"`

	SourceFile   string `json:"source-file"`
	LineNumber   int    `json:"line-number"`
	ColumnNumber int    `json:"column-number"`
}

type reportCount struct {
	report cspReport
	count  int
}

type reportCounts []reportCount

func (rc reportCounts) Len() int      { return len(rc) }
func (rc reportCounts) Swap(i, j int) { rc[i], rc[j] = rc[j], rc[i] }
func (rc reportCounts) Less(i, j int) bool {
	if rc[i].count != rc[j].count {
		return rc[i].count < rc[j].count
	}
	if rc[i].report.DocumentURI != rc[j].report.DocumentURI {
		return rc[i].report.DocumentURI < rc[j].report.DocumentURI
	}
	if rc[i].report.BlockedURI != rc[j].report.BlockedURI {
		return rc[i].report.BlockedURI < rc[j].report.BlockedURI
	}
	if rc[i].report.SourceFile != rc[j].report.SourceFile {
		return rc[i].report.SourceFile < rc[j].report.SourceFile
	}
	return false
}
