package blacktip

import (
	"bufio"
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

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

// TraefikConfig turns on Traefik as the HTTP/HTTPS front end.
//
// Traefik declares its listening ports in static configuration, so it can only
// own the ports for web resources (a handful of ip:443/80 pairs that rarely
// change). Blacktip writes that file, restarts Traefik only when the set of web
// listeners changes, and hot-reloads everything else (routes, logins, targets)
// through Traefik's file provider. Raw TCP, UDP and TLS-passthrough ports stay
// with Blacktip's own listeners, which open instantly with no restart.
type TraefikConfig struct {
	Path    string // path to the traefik binary; empty = use the built-in proxy
	ACMEDNS string // DNS-challenge provider (e.g. "cloudflare") for wildcard certificates
	NoACME  bool   // serve Traefik's built-in self-signed certificate instead of Let's Encrypt
}

type traefikMgr struct {
	b        *Blacktip
	cfg      TraefikConfig
	dir      string
	gatePort int

	mu        sync.Mutex
	staticSig string
	stop      context.CancelFunc // stops the current traefik process
	lastErr   string
	maps      map[string]*mapState
	ports     map[string]int // mapping ID -> relay port, kept stable across edits so Traefik's backend URL never changes
	dynBytes  []byte
}

func newTraefik(b *Blacktip, cfg TraefikConfig) (*traefikMgr, error) {
	if _, err := exec.LookPath(cfg.Path); err != nil {
		if _, err2 := os.Stat(cfg.Path); err2 != nil {
			return nil, fmt.Errorf("traefik binary %q not found: %w", cfg.Path, err)
		}
	}
	t := &traefikMgr{b: b, cfg: cfg, dir: filepath.Join(b.cfg.DataDir, "traefik"), maps: map[string]*mapState{}, ports: map[string]int{}}
	if err := os.MkdirAll(t.dir, 0o700); err != nil {
		return nil, err
	}
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	t.gatePort = ln.Addr().(*net.TCPAddr).Port
	mux := http.NewServeMux()
	mux.HandleFunc("/auth/", t.forwardAuth)
	mux.HandleFunc("/.hammerhead/", t.gateCallback)
	go (&http.Server{Handler: mux, ReadHeaderTimeout: 10 * time.Second}).Serve(ln)
	return t, nil
}

func (t *traefikMgr) run(ctx context.Context) {
	go t.tailAccessLog(ctx)
	go t.probeCerts(ctx)
	<-ctx.Done()
	t.mu.Lock()
	if t.stop != nil {
		t.stop()
	}
	t.mu.Unlock()
}

// ---------- the gate: forward-auth and login callbacks ----------

func (t *traefikMgr) mapByID(id string) *mapState {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.maps[id]
}

func firstIP(xff string) string {
	if i := strings.IndexByte(xff, ','); i >= 0 {
		xff = xff[:i]
	}
	return strings.TrimSpace(xff)
}

// forwardAuth is called by Traefik before every request to a protected resource.
func (t *traefikMgr) forwardAuth(w http.ResponseWriter, r *http.Request) {
	ms := t.mapByID(strings.TrimPrefix(r.URL.Path, "/auth/"))
	if ms == nil {
		http.Error(w, "unknown resource", http.StatusForbidden)
		return
	}
	in := authInput{
		Method: r.Header.Get("X-Forwarded-Method"), Scheme: r.Header.Get("X-Forwarded-Proto"), Host: r.Header.Get("X-Forwarded-Host"),
		URI: r.Header.Get("X-Forwarded-Uri"), ClientIP: firstIP(r.Header.Get("X-Forwarded-For")), Cookie: r.Header.Get("Cookie"),
		Upgrade: strings.EqualFold(r.Header.Get("Upgrade"), "websocket"),
	}
	d := t.b.decide(ms, in)
	if d.Result != "allowed" && d.Result != "bypass" {
		t.b.logAccess(model.LogEntry{Time: time.Now(), Mapping: ms.m.ID, Remote: in.ClientIP, Method: in.Method, Host: hostOnly(in.Host), Path: in.URI, Result: d.Result, Status: d.Status})
	}
	if !d.Allow {
		if d.Location != "" {
			w.Header().Set("Location", d.Location)
		}
		w.WriteHeader(d.Status)
		return
	}
	if d.User != "" {
		w.Header().Set("X-Hammerhead-User", d.User)
	}
	if d.Cookie != "" {
		w.Header().Set("Cookie", d.Cookie)
	}
	w.WriteHeader(http.StatusOK)
}

// gateCallback serves /.hammerhead/callback and /.hammerhead/logout. Traefik routes
// those paths here and a middleware stamps which mapping the request is for
// (Traefik ignores paths in backend URLs, so a header is the reliable way).
func (t *traefikMgr) gateCallback(w http.ResponseWriter, r *http.Request) {
	id := r.Header.Get("X-Hammerhead-Mapping")
	ms := t.mapByID(id)
	if ms == nil {
		http.NotFound(w, r)
		return
	}
	secure := r.Header.Get("X-Forwarded-Proto") == "https"
	switch r.URL.Path {
	case "/.hammerhead/callback":
		user, result := t.b.callback(w, r, ms, secure)
		t.b.logAccess(model.LogEntry{Time: time.Now(), Mapping: id, Remote: firstIP(r.Header.Get("X-Forwarded-For")), Host: hostOnly(r.Header.Get("X-Forwarded-Host")), Path: "/.hammerhead/callback", User: user, Result: result})
	case "/.hammerhead/logout":
		t.b.logout(w, r, secure)
	default:
		http.NotFound(w, r)
	}
}

// ---------- relays: loopback ports Traefik uses to reach targets through the tunnel ----------

func (ms *mapState) startRelay() {
	var ln net.Listener
	var err error
	if ms.b.tr != nil {
		ms.b.tr.mu.Lock()
		pref := ms.b.tr.ports[ms.m.ID]
		ms.b.tr.mu.Unlock()
		if pref != 0 {
			ln, err = net.Listen("tcp", "127.0.0.1:"+strconv.Itoa(pref))
		}
	}
	if ln == nil {
		ln, err = net.Listen("tcp", "127.0.0.1:0")
	}
	if err != nil {
		log.Printf("relay for %s: %v", ms.m.Name, err)
		return
	}
	ms.relayLn = ln
	ms.relayPort = ln.Addr().(*net.TCPAddr).Port
	if ms.b.tr != nil {
		ms.b.tr.mu.Lock()
		ms.b.tr.ports[ms.m.ID] = ms.relayPort
		ms.b.tr.mu.Unlock()
	}
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func() {
				up, _, err := ms.dial("tcp", -1)
				if err != nil {
					c.Close()
					return
				}
				tunnel.Proxy(c, c, up, up)
			}()
		}
	}()
}

