package service

import (
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/oschwald/maxminddb-golang"

	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
	"github.com/mhsanaei/3x-ui/v3/internal/web/service/probe"
)

// Q7 node country flags: local MMDB resolution on the master. The
// known-address expectations below are calibrated against the
// embedded factory snapshot (DB-IP Lite country database, retrieved
// 2026-10-24; provenance in geoipdata/NOTICE): these are long-lived
// anycast/ISP addresses whose country attribution is stable across
// editions — 8.8.8.8 (Google DNS) US, 223.5.5.5 (AliDNS) CN,
// 1.1.1.1 (Cloudflare/APNIC) AU, 2400:3200::1 (AliDNS v6) CN.

// useEmbeddedGeoIP swaps the embedded snapshot in as the live
// database for one test, restoring whatever was live before.
func useEmbeddedGeoIP(t *testing.T) {
	t.Helper()
	prev := currentGeoIP() // fire the lazy load so the Once is spent
	r, err := maxminddb.FromBytes(geoipCountryEmbedded)
	if err != nil {
		t.Fatalf("embedded snapshot will not open: %v", err)
	}
	swapGeoIP(&geoipDB{reader: r, origin: "embedded", cache: map[string]string{}})
	t.Cleanup(func() { swapGeoIP(prev) })
}

func TestCountryCodeForIPKnownAddresses(t *testing.T) {
	useEmbeddedGeoIP(t)
	for ip, want := range map[string]string{
		"8.8.8.8":      "US",
		"223.5.5.5":    "CN",
		"1.1.1.1":      "AU",
		"2400:3200::1": "CN",
		" 8.8.8.8 ":    "US", // surrounding space is trimmed
	} {
		if got := countryCodeForIP(ip); got != want {
			t.Errorf("countryCodeForIP(%q) = %q, want %q", ip, got, want)
		}
	}
	// Non-public and unlisted addresses have no country.
	for _, ip := range []string{
		"192.168.0.1", // RFC1918
		"10.0.0.5",    // RFC1918
		"127.0.0.1",   // loopback
		"100.64.0.1",  // CGNAT shared range
		"203.0.113.9", // public documentation range: unlisted in the DB
		"",            // empty
		"example.com", // not a literal
		"999.1.1.1",   // not an address
	} {
		if got := countryCodeForIP(ip); got != "" {
			t.Errorf("countryCodeForIP(%q) = %q, want empty", ip, got)
		}
	}
}

func TestProbeCountryAutoSourceOrder(t *testing.T) {
	useEmbeddedGeoIP(t)
	// Held v4 leads over held v6 (hub order).
	a := &model.AgentNode{IPv4: "8.8.8.8", IPv6: "2400:3200::1", SeenIPv4: "223.5.5.5"}
	if got := probeCountryAuto(a); got != "US" {
		t.Errorf("held v4 first: got %q, want US", got)
	}
	// A private held v4 is skipped; the held public v6 beats the
	// observed exit v4.
	a = &model.AgentNode{IPv4: "10.0.0.5", IPv6: "2400:3200::1", SeenIPv4: "8.8.8.8"}
	if got := probeCountryAuto(a); got != "CN" {
		t.Errorf("held v6 over exit v4: got %q, want CN", got)
	}
	// Exit addresses are the fallback when the node holds none.
	a = &model.AgentNode{SeenIPv4: "223.5.5.5"}
	if got := probeCountryAuto(a); got != "CN" {
		t.Errorf("exit v4 fallback: got %q, want CN", got)
	}
	a = &model.AgentNode{SeenIPv6: "2400:3200::1"}
	if got := probeCountryAuto(a); got != "CN" {
		t.Errorf("exit v6 fallback: got %q, want CN", got)
	}
	// No public address anywhere: no country.
	a = &model.AgentNode{IPv4: "192.168.1.2", SeenIPv4: "10.1.1.1"}
	if got := probeCountryAuto(a); got != "" {
		t.Errorf("private only: got %q, want empty", got)
	}
}

func TestProbeCountryPinWins(t *testing.T) {
	useEmbeddedGeoIP(t)
	a := &model.AgentNode{IPv4: "223.5.5.5"}
	country, auto := probeCountry("JP", a)
	if country != "JP" || auto != "CN" {
		t.Errorf("pin set: got (%q, %q), want (JP, CN)", country, auto)
	}
	country, auto = probeCountry("", a)
	if country != "CN" || auto != "CN" {
		t.Errorf("no pin: got (%q, %q), want (CN, CN)", country, auto)
	}
	country, auto = probeCountry("JP", &model.AgentNode{})
	if country != "JP" || auto != "" {
		t.Errorf("pin without address: got (%q, %q), want (JP, empty)", country, auto)
	}
}

