package main

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

func TestPromStoreIngestBytesBuckets(t *testing.T) {
	ps := NewPromStore()
	// 512 bytes: > 256, ≤ 1024 → bucket[0] stays 0, buckets[1..N] get 1
	ps.Ingest(LogRecord{Website: "example.com", Method: "GET", Status: "200", BytesSent: 512, SourceTag: "direct"})

	ps.mu.Lock()
	be := ps.bytesSent[hostSourceKey{"example.com", "direct"}]
	ps.mu.Unlock()

	if be == nil {
		t.Fatal("expected bytes entry, got nil")
	}
	if be.buckets[0] != 0 {
		t.Errorf("le=256 bucket = %d, want 0", be.buckets[0])
	}
	if be.buckets[1] != 1 {
		t.Errorf("le=1024 bucket = %d, want 1", be.buckets[1])
	}
	for i := 2; i <= promNumBodyBounds; i++ {
		if be.buckets[i] != 1 {
			t.Errorf("bucket[%d] = %d, want 1", i, be.buckets[i])
		}
	}
	if be.sum != 512 {
		t.Errorf("sum = %d, want 512", be.sum)
	}
}

func TestPromStoreIngestTimeBuckets(t *testing.T) {
	ps := NewPromStore()
	// 0.075s: > 0.05, ≤ 0.1
	ps.Ingest(LogRecord{Website: "example.com", Method: "GET", Status: "200", RequestTime: 0.075, SourceTag: "direct"})

	ps.mu.Lock()
	te := ps.requestDuration[hostSourceKey{"example.com", "direct"}]
	ps.mu.Unlock()

	if te == nil {
		t.Fatal("expected request_duration entry, got nil")
	}
	if te.buckets[3] != 0 {
		t.Errorf("le=0.05 bucket = %d, want 0", te.buckets[3])
	}
	if te.buckets[4] != 1 {
		t.Errorf("le=0.1 bucket = %d, want 1", te.buckets[4])
	}
	if te.buckets[promNumTimeBounds] != 1 {
		t.Errorf("+Inf bucket = %d, want 1", te.buckets[promNumTimeBounds])
	}
}

func TestPromStoreCounter(t *testing.T) {
	ps := NewPromStore()
	ps.Ingest(LogRecord{Website: "a.com", Method: "GET", Status: "200"})
	ps.Ingest(LogRecord{Website: "a.com", Method: "GET", Status: "200"})
	ps.Ingest(LogRecord{Website: "a.com", Method: "POST", Status: "201"})

	ps.mu.Lock()
	c1 := ps.counters[promCounterKey{"a.com", "GET", "200"}]
	c2 := ps.counters[promCounterKey{"a.com", "POST", "201"}]
	ps.mu.Unlock()

	if c1 != 2 {
		t.Errorf("GET/200 count = %d, want 2", c1)
	}
	if c2 != 1 {
		t.Errorf("POST/201 count = %d, want 1", c2)
	}
}

func TestPromStoreServeHTTP(t *testing.T) {
	ps := NewPromStore()
	ps.Ingest(LogRecord{
		Website: "example.com", Method: "GET", Status: "200",
		BytesSent: 100, RequestTime: 0.042, SourceTag: "direct",
	})

	req := httptest.NewRequest("GET", "/metrics", nil)
	rec := httptest.NewRecorder()
	ps.ServeHTTP(rec, req)

	body := rec.Body.String()

	checks := []string{
		"# TYPE nginx_http_requests_total counter",
		`nginx_http_requests_total{host="example.com",method="GET",status="200"} 1`,
		"# TYPE nginx_http_bytes_sent histogram",
		`nginx_http_bytes_sent_bucket{host="example.com",source_tag="direct",le="256"} 1`,
		`nginx_http_bytes_sent_count{host="example.com",source_tag="direct"} 1`,
		`nginx_http_bytes_sent_sum{host="example.com",source_tag="direct"} 100`,
		"# TYPE nginx_http_request_duration_seconds histogram",
		`nginx_http_request_duration_seconds_bucket{host="example.com",source_tag="direct",le="0.05"} 1`,
		`nginx_http_request_duration_seconds_count{host="example.com",source_tag="direct"} 1`,
		"# TYPE nginx_http_requests_by_source_total counter",
		`nginx_http_requests_by_source_total{source_tag="direct",status_class="2xx"} 1`,
	}
	for _, want := range checks {
		if !strings.Contains(body, want) {
			t.Errorf("missing %q in output:\n%s", want, body)
		}
	}
}

