package htpack

import (
	"fmt"
	"net/http"
	"path"
	"strconv"
	"strings"
	"time"

	"src.lwithers.me.uk/go/htpack/packed"
)

const (
	encodingGzip   = "gzip"
	encodingBrotli = "br"
	encodingZstd   = "zstd"
)

// New returns a new handler. Standard security headers are set.
func New(packfile string) (*Handler, error) {
	mapped, err := packed.MapAndLoad(packfile)
	if err != nil {
		return nil, err
	}
	return NewMapped("/", mapped), nil
}

// NewMapped returns a new handler, given an already-mapped datafile. Standard
// security headers are set.
//
// Files are served under the directory prefix given by path, which may be
// empty or "/" for the root path. Note the [Handler.ServeMapped] method may specify
// additional files to be served under different paths, not affected by this
// path.
func NewMapped(path string, mapped *packed.Mapped) *Handler {
	h := &Handler{
		files:     make(map[string]handlerFile),
		headers:   make(map[string]string),
		startTime: time.Now(),
	}

	// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/X-Frame-Options
	h.SetHeader("X-Frame-Options", "SAMEORIGIN")

	// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/X-Content-Type-Options
	h.SetHeader("X-Content-Type-Options", "nosniff")

	h.addFiles(path, mapped)

	return h
}

// Handler implements http.Handler and allows options to be set.
type Handler struct {
	files     map[string]handlerFile
	headers   map[string]string
	startTime time.Time
	notFound  *handlerFile
}

// AddMapped serves up additional files from the given mapped packfile. The
// given path is taken as a serving prefix. If a file for a given path already
// exists, the existing file takes precedence.
func (h *Handler) AddMapped(path string, mapped *packed.Mapped) {
	h.addFiles(path, mapped)
}

// SetHeader allows a custom header to be set on HTTP responses. These are
// always emitted by ServeHTTP, whether the response status is success or
// otherwise. Note that you can override the standard security headers
// (X-Frame-Options and X-Content-Type-Options) using this function. You can
// remove previously-set headers altogether by passing an empty string for
// value.
func (h *Handler) SetHeader(key, value string) {
	if value == "" {
		delete(h.headers, key)
	} else {
		h.headers[key] = value
	}
}

// SetIndex allows setting an index.html (or equivalent) that can be used to
// serve requests landing at a directory. For instance, if a file named
// "/foo/index.html" exists, and this function is called with "index.html",
// then a route will be registered to serve the contents of this file at
// "/foo". Noting that the ServeHTTP handler discards a trailing "/" on non
// root URLs, this means that it will serve equivalent content for requests
// to "/foo/index.html", "/foo/" and "/foo".
//
// Existing routes are not overwritten, and this function could be called
// multiple times with different filenames (noting later calls would not
// overwrite files matching earlier calls).
func (h *Handler) SetIndex(filename string) {
	for k, v := range h.files {
		if path.Base(k) == filename {
			routeToAdd := path.Dir(k)
			if _, exists := h.files[routeToAdd]; !exists {
				h.files[routeToAdd] = v
			}
		}
	}
}

// SetNotFound allows overriding the returned resource when a request is made
// for a resource that does not exist. The default behaviour would be to return
// a standard HTTP 404 Not Found response; calling this function with an empty
// string will restore that behaviour.
//
// This function will return an error if the named resource is not present in
// the packfile.
func (h *Handler) SetNotFound(notFound string) error {
	if notFound == "" {
		h.notFound = nil
		return nil
	}

	notFound = path.Clean(notFound)
	f, ok := h.files[path.Clean(notFound)]
	if !ok {
		return fmt.Errorf("no such resource %q", notFound)
	}
	h.notFound = &f
	return nil
}

type handlerFile struct {
	contentType, etag                string
	uncompressed, gzip, brotli, zstd []byte
}

func (h *Handler) addFiles(prefix string, mapped *packed.Mapped) {
	// our directory contains files like "/index.html". We will add
	// a prefix, which might be "" or "/foo".
	prefix = path.Clean(prefix)
	prefix = strings.TrimSuffix(prefix, "/")
	switch {
	case prefix == "":
	case prefix == ".":
		prefix = ""
	case !strings.HasPrefix(prefix, "/"):
		prefix = "/" + prefix
	}

	for filePath, file := range mapped.Directory.Files {
		handlerPath := prefix + filePath
		if _, exists := h.files[handlerPath]; exists {
			// don't overwrite existing entries
			continue
		}
		h.files[handlerPath] = buildHandlerFile(mapped.Data, file)
	}
}

