package edge

import (
	"bufio"
	"context"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/carlosframework/platform/internal/geo"
	"github.com/carlosframework/platform/internal/geoip"
	"github.com/carlosframework/platform/internal/registry"
)

var geoHeaderNames = []string{geoip.HeaderCountry, geoip.HeaderSubdivision, geoip.HeaderLat, geoip.HeaderLon, geoip.HeaderSource, geoip.HeaderHop}

// geoEchoBackend answers with the geo headers it received, as JSON.
func geoEchoBackend(t *testing.T) *httptest.Server {
	t.Helper()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got := map[string]string{}
		for _, k := range geoHeaderNames {
			if v := r.Header.Get(k); v != "" {
				got[k] = v
			}
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(got)
	}))
	t.Cleanup(ts.Close)
	return ts
}

// geoEdge serves host from a geo-echo backend.
func geoEdge(t *testing.T, host string) *Edge {
	t.Helper()
	backend := geoEchoBackend(t)
	reg := openReg(t)
	if err := reg.Upsert(registry.Route{Host: host, Kind: "service", Addr: backend.Listener.Addr().String(), Status: "ready"}); err != nil {
		t.Fatal(err)
	}
	e, err := New(reg, nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(e.probers.Stop)
	return e
}

// spyGeo makes every lookup a hit for country and counts the calls.
func spyGeo(e *Edge, country string) *atomic.Int64 {
	var n atomic.Int64
	e.geoLookup = func(net.IP) geo.Snapshot {
		n.Add(1)
		return geo.Snapshot{Loaded: true, Found: true, Source: "TestDB-" + country,
			Result: geo.Result{Country: country, Point: geo.Point{Lat: 1.234, Lon: 5.678}, HasPoint: true}}
	}
	return &n
}

// serveGeo drives one request through ServeHTTP and returns the headers the
// app saw.
func serveGeo(t *testing.T, e *Edge, r *http.Request) (int, map[string]string) {
	t.Helper()
	w := httptest.NewRecorder()
	e.ServeHTTP(w, r)
	got := map[string]string{}
	if w.Code == http.StatusOK {
		if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil {
			t.Fatalf("backend body %q: %v", w.Body.String(), err)
		}
	}
	return w.Code, got
}

func geoReq(host, from string) *http.Request {
	r := httptest.NewRequest("GET", "http://"+host+"/", nil)
	r.RemoteAddr = net.JoinHostPort(from, "4444")
	return r
}

func TestDirectRequestGetsGeoHeadersFromTheRealReader(t *testing.T) {
	e := geoEdge(t, "app.test")
	db, err := geo.OpenMMDB("../geo/testdata/GeoIP2-City-Test.mmdb")
	if err != nil {
		t.Fatal(err)
	}
	e.geo.Swap(db)
	code, h := serveGeo(t, e, geoReq("app.test", "81.2.69.142"))
	if code != 200 || h[geoip.HeaderCountry] != "GB" || h[geoip.HeaderSubdivision] != "GB-ENG" ||
		h[geoip.HeaderSource] != "GeoIP2-City" || h[geoip.HeaderLat] != "51.51" || h[geoip.HeaderLon] != "-0.09" {
		t.Fatalf("app saw %v", h)
	}
	if _, leaked := h[geoip.HeaderHop]; leaked {
		t.Fatal("X-Carlos-Hop-Geo reached the app")
	}
}

func TestNoDatabaseMeansNoHeaders(t *testing.T) {
	e := geoEdge(t, "app.test")
	r := geoReq("app.test", "81.2.69.142")
	r.Header.Set(geoip.HeaderCountry, "KP")
	if _, h := serveGeo(t, e, r); len(h) != 0 {
		t.Fatalf("app saw %v with no database loaded", h)
	}
}

// Review Focus 1.
func TestForgedGeoHeadersNeverReachTheApp(t *testing.T) {
	e := geoEdge(t, "app.test")
	spyGeo(e, "GB")
	for name, mutate := range map[string]func(*http.Request){
		"plain forgery": func(r *http.Request) {
			r.Header.Set(geoip.HeaderCountry, "KP")
			r.Header.Set(geoip.HeaderSource, "forged")
			r.Header.Set(geoip.HeaderHop, "hit;src=forged")
		},
		"connection nomination": func(r *http.Request) {
			r.Header.Set("Connection", "X-Carlos-Geo-Country, X-Carlos-Geo-Source")
		},
	} {
		r := geoReq("app.test", "81.2.69.142")
		mutate(r)
		_, h := serveGeo(t, e, r)
		if h[geoip.HeaderCountry] != "GB" || h[geoip.HeaderSource] != "TestDB-GB" {
			t.Errorf("%s: app saw %v, want the edge's own GB answer", name, h)
		}
		if _, leaked := h[geoip.HeaderHop]; leaked {
			t.Errorf("%s: hop header reached the app", name)
		}
	}
}

// meshGeoEdge is geoEdge plus a minted mesh key, so hop requests can be
// signed (homemap_test.go's hopEdge, with a geo-echo backend).
func meshGeoEdge(t *testing.T, host string) (*Edge, []byte) {
	t.Helper()
	e := geoEdge(t, host)
	s := meshTestStore(t)
	if err := mintMeshKey(context.Background(), s); err != nil {
		t.Fatal(err)
	}
	m := newHomeMap(s, "home")
	refreshHomeMap(t, m)
	e.homes.Store(m)
	key, _ := m.meshKeys()
	return e, key
}

func TestVerifiedHopAdoptsTheSendersAnswerWithoutLookingUp(t *testing.T) {
	e, key := meshGeoEdge(t, "app.test")
	spy := spyGeo(e, "XX")                  // a lookup here would show as XX
	r := geoReq("app.test", "198.51.100.7") // the sending edge's address
	r.Header.Set(hopHeader, "1")
	r.Header.Set(hopAuthHeader, signHop(key, "app.test", time.Now()))
	r.Header.Set("X-Forwarded-For", "81.2.69.142")
	r.Header.Set(geoip.HeaderHop, "hit;src=GeoLite2-City")
	r.Header.Set(geoip.HeaderCountry, "FR")
	code, h := serveGeo(t, e, r)
	if code != 200 || h[geoip.HeaderCountry] != "FR" || h[geoip.HeaderSource] != "GeoLite2-City" {
		t.Fatalf("app saw %d %v, want the adopted FR", code, h)
	}
	if spy.Load() != 0 {
		t.Fatal("the home edge looked the visitor up instead of adopting")
	}
}

// Review Focus 2, codex R2-1.
func TestRejectedHopIsServedWithNoGeo(t *testing.T) {
	e, _ := meshGeoEdge(t, "app.test")
	spy := spyGeo(e, "XX")
	r := geoReq("app.test", "81.2.69.142")
	r.Header.Set(hopHeader, "1")
	r.Header.Set(hopAuthHeader, signHop([]byte("not-the-mesh-key-not-the-mesh-key"), "app.test", time.Now()))
	r.Header.Set(geoip.HeaderHop, "hit;src=GeoLite2-City")
	r.Header.Set(geoip.HeaderCountry, "FR")
	code, h := serveGeo(t, e, r)
	if code != 200 {
		t.Fatalf("a rejected hop must still be served (as today), got %d", code)
	}
	if len(h) != 0 || spy.Load() != 0 {
		t.Fatalf("rejected hop: app saw %v, lookups %d; want nothing and no lookup", h, spy.Load())
	}
}

func TestDecideGeoDoorIsOffWithoutALookup(t *testing.T) {
	e := geoEdge(t, "app.test")
	spy := spyGeo(e, "GB")
	a := e.decideGeo(geoReq("app.test", "81.2.69.142"), geoip.Captured{}, false, false, true)
	if a.State != geoip.StateOff || spy.Load() != 0 {
		t.Fatalf("door = %+v, lookups %d", a, spy.Load())
	}
}

func TestDoorCallReachesTheAppWithNoGeo(t *testing.T) {
	s := newDoorScene(t)
	spy := spyGeo(s.e, "GB")
	rec := s.do(t, "GET", doorOwnHost, "/", s.key, map[string]string{geoip.HeaderCountry: "KP"}, nil)
	if rec.Code != http.StatusOK {
		t.Fatalf("door call = %d %s", rec.Code, rec.Body.String())
	}
	if spy.Load() != 0 {
		t.Fatal("a door call was looked up: it would locate the box itself")
	}
	s.own.mu.Lock()
	hdr := s.own.hdr
	s.own.mu.Unlock()
	for _, k := range geoHeaderNames {
		if v := hdr.Get(k); v != "" {
			t.Errorf("door call reached the app with %s: %q", k, v)
		}
	}
}

func TestSetGeoForBackhaulCarriesStateAndAppsNeverSeeIt(t *testing.T) {
	in := httptest.NewRequest("GET", "http://app.test/", nil)
	for _, c := range []struct {
		a       geoip.Answer
		forward bool
		hop     string
	}{
		{geoip.Answer{State: geoip.StateHit, Country: "GB", Source: "S"}, true, "hit;src=S"},
		{geoip.Off(), true, "off"}, // a door call being backhauled
		{geoip.Answer{State: geoip.StateHit, Country: "GB", Source: "S"}, false, ""},
	} {
		pr := &httputil.ProxyRequest{In: in.WithContext(context.WithValue(in.Context(), geoResultKey{}, c.a)), Out: in.Clone(context.Background())}
		pr.Out.Header.Set(geoip.HeaderHop, "hit;src=forged")
		pr.Out.Header.Set(geoip.HeaderCountry, "KP")
		setGeo(pr, c.forward)
		if got := pr.Out.Header.Get(geoip.HeaderHop); got != c.hop {
			t.Errorf("%+v forward=%v: hop = %q, want %q", c.a, c.forward, got, c.hop)
		}
		if got := pr.Out.Header.Get(geoip.HeaderCountry); got == "KP" {
			t.Errorf("%+v: forged country survived setGeo", c.a)
		}
	}
	// A request that never passed ServeHTTP (the scheduler's UnixProxy) gets nothing.
	pr := &httputil.ProxyRequest{In: in, Out: in.Clone(context.Background())}
	pr.Out.Header.Set(geoip.HeaderCountry, "KP")
	setGeo(pr, false)
	if pr.Out.Header.Get(geoip.HeaderCountry) != "" {
		t.Fatal("no context answer must mean no geo headers")
	}
}

func TestBackhaulDeliversTheSendersAnswer(t *testing.T) {
	s := newBackhaulScene(t)
	s.converge(t)
	homeSpy, senderSpy := spyGeo(s.a.edge, "XX"), spyGeo(s.b.edge, "GB")
	_, body := s.get(t, s.b, "https://"+e2eHomedHost+"/", nil, true)
	seen := e2eParseApp(t, body)
	if seen.geoCountry != "GB" || seen.hopGeo != "" {
		t.Fatalf("home app saw country=%q hopGeo=%q, want GB and no hop header", seen.geoCountry, seen.hopGeo)
	}
	if homeSpy.Load() != 0 || senderSpy.Load() != 1 {
		t.Fatalf("lookups home=%d sender=%d, want 0 and 1 (ingress computes, spec §8.2)", homeSpy.Load(), senderSpy.Load())
	}
}

// Review Focus 1: Vary on the geo header picks the right variant.
func TestCacheVariesOnTheEdgesCountry(t *testing.T) {
	var hits atomic.Int64
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		w.Header().Set("Cache-Control", "max-age=60")
		w.Header().Set("Vary", geoip.HeaderCountry)
		w.Write([]byte("country=" + r.Header.Get(geoip.HeaderCountry)))
	}))
	t.Cleanup(backend.Close)
	reg := openReg(t)
	reg.Upsert(registry.Route{Host: "app.test", Kind: "instance", Addr: backend.Listener.Addr().String(), Status: "ready"})
	e, err := New(reg, nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(e.probers.Stop)
	e.EnableCache(NewCache(DefaultCacheMaxEntries, DefaultCacheMaxBodyBytes))
	db, _ := geo.OpenMMDB("../geo/testdata/GeoIP2-City-Test.mmdb")
	e.geo.Swap(db)
	get := func(from string, forged string) string {
		r := geoReq("app.test", from)
		if forged != "" {
			r.Header.Set(geoip.HeaderCountry, forged)
		}
		w := httptest.NewRecorder()
		e.ServeHTTP(w, r)
		return w.Body.String()
	}
	if b := get("81.2.69.142", ""); b != "country=GB" {
		t.Fatalf("GB = %q", b)
	}
	if b := get("216.160.83.56", ""); b != "country=US" {
		t.Fatalf("US = %q", b)
	}
	if b := get("81.2.69.142", "US"); b != "country=GB" {
		t.Fatalf("a forged header picked another country's variant: %q", b)
	}
	if hits.Load() != 2 || cachedVariants(e, "app.test", "/") != 2 {
		t.Fatalf("hits=%d variants=%d, want 2 and 2", hits.Load(), cachedVariants(e, "app.test", "/"))
	}
}