// ---------- configuration ----------

var nonAlnum = regexp.MustCompile(`[^a-zA-Z0-9]+`)

func epName(ip string, port int) string {
	return "ep-" + nonAlnum.ReplaceAllString(ip, "_") + "-" + strconv.Itoa(port)
}

func hostRule(host string) (string, bool) {
	if strings.HasPrefix(host, "*.") {
		return "HostRegexp(`^[^.]+\\." + regexp.QuoteMeta(host[2:]) + "$`)", true
	}
	return "Host(`" + host + "`)", false
}

// sync rewrites Traefik's configuration for the given web mappings and
// (re)starts the process if its static part changed. Called with b.mu held.
func (t *traefikMgr) sync(ctx context.Context, maps []*mapState, acmeEmail string, staging bool) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.maps = map[string]*mapState{}
	eps := map[string]string{} // name -> address
	needACME := false
	for _, ms := range maps {
		t.maps[ms.m.ID] = ms
		eps[epName(ms.m.BindIP, ms.m.Port)] = net.JoinHostPort(ms.m.BindIP, strconv.Itoa(ms.m.Port))
		if ms.m.Proto == model.ProtoHTTPS {
			needACME = true
			if ms.m.RedirectHTTP {
				eps[epName(ms.m.BindIP, 80)] = net.JoinHostPort(ms.m.BindIP, "80")
			}
		}
	}
	routers := map[string]any{}
	services := map[string]any{}
	mws := map[string]any{
		"strip-user": map[string]any{"headers": map[string]any{"customRequestHeaders": map[string]string{"X-Hammerhead-User": ""}}},
	}
	for _, ms := range maps {
		m := ms.m
		if ms.relayPort == 0 {
			continue
		}
		rule, wildcard := hostRule(m.Hostname)
		ep := epName(m.BindIP, m.Port)
		prio := 5000
		if wildcard {
			prio = 2000
		}
		scheme := "http"
		if m.BackendHTTPS {
			scheme = "https"
		}
		services["s-"+m.ID] = map[string]any{"loadBalancer": map[string]any{
			"servers":          []map[string]string{{"url": fmt.Sprintf("%s://127.0.0.1:%d", scheme, ms.relayPort)}},
			"passHostHeader":   true,
			"serversTransport": "insecure",
		}}
		chain := []string{}
		if m.Auth.Protected() {
			mws["auth-"+m.ID] = map[string]any{"forwardAuth": map[string]any{
				"address":             fmt.Sprintf("http://127.0.0.1:%d/auth/%s", t.gatePort, m.ID),
				"authResponseHeaders": []string{"X-Hammerhead-User", "Cookie"},
				"trustForwardHeader":  true,
				"maxResponseBodySize": 8192,
			}}
			chain = append(chain, "auth-"+m.ID)
		} else {
			chain = append(chain, "strip-user")
		}
		if m.RewriteHost != "" {
			mws["host-"+m.ID] = map[string]any{"headers": map[string]any{"customRequestHeaders": map[string]string{"Host": m.RewriteHost}}}
			chain = append(chain, "host-"+m.ID)
		}
		router := map[string]any{"rule": rule, "entryPoints": []string{ep}, "service": "s-" + m.ID, "middlewares": chain, "priority": prio}
		var tlsCfg map[string]any
		if m.Proto == model.ProtoHTTPS {
			tlsCfg = map[string]any{}
			if !t.cfg.NoACME {
				tlsCfg["certResolver"] = "le"
				if wildcard {
					tlsCfg["domains"] = []map[string]any{{"main": m.Hostname}}
				}
			}
			router["tls"] = tlsCfg
		}
		routers["m-"+m.ID] = router

		if m.Auth.Protected() { // login callbacks are answered by the gate, ahead of the app
			services["g-"+m.ID] = map[string]any{"loadBalancer": map[string]any{
				"servers": []map[string]string{{"url": fmt.Sprintf("http://127.0.0.1:%d", t.gatePort)}}, "passHostHeader": true}}
			mws["gate-"+m.ID] = map[string]any{"headers": map[string]any{"customRequestHeaders": map[string]string{"X-Hammerhead-Mapping": m.ID}}}
			g := map[string]any{"rule": rule + " && PathPrefix(`/.hammerhead/`)", "entryPoints": []string{ep}, "service": "g-" + m.ID,
				"middlewares": []string{"gate-" + m.ID}, "priority": prio + 1000}
			if tlsCfg != nil {
				g["tls"] = tlsCfg
			}
			routers["g-"+m.ID] = g
		}
		if m.Proto == model.ProtoHTTPS && m.RedirectHTTP {
			redirPort := ""
			if m.Port != 443 {
				redirPort = strconv.Itoa(m.Port)
			}
			mws["redir-"+m.ID] = map[string]any{"redirectScheme": map[string]any{"scheme": "https", "port": redirPort, "permanent": true}}
			routers["r-"+m.ID] = map[string]any{"rule": rule, "entryPoints": []string{epName(m.BindIP, 80)}, "service": "noop@internal",
				"middlewares": []string{"redir-" + m.ID}, "priority": prio}
		}
	}
	dynamic := map[string]any{"http": map[string]any{
		"routers": routers, "services": services, "middlewares": mws,
		"serversTransports": map[string]any{"insecure": map[string]any{"insecureSkipVerify": true}},
	}}
	t.writeDynamic(dynamic) // JSON is valid YAML

	// static part: a change here needs a restart
	names := make([]string, 0, len(eps))
	for n := range eps {
		names = append(names, n)
	}
	sort.Strings(names)
	entry := map[string]any{}
	for _, n := range names {
		// "delete" drops header aliases such as X_Hammerhead_User, which CGI/PHP style
		// backends would otherwise read as the trusted X-Hammerhead-User
		entry[n] = map[string]any{"address": eps[n], "http": map[string]any{"aliasHeadersStrategy": "delete"}}
	}
	static := map[string]any{
		"global": map[string]any{"checkNewVersion": false, "sendAnonymousUsage": false},
		"log":    map[string]any{"level": "INFO"},
		"accessLog": map[string]any{"filePath": filepath.Join(t.dir, "access.log"), "format": "json", "fields": map[string]any{
			"defaultMode": "drop",
			"names": map[string]string{"StartUTC": "keep", "RouterName": "keep", "ClientHost": "keep", "RequestMethod": "keep", "RequestHost": "keep",
				"RequestPath": "keep", "DownstreamStatus": "keep", "DownstreamContentSize": "keep", "RequestContentSize": "keep", "Duration": "keep"}}},
		"entryPoints": entry,
		"providers":   map[string]any{"file": map[string]any{"filename": filepath.Join(t.dir, "dynamic.yml"), "watch": true}},
	}
	if needACME && !t.cfg.NoACME {
		acme := map[string]any{"email": acmeEmail, "storage": filepath.Join(t.dir, "acme.json")}
		if staging {
			acme["caServer"] = "https://acme-staging-v02.api.letsencrypt.org/directory"
		}
		if t.cfg.ACMEDNS != "" {
			acme["dnsChallenge"] = map[string]any{"provider": t.cfg.ACMEDNS}
		} else {
			acme["tlsChallenge"] = map[string]any{}
		}
		static["certificatesResolvers"] = map[string]any{"le": map[string]any{"acme": acme}}
	}
	sb, _ := json.Marshal(static)
	sig := string(sb)
	if len(eps) == 0 { // nothing to serve: stop Traefik until a web port exists
		if t.stop != nil {
			t.stop()
			t.stop = nil
		}
		t.staticSig = ""
		return
	}
	if sig == t.staticSig && t.stop != nil {
		return
	}
	if t.stop != nil {
		t.stop()
	}
	writeJSONAtomic(filepath.Join(t.dir, "static.yml"), static)
	t.staticSig = sig
	pctx, cancel := context.WithCancel(ctx)
	t.stop = cancel
	go t.supervise(pctx)
}

