package blacktip

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"time"

	"git.veillax.com/Veillax/reef/authtoken"
	"git.veillax.com/Veillax/reef/model"
)

const authCookie = "hh_auth"

// recorder captures what a handler sent, without breaking websockets or
// flushing (http.ResponseController follows Unwrap).
type recorder struct {
	http.ResponseWriter
	status int
	bytes  int64
}

func (r *recorder) WriteHeader(code int) {
	if r.status == 0 {
		r.status = code
	}
	r.ResponseWriter.WriteHeader(code)
}

func (r *recorder) Write(p []byte) (int, error) {
	if r.status == 0 {
		r.status = 200
	}
	n, err := r.ResponseWriter.Write(p)
	r.bytes += int64(n)
	return n, err
}

func (r *recorder) Unwrap() http.ResponseWriter { return r.ResponseWriter }

// initProxy builds the built-in reverse proxy (used when Traefik isn't).
func (ms *mapState) initProxy() {
	scheme, port := "http", 80
	if ms.m.BackendHTTPS {
		scheme, port = "https", 443
	}
	ms.transport = &http.Transport{
		// The request URL host is "t<index>": it names the target and keeps
		// pooled connections to different targets apart.
		DialContext: func(ctx context.Context, _, addr string) (net.Conn, error) {
			idx := -1
			if h, _, err := net.SplitHostPort(addr); err == nil && strings.HasPrefix(h, "t") {
				if n, err := strconv.Atoi(h[1:]); err == nil {
					idx = n
				}
			}
			c, _, err := ms.dial("tcp", idx)
			return c, err
		},
		TLSClientConfig:     &tls.Config{InsecureSkipVerify: true}, // backends are usually self-signed
		TLSHandshakeTimeout: 10 * time.Second,
		MaxIdleConnsPerHost: 32,
		IdleConnTimeout:     60 * time.Second,
	}
	ms.proxy = &httputil.ReverseProxy{
		Rewrite: func(pr *httputil.ProxyRequest) {
			pr.SetXForwarded()
			idx := ms.order()[0]
			pr.Out.URL.Scheme = scheme
			pr.Out.URL.Host = fmt.Sprintf("t%d:%d", idx, port)
			if ms.m.RewriteHost != "" {
				pr.Out.Host = ms.m.RewriteHost
			}
		},
		Transport:     ms.transport,
		FlushInterval: -1,
		ErrorLog:      log.New(io.Discard, "", 0),
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			http.Error(w, "Bad gateway: the service behind this address isn't answering.", http.StatusBadGateway)
		},
	}
}

func ruleMatches(ru model.Rule, path, clientIP string) bool {
	switch ru.Match {
	case "path":
		return strings.HasPrefix(path, ru.Value)
	case "ip":
		ip := net.ParseIP(clientIP)
		if ip == nil {
			return false
		}
		if p := net.ParseIP(ru.Value); p != nil {
			return p.Equal(ip)
		}
		if _, n, err := net.ParseCIDR(ru.Value); err == nil {
			return n.Contains(ip)
		}
	}
	return false
}

func safeDest(p string) string {
	if p == "" || p[0] != '/' || strings.HasPrefix(p, "//") || strings.HasPrefix(p, "/\\") {
		return "/"
	}
	return p
}

func hostFromAddr(a string) string {
	if h, _, err := net.SplitHostPort(a); err == nil {
		return h
	}
	return a
}

// ---------- the access decision, shared by the built-in proxy and the Traefik gate ----------

type authInput struct {
	Method, Scheme, Host, URI, ClientIP, Cookie string
	Upgrade                                     bool
}

type authDecision struct {
	Allow    bool
	Status   int    // when !Allow: 302, 401 or 403
	Location string // for 302
	User     string
	Cookie   string // the request's Cookie header with our own cookie removed
	Result   string
}

func cookieValue(header, name string) string {
	for _, c := range (&http.Request{Header: http.Header{"Cookie": {header}}}).Cookies() {
		if c.Name == name {
			return c.Value
		}
	}
	return ""
}

func stripCookie(header string) string {
	var keep []string
	for _, c := range (&http.Request{Header: http.Header{"Cookie": {header}}}).Cookies() {
		if c.Name != authCookie {
			keep = append(keep, c.Name+"="+c.Value)
		}
	}
	return strings.Join(keep, "; ")
}

// decide applies a web mapping's rules and login policy to one request.
func (b *Blacktip) decide(ms *mapState, in authInput) authDecision {
	m := ms.m
	path, query, rawPath := "/", url.Values{}, "/"
	if u, err := url.ParseRequestURI(in.URI); err == nil {
		path, rawPath, query = u.Path, u.EscapedPath(), u.Query()
	}
	d := authDecision{Cookie: stripCookie(in.Cookie), Result: "allowed"}
	bypass := false
	for _, ru := range m.Rules {
		if !ruleMatches(ru, path, in.ClientIP) {
			continue
		}
		if ru.Action == "deny" {
			return authDecision{Status: http.StatusForbidden, Result: "denied"}
		}
		bypass, d.Result = true, "bypass"
		break
	}
	if !m.Auth.Protected() || bypass {
		d.Allow = true
		return d
	}
	b.mu.Lock()
	key, dash := b.authKey, b.dashboard
	b.mu.Unlock()
	if t, ok := authtoken.Verify(key, cookieValue(in.Cookie, authCookie), authtoken.KindSession, m.ID, ms.epoch.Load()); ok {
		d.Allow, d.User = true, t.Sub
		return d
	}
	share := query.Get("hh_share")
	query.Del("hh_share")
	back := in.Scheme + "://" + in.Host + rawPath
	if q := query.Encode(); q != "" {
		back += "?" + q
	}
	login := fmt.Sprintf("%s/auth/r/%s?return=%s", strings.TrimRight(dash, "/"), m.ID, url.QueryEscape(back))
	if share != "" {
		login += "&share=" + url.QueryEscape(share)
	}
	d.Cookie, d.Result = "", "auth-required"
	if (in.Method == http.MethodGet || in.Method == http.MethodHead) && !in.Upgrade {
		d.Status, d.Location = http.StatusFound, login
	} else {
		d.Status = http.StatusUnauthorized
	}
	return d
}