func buildHandlerFile(data []byte, file *packed.File) handlerFile {
	byteSlice := func(fileData *packed.FileData) []byte {
		if fileData == nil {
			return nil
		}
		start := fileData.Offset
		end := start + fileData.Length
		return data[start:end]
	}
	hf := handlerFile{
		contentType: file.ContentType,
		etag:        file.Etag,
	}
	hf.uncompressed = byteSlice(file.Uncompressed)
	hf.gzip = byteSlice(file.Gzip)
	hf.brotli = byteSlice(file.Brotli)
	hf.zstd = byteSlice(file.Zstd)
	return hf
}

// ServeHTTP handles requests for files. It supports GET and HEAD methods, with
// anything else returning a 405. Exact path matches are required, else a 404 is
// returned.
func (h *Handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// set custom headers before any processing; ensures these are set even
	// on error responses
	for hkey, hval := range h.headers {
		w.Header().Set(hkey, hval)
	}

	switch req.Method {
	case "HEAD", "GET":
		// OK
	default:
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	info, ok := h.files[path.Clean(req.URL.Path)]
	if !ok {
		if h.notFound == nil {
			http.NotFound(w, req)
			return
		}
		info = *h.notFound
	}

	// set standard headers
	w.Header().Set("Vary", "Accept-Encoding")
	w.Header().Set("Etag", info.etag)
	w.Header().Set("Content-Type", info.contentType)
	w.Header().Set("Accept-Ranges", "bytes")

	// process etag / modtime
	if clientHasCachedVersion(info.etag, h.startTime, req) {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	// select compression
	data := info.uncompressed
	gzip, brotli, zstd := acceptedEncodings(req)
	switch {
	case zstd && info.zstd != nil:
		data = info.zstd
		w.Header().Set("Content-Encoding", encodingZstd)
	case brotli && info.brotli != nil:
		data = info.brotli
		w.Header().Set("Content-Encoding", encodingBrotli)
	case gzip && info.gzip != nil:
		data = info.gzip
		w.Header().Set("Content-Encoding", encodingGzip)
	}

	// range support (single-part ranges only)
	// https://developer.mozilla.org/en-US/docs/Web/HTTP/Range_requests#Single_part_ranges
	offset, length, isPartial := getFileRange(uint64(len(data)), req)
	if isPartial {
		w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d",
			offset, offset+length-1, len(data)))
	}

	// now we know exactly what we're writing, finalise HTTP header
	w.Header().Set("Content-Length", strconv.FormatUint(length, 10))
	if isPartial {
		w.WriteHeader(http.StatusPartialContent)
	} else {
		w.WriteHeader(http.StatusOK)
	}

	// send body (though not for HEAD)
	if req.Method == "HEAD" {
		return
	}

	w.Write(data[offset : offset+length])
}

func acceptedEncodings(req *http.Request) (gzip, brotli, zstd bool) {
	encodings := req.Header.Get("Accept-Encoding")
	for _, enc := range strings.Split(encodings, ",") {
		switch strings.TrimSpace(enc) {
		case encodingGzip:
			gzip = true
		case encodingBrotli:
			brotli = true
		}
	}
	return
}

// clientHasCachedVersion returns true if the client has a cached version of
// the resource. We'll check the etags presented by the client, but if etags
// are not present then we'll check the if-modified-since date.
func clientHasCachedVersion(etag string, startTime time.Time, req *http.Request,
) bool {
	checkEtags := req.Header.Get("If-None-Match")
	for _, check := range strings.Split(checkEtags, ",") {
		if etag == strings.TrimSpace(check) {
			// client knows the etag, so it has this version of the
			// resource cached already
			return true
		}
	}

	// if the client presented etags at all, we use that as our definitive
	// answer
	if _, sawEtags := req.Header["If-None-Match"]; sawEtags {
		return false
	}

	// check the timestamp the client last grabbed the resource
	cachedTime, err := http.ParseTime(req.Header.Get("If-Modified-Since"))
	if err != nil {
		return false
	}
	return cachedTime.After(startTime)
}

// getFileRange returns the byte offset and length of the file to serve, along
// with whether or not it's partial content.
func getFileRange(dataLength uint64, req *http.Request) (offset, length uint64, isPartial bool) {
	length = dataLength

	// only accept "Range: bytes=…"
	r := req.Header.Get("Range")
	if !strings.HasPrefix(r, "bytes=") {
		return
	}
	r = strings.TrimPrefix(r, "bytes=")

	// only accept a single range, "from-to", mapping to interval [from,to]
	pos := strings.IndexByte(r, '-')
	if pos == -1 {
		return
	}
	sfrom, sto := r[:pos], r[pos+1:]
	from, err := strconv.ParseUint(sfrom, 10, 64)
	if err != nil {
		return
	}
	to, err := strconv.ParseUint(sto, 10, 64)
	if err != nil {
		return
	}

	// validate the interval lies within the file
	switch {
	case from > to,
		from >= dataLength,
		to >= dataLength:
		return
	}

	// all good
	offset = from
	length = to - from + 1
	isPartial = true
	return
}