// writeDynamic updates Traefik's dynamic configuration in place. File watchers
// on some platforms miss atomic renames, so the same bytes are written again
// shortly after: a missed event can never leave stale routes behind.
func (t *traefikMgr) writeDynamic(v any) {
	b, _ := json.MarshalIndent(v, "", "  ")
	path := filepath.Join(t.dir, "dynamic.yml")
	t.dynBytes = b
	os.WriteFile(path, b, 0o600)
	for _, d := range []time.Duration{time.Second, 4 * time.Second} {
		time.AfterFunc(d, func() {
			t.mu.Lock()
			defer t.mu.Unlock()
			if string(t.dynBytes) == string(b) { // still the latest
				os.WriteFile(path, b, 0o600)
			}
		})
	}
}

func writeJSONAtomic(path string, v any) {
	b, _ := json.MarshalIndent(v, "", "  ")
	tmp := path + ".tmp"
	if os.WriteFile(tmp, b, 0o600) == nil {
		os.Rename(tmp, path)
	}
}

// ---------- process supervision ----------

type lineLogger struct{ prefix string }

func (l lineLogger) Write(p []byte) (int, error) {
	sc := bufio.NewScanner(strings.NewReader(string(p)))
	for sc.Scan() {
		log.Printf("%s%s", l.prefix, sc.Text())
	}
	return len(p), nil
}