// callback trades hammerhead's short-lived login token for this visitor's cookie.
func (b *Blacktip) callback(w http.ResponseWriter, r *http.Request, ms *mapState, secure bool) (user, result string) {
	b.mu.Lock()
	key := b.authKey
	b.mu.Unlock()
	t, ok := authtoken.Verify(key, r.URL.Query().Get("t"), authtoken.KindCallback, ms.m.ID, ms.epoch.Load())
	if !ok {
		http.Error(w, "That login link is invalid or has expired. Go back and sign in again.", http.StatusForbidden)
		return "", "bad-login-token"
	}
	hours := ms.m.Auth.SessionHours
	if hours <= 0 {
		hours = 24
	}
	sess := authtoken.Sign(key, authtoken.Token{Mapping: ms.m.ID, Sub: t.Sub, Kind: authtoken.KindSession, Epoch: t.Epoch,
		Exp: time.Now().Add(time.Duration(hours) * time.Hour).Unix()})
	http.SetCookie(w, &http.Cookie{Name: authCookie, Value: sess, Path: "/", HttpOnly: true, Secure: secure,
		SameSite: http.SameSiteLaxMode, MaxAge: hours * 3600})
	http.Redirect(w, r, safeDest(r.URL.Query().Get("r")), http.StatusFound)
	return t.Sub, "login"
}

func (b *Blacktip) logout(w http.ResponseWriter, r *http.Request, secure bool) {
	http.SetCookie(w, &http.Cookie{Name: authCookie, Value: "", Path: "/", MaxAge: -1, HttpOnly: true, Secure: secure})
	http.Redirect(w, r, "/", http.StatusFound)
}

// ServeHTTP handles every request on a built-in http/https listener.
func (l *listener) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	b := l.b
	start := time.Now()
	rec := &recorder{ResponseWriter: w}
	host := hostOnly(r.Host)
	ms, found := route(l.snapshot(), host)
	if !found {
		http.Error(rec, "Nothing is served at this address.", http.StatusNotFound)
		return
	}
	entry := model.LogEntry{Time: start, Mapping: ms.m.ID, Remote: r.RemoteAddr, Method: r.Method, Host: host, Path: r.URL.Path, Result: "allowed"}
	defer func() {
		entry.Status = rec.status
		if entry.Status == 0 {
			entry.Status = 200
			if r.Header.Get("Upgrade") != "" {
				entry.Status = 101
			}
		}
		if r.ContentLength > 0 {
			entry.BytesIn = r.ContentLength
		}
		entry.BytesOut = rec.bytes
		entry.Millis = time.Since(start).Milliseconds()
		ms.conns.Add(1)
		ms.bytesIn.Add(entry.BytesIn)
		ms.bytesOut.Add(entry.BytesOut)
		b.logAccess(entry)
	}()

	if ms.m.IsRedirect() {
		code := http.StatusFound
		if ms.m.RedirectPermanent {
			code = http.StatusMovedPermanently
		}
		entry.Result = "redirect"
		http.Redirect(rec, r, ms.m.RedirectLocation(r.URL.RequestURI()), code)
		return
	}

	secure := l.kind == model.ProtoHTTPS
	scheme := "http"
	if secure {
		scheme = "https"
	}
	if ms.m.Auth.Protected() {
		switch r.URL.Path {
		case "/.hammerhead/callback":
			entry.User, entry.Result = b.callback(rec, r, ms, secure)
			return
		case "/.hammerhead/logout":
			b.logout(rec, r, secure)
			entry.Result = "logout"
			return
		}
	}

	d := b.decide(ms, authInput{Method: r.Method, Scheme: scheme, Host: r.Host, URI: r.URL.RequestURI(),
		ClientIP: hostFromAddr(r.RemoteAddr), Cookie: r.Header.Get("Cookie"), Upgrade: r.Header.Get("Upgrade") != ""})
	entry.Result, entry.User = d.Result, d.User
	if !d.Allow {
		switch d.Status {
		case http.StatusFound:
			http.Redirect(rec, r, d.Location, http.StatusFound)
		case http.StatusForbidden:
			http.Error(rec, "Forbidden", http.StatusForbidden)
		default:
			http.Error(rec, "Authentication required.", http.StatusUnauthorized)
		}
		return
	}

	// the backend never sees our cookie, and can trust X-Hammerhead-User
	r.Header.Del("X-Hammerhead-User")
	if d.User != "" {
		r.Header.Set("X-Hammerhead-User", d.User)
	}
	if d.Cookie != "" {
		r.Header.Set("Cookie", d.Cookie)
	} else {
		r.Header.Del("Cookie")
	}
	ms.proxy.ServeHTTP(rec, r)
}