func TestUpgradeRequestCarriesGeo(t *testing.T) {
	// The backend answers the upgrade with the country it was sent, in a
	// 101 response header (ReverseProxy copies those back to the client).
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hj, ok := w.(http.Hijacker)
		if !ok {
			http.Error(w, "not hijackable", http.StatusInternalServerError)
			return
		}
		conn, buf, err := hj.Hijack()
		if err != nil {
			return
		}
		defer conn.Close()
		buf.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nX-Seen-Country: " +
			r.Header.Get(geoip.HeaderCountry) + "\r\n\r\n")
		buf.Flush()
	}))
	t.Cleanup(backend.Close)
	reg := openReg(t)
	if err := reg.Upsert(registry.Route{Host: "app.test", Kind: "service", Addr: backend.Listener.Addr().String(), Status: "ready"}); err != nil {
		t.Fatal(err)
	}
	e, err := New(reg, nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(e.probers.Stop)
	spyGeo(e, "GB")
	srv := httptest.NewServer(e)
	t.Cleanup(srv.Close)
	conn, err := net.Dial("tcp", srv.Listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	io.WriteString(conn, "GET / HTTP/1.1\r\nHost: app.test\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n"+
		"X-Carlos-Geo-Country: KP\r\n\r\n")
	resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
	if err != nil {
		t.Fatal(err)
	}
	if resp.StatusCode != http.StatusSwitchingProtocols || resp.Header.Get("X-Seen-Country") != "GB" {
		t.Fatalf("upgrade = %d, backend saw country %q, want 101 and GB", resp.StatusCode, resp.Header.Get("X-Seen-Country"))
	}
}

// Plan review R1-18 / spec §8.4: forgery, Connection nomination and an
// unknown X-Carlos-Geo-* suffix, through every proxy type a local route can
// use (unix, tcp, multi). Backhaul is covered by
// TestBackhaulDeliversTheSendersAnswer and the tunnel by Task 10's e2e.
func TestForgeryMatrixAcrossLocalProxyTypes(t *testing.T) {
	echo := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got := map[string]string{}
		for k := range r.Header {
			if n := strings.ToLower(strings.ReplaceAll(k, "_", "-")); strings.HasPrefix(n, "x-carlos-geo-") || n == strings.ToLower(geoip.HeaderHop) {
				got[k] = r.Header.Get(k)
			}
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(got)
	})
	tcp := httptest.NewServer(echo)
	t.Cleanup(tcp.Close)
	routes := map[string]registry.Route{
		"tcp":   {Host: "tcp.test", Kind: "service", Addr: tcp.Listener.Addr().String(), Status: "ready"},
		"multi": {Host: "multi.test", Kind: "service", Addrs: tcp.Listener.Addr().String(), Health: healthJSON(t, frozenHealth()), Status: "ready"},
	}
	sock := filepath.Join(t.TempDir(), "g.sock")
	if ln, err := net.Listen("unix", sock); err == nil {
		us := &httptest.Server{Listener: ln, Config: &http.Server{Handler: echo}}
		us.Start()
		t.Cleanup(us.Close)
		routes["unix"] = registry.Route{Host: "unix.test", Kind: "instance", Socket: sock, Status: "ready"}
	} else {
		t.Logf("unix listener unavailable here (%v); unix leg skipped", err)
	}
	reg := openReg(t)
	for _, rt := range routes {
		if err := reg.Upsert(rt); err != nil {
			t.Fatal(err)
		}
	}
	e, err := New(reg, nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(e.probers.Stop)
	spyGeo(e, "GB")
	for name, rt := range routes {
		r := geoReq(rt.Host, "81.2.69.142")
		r.Header.Set(geoip.HeaderCountry, "KP")
		r.Header.Set("X-Carlos-Geo-City", "Atlantis")
		r.Header.Set(geoip.HeaderHop, "hit;src=forged")
		r.Header.Set("Connection", "X-Carlos-Geo-Country, X-Carlos-Geo-Source")
		r.Header["X_Carlos_Geo_Country"] = []string{"KP"} // WSGI reads this as HTTP_X_CARLOS_GEO_COUNTRY
		r.Header["X_Carlos_Hop_Geo"] = []string{"hit;src=forged"}
		code, h := serveGeo(t, e, r)
		if code != 200 {
			t.Errorf("%s: status %d", name, code)
			continue
		}
		if h[geoip.HeaderCountry] != "GB" || h[geoip.HeaderSource] != "TestDB-GB" {
			t.Errorf("%s: app saw %v, want the edge's GB answer", name, h)
		}
		// Only the edge's own five, under their exact names, may arrive:
		// any other spelling of the namespace (a city, the hop header, an
		// underscore alias WSGI would read as the real one) is a forgery.
		allowed := map[string]bool{geoip.HeaderCountry: true, geoip.HeaderSubdivision: true,
			geoip.HeaderLat: true, geoip.HeaderLon: true, geoip.HeaderSource: true}
		for k := range h {
			if !allowed[k] {
				t.Errorf("%s: %s reached the app", name, k)
			}
		}
	}
}