func (t *traefikMgr) supervise(ctx context.Context) {
	for ctx.Err() == nil {
		cmd := exec.CommandContext(ctx, t.cfg.Path, "--configFile="+filepath.Join(t.dir, "static.yml"))
		cmd.Dir = t.dir
		cmd.Stdout, cmd.Stderr = lineLogger{"traefik: "}, lineLogger{"traefik: "}
		err := cmd.Start()
		if err == nil {
			t.setErr("")
			log.Printf("traefik started (pid %d)", cmd.Process.Pid)
			err = cmd.Wait()
		}
		if ctx.Err() != nil {
			return
		}
		t.setErr(fmt.Sprintf("traefik stopped: %v", err))
		log.Printf("traefik stopped: %v; restarting in 3s", err)
		select {
		case <-time.After(3 * time.Second):
		case <-ctx.Done():
		}
	}
}

func (t *traefikMgr) setErr(s string) {
	t.mu.Lock()
	t.lastErr = s
	t.mu.Unlock()
}

func (t *traefikMgr) err() string {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.lastErr
}

// ---------- access log import ----------

type traefikLine struct {
	StartUTC         time.Time `json:"StartUTC"`
	RouterName       string    `json:"RouterName"`
	ClientHost       string    `json:"ClientHost"`
	RequestMethod    string    `json:"RequestMethod"`
	RequestHost      string    `json:"RequestHost"`
	RequestPath      string    `json:"RequestPath"`
	DownstreamStatus int       `json:"DownstreamStatus"`
	DownstreamSize   int64     `json:"DownstreamContentSize"`
	RequestSize      int64     `json:"RequestContentSize"`
	DurationNanos    int64     `json:"Duration"`
}