func TestBuildProbeNodeViewCountry(t *testing.T) {
	useEmbeddedGeoIP(t)
	src := probeNodeSource{
		name:  "flagged",
		agent: model.AgentNode{IPv4: "8.8.8.8"},
		meta:  probe.Meta{Public: true},
	}
	anon := buildProbeNodeView(src, false)
	if anon["country"] != "US" {
		t.Errorf("anonymous country = %v, want US", anon["country"])
	}
	admin := buildProbeNodeView(src, true)
	if admin["country"] != "US" || admin["country_auto"] != "US" || admin["country_pin"] != "" {
		t.Errorf("admin pair = (%v, %v, pin %v), want (US, US, empty)",
			admin["country"], admin["country_auto"], admin["country_pin"])
	}
	src.meta.CountryPin = "DE"
	if got := buildProbeNodeView(src, false)["country"]; got != "DE" {
		t.Errorf("pinned country = %v, want DE", got)
	}
	if got := buildProbeNodeView(src, true)["country_auto"]; got != "US" {
		t.Errorf("country_auto under a pin = %v, want US", got)
	}
	// A node with no public address keeps the pre-flag empty country.
	src = probeNodeSource{name: "dark", agent: model.AgentNode{IPv4: "192.168.0.9"}, meta: probe.Meta{Public: true}}
	if got := buildProbeNodeView(src, false)["country"]; got != "" {
		t.Errorf("private node country = %v, want empty", got)
	}
}

func TestOpenGeoIPFileValidation(t *testing.T) {
	dir := t.TempDir()
	bad := filepath.Join(dir, "bad.mmdb")
	if err := os.WriteFile(bad, []byte("this is not an mmdb"), 0o644); err != nil {
		t.Fatal(err)
	}
	if _, err := openGeoIPFile(bad); err == nil {
		t.Fatal("garbage file opened as a country database")
	}
	good := filepath.Join(dir, "good.mmdb")
	if err := os.WriteFile(good, geoipCountryEmbedded, 0o644); err != nil {
		t.Fatal(err)
	}
	r, err := openGeoIPFile(good)
	if err != nil {
		t.Fatalf("embedded bytes as a file rejected: %v", err)
	}
	defer func() { _ = r.Close() }()
	db := &geoipDB{reader: r, cache: map[string]string{}}
	if got := db.lookup("223.5.5.5"); got != "CN" {
		t.Errorf("file-backed lookup = %q, want CN", got)
	}
}

// ---- the weekly updater's download/validate/replace core ----

func TestDownloadGeoIPCountryUpdates(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("ETag", `"db-v1"`)
		_, _ = w.Write(geoipCountryEmbedded)
	}))
	defer srv.Close()

	dest := filepath.Join(t.TempDir(), "geoip-country.mmdb")
	res, etag, err := downloadGeoIPCountry(srv.Client(), srv.URL, dest, "", time.Time{})
	if err != nil {
		t.Fatalf("download: %v", err)
	}
	if res != geoipFetchUpdated || etag != `"db-v1"` {
		t.Fatalf("got (res %v, etag %q), want (updated, \"db-v1\")", res, etag)
	}
	fi, err := os.Stat(dest)
	if err != nil || fi.Size() != int64(len(geoipCountryEmbedded)) {
		t.Fatalf("dest stat = %v (err %v), want %d bytes", fi, err, len(geoipCountryEmbedded))
	}
	// The replaced file is a working database.
	r, err := openGeoIPFile(dest)
	if err != nil {
		t.Fatalf("replaced file will not open: %v", err)
	}
	_ = r.Close()
}

func TestDownloadGeoIPCountryNotModified(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("If-None-Match") == `"db-v1"` || r.Header.Get("If-Modified-Since") != "" {
			w.WriteHeader(http.StatusNotModified)
			return
		}
		_, _ = w.Write(geoipCountryEmbedded)
	}))
	defer srv.Close()

	dir := t.TempDir()
	dest := filepath.Join(dir, "geoip-country.mmdb")
	if err := os.WriteFile(dest, []byte("ORIGINAL"), 0o644); err != nil {
		t.Fatal(err)
	}
	// ETag path.
	res, _, err := downloadGeoIPCountry(srv.Client(), srv.URL, dest, `"db-v1"`, time.Time{})
	if err != nil || res != geoipFetchUnchanged {
		t.Fatalf("etag path: got (res %v, err %v), want unchanged", res, err)
	}
	// If-Modified-Since path.
	res, _, err = downloadGeoIPCountry(srv.Client(), srv.URL, dest, "", time.Now())
	if err != nil || res != geoipFetchUnchanged {
		t.Fatalf("modified-since path: got (res %v, err %v), want unchanged", res, err)
	}
	body, _ := os.ReadFile(dest)
	if string(body) != "ORIGINAL" {
		t.Fatalf("dest rewritten on a 304: %q", body)
	}
}

func TestDownloadGeoIPCountryRejectsBad(t *testing.T) {
	dir := t.TempDir()
	dest := filepath.Join(dir, "geoip-country.mmdb")
	if err := os.WriteFile(dest, []byte("ORIGINAL"), 0o644); err != nil {
		t.Fatal(err)
	}
	assertKept := func(name string) {
		t.Helper()
		body, _ := os.ReadFile(dest)
		if string(body) != "ORIGINAL" {
			t.Fatalf("%s: dest replaced despite the failure: %q", name, body)
		}
	}

	garbage := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte("definitely not an mmdb, just bytes"))
	}))
	defer garbage.Close()
	if _, _, err := downloadGeoIPCountry(garbage.Client(), garbage.URL, dest, "", time.Time{}); err == nil {
		t.Fatal("garbage download accepted")
	}
	assertKept("garbage")

	broken := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	}))
	defer broken.Close()
	if _, _, err := downloadGeoIPCountry(broken.Client(), broken.URL, dest, "", time.Time{}); err == nil {
		t.Fatal("500 download accepted")
	}
	assertKept("500")
}