package router

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

func serveWith(h Handler, remoteAddr, xff string) int {
	ok := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})
	req := httptest.NewRequest("GET", "http://example.com/", nil)
	req.RemoteAddr = remoteAddr
	if xff != "" {
		req.Header.Set("X-Forwarded-For", xff)
	}
	rec := httptest.NewRecorder()
	h(ok).ServeHTTP(rec, req)
	return rec.Code
}

func TestIpBlockList(t *testing.T) {
	block := IpBlockList("203.0.113.0/24", "198.51.100.7")

	if code := serveWith(block, "203.0.113.5:1234", ""); code != http.StatusForbidden {
		t.Errorf("CIDR match should be blocked, got %d", code)
	}
	if code := serveWith(block, "198.51.100.7:1234", ""); code != http.StatusForbidden {
		t.Errorf("bare IP match should be blocked, got %d", code)
	}
	if code := serveWith(block, "10.0.0.1:1234", ""); code != http.StatusOK {
		t.Errorf("non-match should pass, got %d", code)
	}
	// XFF is honored: RemoteAddr is a proxy, real client is blocked.
	if code := serveWith(block, "10.0.0.1:1234", "203.0.113.9"); code != http.StatusForbidden {
		t.Errorf("XFF client should be blocked, got %d", code)
	}
}

func TestIpAllowList(t *testing.T) {
	allow := IpAllowList("10.0.0.0/8")

	if code := serveWith(allow, "10.9.9.9:1234", ""); code != http.StatusOK {
		t.Errorf("in-range should pass, got %d", code)
	}
	if code := serveWith(allow, "203.0.113.1:1234", ""); code != http.StatusForbidden {
		t.Errorf("out-of-range should be blocked, got %d", code)
	}
	// Unparseable client IP is not on the allow list.
	if code := serveWith(allow, "garbage", ""); code != http.StatusForbidden {
		t.Errorf("unparseable IP should be blocked, got %d", code)
	}
}

func TestParsePrefixesPanicsOnGarbage(t *testing.T) {
	defer func() {
		if recover() == nil {
			t.Error("expected panic on invalid CIDR")
		}
	}()
	IpBlockList("not-an-ip")
}