// tailAccessLog turns Traefik's JSON access log into Blacktip's log entries and traffic counters.
// Lines older than this process are ignored (by timestamp, so requests served
// before the first poll still count).
func (t *traefikMgr) tailAccessLog(ctx context.Context) {
	path := filepath.Join(t.dir, "access.log")
	startedAt := time.Now().Add(-2 * time.Second)
	var off int64
	first := true
	tick := time.NewTicker(time.Second)
	defer tick.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-tick.C:
		}
		f, err := os.Open(path)
		if err != nil {
			continue
		}
		st, _ := f.Stat()
		switch {
		case first && st.Size() > 4<<20:
			off = st.Size() - 4<<20 // only the recent tail of an old log
		case st.Size() < off:
			off = 0 // truncated
		}
		first = false
		f.Seek(off, io.SeekStart)
		r := bufio.NewReader(f)
		for {
			line, err := r.ReadBytes('\n')
			if err != nil {
				break // partial line: read it next time
			}
			off += int64(len(line))
			var tl traefikLine
			if json.Unmarshal(line, &tl) != nil || !strings.HasPrefix(tl.RouterName, "m-") || tl.StartUTC.Before(startedAt) {
				continue
			}
			id := strings.TrimSuffix(strings.TrimPrefix(tl.RouterName, "m-"), "@file")
			ms := t.mapByID(id)
			if ms == nil {
				continue
			}
			ms.conns.Add(1)
			ms.bytesIn.Add(tl.RequestSize)
			ms.bytesOut.Add(tl.DownstreamSize)
			t.b.logAccess(model.LogEntry{Time: tl.StartUTC, Mapping: id, Remote: tl.ClientHost, Method: tl.RequestMethod, Host: hostOnly(tl.RequestHost),
				Path: tl.RequestPath, Status: tl.DownstreamStatus, Result: "allowed", BytesIn: tl.RequestSize, BytesOut: tl.DownstreamSize, Millis: tl.DurationNanos / 1e6})
		}
		f.Close()
		if st.Size() > 64<<20 { // keep the log from growing forever (best effort)
			if os.Truncate(path, 0) == nil {
				off = 0
			}
		}
	}
}

// ---------- certificate status ----------

// probeCerts handshakes with Traefik to see which certificate it serves for each https hostname.
func (t *traefikMgr) probeCerts(ctx context.Context) {
	tick := time.NewTicker(30 * time.Second)
	defer tick.Stop()
	first := time.After(4 * time.Second)
	for {
		select {
		case <-ctx.Done():
			return
		case <-first:
		case <-tick.C:
		}
		t.mu.Lock()
		var targets []*mapState
		for _, ms := range t.maps {
			if ms.m.Proto == model.ProtoHTTPS && !strings.HasPrefix(ms.m.Hostname, "*") {
				targets = append(targets, ms)
			}
		}
		t.mu.Unlock()
		changed := false
		for _, ms := range targets {
			ip := ms.m.BindIP
			if pip := net.ParseIP(ip); pip != nil && pip.IsUnspecified() {
				ip = "127.0.0.1"
			}
			status := t.probeOne(net.JoinHostPort(ip, strconv.Itoa(ms.m.Port)), ms.m.Hostname)
			t.b.mu.Lock()
			if t.b.acme.certs == nil {
				t.b.acme.certs = map[string]string{}
			}
			if t.b.acme.certs[ms.m.Hostname] != status {
				t.b.acme.certs[ms.m.Hostname] = status
				changed = true
			}
			t.b.mu.Unlock()
		}
		if changed {
			go t.b.report()
		}
	}
}

func (t *traefikMgr) probeOne(addr, host string) string {
	d := &net.Dialer{Timeout: 3 * time.Second}
	c, err := tls.DialWithDialer(d, "tcp", addr, &tls.Config{ServerName: host, InsecureSkipVerify: true})
	if err != nil {
		return "error: " + err.Error()
	}
	defer c.Close()
	certs := c.ConnectionState().PeerCertificates
	if len(certs) == 0 {
		return "error: no certificate"
	}
	leaf := certs[0]
	switch {
	case leaf.Subject.CommonName == "TRAEFIK DEFAULT CERT":
		if t.cfg.NoACME {
			return "self-signed (Traefik default certificate)"
		}
		return "pending (still serving Traefik's default certificate)"
	case leaf.Subject.CommonName == host || containsStr(leaf.DNSNames, host):
		return "valid until " + leaf.NotAfter.Format("2006-01-02") + " (" + strings.Join(leaf.Issuer.Organization, ", ") + ")"
	}
	return "error: certificate is for " + leaf.Subject.CommonName
}

func containsStr(list []string, s string) bool {
	for _, x := range list {
		if x == s {
			return true
		}
	}
	return false
}