func TestPromStoreSourceTagAndStatusClass(t *testing.T) {
	ps := NewPromStore()
	ps.Ingest(LogRecord{Website: "h", Method: "GET", Status: "200", BytesSent: 100, SourceTag: "direct"})
	ps.Ingest(LogRecord{Website: "h", Method: "GET", Status: "200", BytesSent: 300, SourceTag: "cdn"})
	ps.Ingest(LogRecord{Website: "h", Method: "GET", Status: "404", BytesSent: 100, SourceTag: "cdn"})
	ps.Ingest(LogRecord{Website: "h", Method: "GET", Status: "500", BytesSent: 50, SourceTag: "cdn"})

	req := httptest.NewRequest("GET", "/metrics", nil)
	rec := httptest.NewRecorder()
	ps.ServeHTTP(rec, req)
	body := rec.Body.String()

	checks := []string{
		"# TYPE nginx_http_requests_by_source_total counter",
		`nginx_http_requests_by_source_total{source_tag="direct",status_class="2xx"} 1`,
		`nginx_http_requests_by_source_total{source_tag="cdn",status_class="2xx"} 1`,
		`nginx_http_requests_by_source_total{source_tag="cdn",status_class="4xx"} 1`,
		`nginx_http_requests_by_source_total{source_tag="cdn",status_class="5xx"} 1`,
		// dual-labeled histogram subsumes the old by-source bytes metric
		`nginx_http_bytes_sent_sum{host="h",source_tag="direct"} 100`,
		`nginx_http_bytes_sent_sum{host="h",source_tag="cdn"} 450`,
		// host counter unchanged
		`nginx_http_requests_total{host="h",method="GET",status="200"} 2`,
		`nginx_http_requests_total{host="h",method="GET",status="404"} 1`,
		`nginx_http_requests_total{host="h",method="GET",status="500"} 1`,
	}
	for _, want := range checks {
		if !strings.Contains(body, want) {
			t.Errorf("missing %q in output:\n%s", want, body)
		}
	}
}

func TestPromStoreUpstreamMetrics(t *testing.T) {
	ps := NewPromStore()
	// One upstream-served 200, one upstream-served 502, one no-upstream 200.
	ps.Ingest(LogRecord{
		Website: "h", Method: "GET", Status: "200", BytesSent: 100, SourceTag: "cdn",
		RequestLength:        500,
		HasUpstream:          true,
		UpstreamResponseTime: 0.020,
		UpstreamStatus:       "200",
	})
	ps.Ingest(LogRecord{
		Website: "h", Method: "GET", Status: "502", BytesSent: 100, SourceTag: "cdn",
		RequestLength:        500,
		HasUpstream:          true,
		UpstreamResponseTime: 0.150,
		UpstreamStatus:       "502",
	})
	ps.Ingest(LogRecord{
		Website: "h", Method: "GET", Status: "200", BytesSent: 100, SourceTag: "direct",
		RequestLength: 200,
	})

	req := httptest.NewRequest("GET", "/metrics", nil)
	rec := httptest.NewRecorder()
	ps.ServeHTTP(rec, req)
	body := rec.Body.String()

	checks := []string{
		// upstream counter: only the two upstream-served requests
		`nginx_http_upstream_requests_total{host="h",source_tag="cdn",status_class="2xx"} 1`,
		`nginx_http_upstream_requests_total{host="h",source_tag="cdn",status_class="5xx"} 1`,
		// upstream duration histogram only for upstream-served requests
		`nginx_http_upstream_duration_seconds_count{host="h",source_tag="cdn"} 2`,
		// request_bytes only observed when RequestLength > 0 — all three here
		`nginx_http_request_bytes_count{host="h",source_tag="cdn"} 2`,
		`nginx_http_request_bytes_count{host="h",source_tag="direct"} 1`,
	}
	for _, want := range checks {
		if !strings.Contains(body, want) {
			t.Errorf("missing %q in output:\n%s", want, body)
		}
	}

	// no-upstream request must not bump upstream counters
	if strings.Contains(body, `source_tag="direct",status_class=`) &&
		strings.Contains(body, "nginx_http_upstream_requests_total") {
		// Better check: ensure no direct-tagged upstream entry
		for _, bad := range []string{
			`nginx_http_upstream_requests_total{host="h",source_tag="direct"`,
			`nginx_http_upstream_duration_seconds_bucket{host="h",source_tag="direct"`,
		} {
			if strings.Contains(body, bad) {
				t.Errorf("unexpected %q in output:\n%s", bad, body)
			}
		}
	}
}

func TestPromStoreRequestBytesSkippedWhenZero(t *testing.T) {
	ps := NewPromStore()
	// v1 record — RequestLength=0, so request_bytes histogram should be empty.
	ps.Ingest(LogRecord{Website: "h", Method: "GET", Status: "200", BytesSent: 100, SourceTag: "cdn"})

	req := httptest.NewRequest("GET", "/metrics", nil)
	rec := httptest.NewRecorder()
	ps.ServeHTTP(rec, req)
	body := rec.Body.String()

	if strings.Contains(body, "nginx_http_request_bytes_bucket{") {
		t.Errorf("expected no request_bytes series for v1 record:\n%s", body)
	}
}

func TestPromStoreUDPCounters(t *testing.T) {
	ps := NewPromStore()
	ps.IncUDPPacket()
	ps.IncUDPPacket()
	ps.IncUDPPacket()
	ps.IncUDPSuccess()
	ps.IncUDPSuccess()
	ps.IncUDPConsumed()

	req := httptest.NewRequest("GET", "/metrics", nil)
	rec := httptest.NewRecorder()
	ps.ServeHTTP(rec, req)
	body := rec.Body.String()

	checks := []string{
		"logtail_udp_packets_received_total 3",
		"logtail_udp_loglines_success_total 2",
		"logtail_udp_loglines_consumed_total 1",
	}
	for _, want := range checks {
		if !strings.Contains(body, want) {
			t.Errorf("missing %q in output:\n%s", want, body)
		}
	}
}

func TestPromStoreCounterCap(t *testing.T) {
	ps := NewPromStore()
	for i := 0; i < promCounterCap+10; i++ {
		host := strings.Repeat("x", i%10+1) + ".com"
		status := "200"
		if i%3 == 0 {
			status = "404"
		}
		ps.Ingest(LogRecord{Website: host, Method: "GET", Status: status})
	}
	ps.mu.Lock()
	n := len(ps.counters)
	ps.mu.Unlock()
	if n > promCounterCap {
		t.Errorf("counter map size %d exceeds cap %d", n, promCounterCap)
	}
}

func TestStatusClass(t *testing.T) {
	cases := map[string]string{
		"200": "2xx", "201": "2xx", "299": "2xx",
		"301": "3xx",
		"404": "4xx", "418": "4xx",
		"500": "5xx", "504": "5xx",
		"100": "other", "": "other", "abc": "other",
	}
	for in, want := range cases {
		if got := statusClass(in); got != want {
			t.Errorf("statusClass(%q) = %q, want %q", in, got, want)
		}
	}
}