package aggregate

import (
	"bytes"
	"strings"
	"testing"

	"github.com/netsynth/netsynth/classify"
)

func TestPrintSummary(t *testing.T) {
	totals := map[classify.TrafficClass]int64{
		classify.ClassICMP:  10,
		classify.ClassDNS:   50,
		classify.ClassHTTPS: 100,
	}

	var buf bytes.Buffer
	PrintSummary(&buf, totals)
	output := buf.String()

	// Must contain class names.
	for _, class := range []string{"ICMP", "DNS", "HTTPS"} {
		if !strings.Contains(output, class) {
			t.Errorf("expected output to contain %q, got:\n%s", class, output)
		}
	}

	// Must contain counts.
	for _, count := range []string{"10", "50", "100"} {
		if !strings.Contains(output, count) {
			t.Errorf("expected output to contain count %q, got:\n%s", count, output)
		}
	}

	// Must contain percentage symbols.
	if !strings.Contains(output, "%") {
		t.Errorf("expected output to contain percentages, got:\n%s", output)
	}

	// Must contain TOTAL line.
	if !strings.Contains(output, "TOTAL") {
		t.Errorf("expected output to contain TOTAL line, got:\n%s", output)
	}

	// TOTAL should be 160 packets.
	if !strings.Contains(output, "160") {
		t.Errorf("expected output to contain total 160, got:\n%s", output)
	}

	// Must contain the section header.
	if !strings.Contains(output, "--- Protocol Summary ---") {
		t.Errorf("expected output to contain '--- Protocol Summary ---', got:\n%s", output)
	}
}

func TestPrintSummaryEmpty(t *testing.T) {
	var buf bytes.Buffer
	PrintSummary(&buf, map[classify.TrafficClass]int64{})
	output := buf.String()

	if !strings.Contains(output, "TOTAL") {
		t.Errorf("expected empty summary to contain TOTAL line, got:\n%s", output)
	}

	if !strings.Contains(output, "0") {
		t.Errorf("expected empty summary to contain 0 count, got:\n%s", output)
	}
}

func TestPrintSummarySorted(t *testing.T) {
	totals := map[classify.TrafficClass]int64{
		classify.ClassHTTPS:    100,
		classify.ClassDNS:      50,
		classify.ClassICMP:     10,
		classify.ClassSSH:      5,
		classify.ClassOtherTCP: 2,
	}

	var buf bytes.Buffer
	PrintSummary(&buf, totals)
	output := buf.String()

	// Extract lines containing class entries (not the header or TOTAL).
	lines := strings.Split(output, "\n")
	var classLines []string
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || strings.Contains(trimmed, "---") || strings.Contains(trimmed, "TOTAL") {
			continue
		}
		classLines = append(classLines, trimmed)
	}

	// Verify lines are in sorted order by comparing adjacent entries.
	for i := 1; i < len(classLines); i++ {
		// Each line starts with the class name. Extract first token.
		prevClass := strings.Fields(classLines[i-1])[0]
		curClass := strings.Fields(classLines[i])[0]
		if prevClass > curClass {
			t.Errorf("output is not sorted: %q appears before %q", prevClass, curClass)
		}
	}
}

func TestPrintWindowLine(t *testing.T) {
	snap := classify.WindowSnapshot{
		Counts: map[classify.TrafficClass]int64{
			classify.ClassDNS:   5,
			classify.ClassHTTPS: 10,
		},
		TotalPackets: 15,
		WindowIndex:  0,
	}

	var buf bytes.Buffer
	PrintWindowLine(&buf, snap)
	output := buf.String()

	if !strings.Contains(output, "[window 0]") {
		t.Errorf("expected output to contain '[window 0]', got: %q", output)
	}
	if !strings.Contains(output, "DNS:5") {
		t.Errorf("expected output to contain 'DNS:5', got: %q", output)
	}
	if !strings.Contains(output, "HTTPS:10") {
		t.Errorf("expected output to contain 'HTTPS:10', got: %q", output)
	}
	if !strings.Contains(output, "total: 15") {
		t.Errorf("expected output to contain 'total: 15', got: %q", output)
	}
	// Must be a single line ending with newline.
	trimmed := strings.TrimRight(output, "\n")
	if strings.Contains(trimmed, "\n") {
		t.Errorf("expected single line output, got multiple lines: %q", output)
	}
}

func TestPrintWindowLineEmpty(t *testing.T) {
	snap := classify.WindowSnapshot{
		Counts:       map[classify.TrafficClass]int64{},
		TotalPackets: 0,
		WindowIndex:  3,
	}

	var buf bytes.Buffer
	PrintWindowLine(&buf, snap)
	output := buf.String()

	if !strings.Contains(output, "[window 3]") {
		t.Errorf("expected output to contain '[window 3]', got: %q", output)
	}
	if !strings.Contains(output, "total: 0") {
		t.Errorf("expected output to contain 'total: 0', got: %q", output)
	}
}

func TestAccumulateTotals(t *testing.T) {
	totals := make(map[classify.TrafficClass]int64)

	snap1 := classify.WindowSnapshot{
		Counts: map[classify.TrafficClass]int64{
			classify.ClassICMP: 3,
			classify.ClassDNS:  7,
		},
		TotalPackets: 10,
		WindowIndex:  0,
	}
	snap2 := classify.WindowSnapshot{
		Counts: map[classify.TrafficClass]int64{
			classify.ClassDNS:   2,
			classify.ClassHTTPS: 5,
		},
		TotalPackets: 7,
		WindowIndex:  1,
	}

	AccumulateTotals(totals, snap1)
	AccumulateTotals(totals, snap2)

	if totals[classify.ClassICMP] != 3 {
		t.Errorf("expected ICMP=3, got %d", totals[classify.ClassICMP])
	}
	if totals[classify.ClassDNS] != 9 {
		t.Errorf("expected DNS=9, got %d", totals[classify.ClassDNS])
	}
	if totals[classify.ClassHTTPS] != 5 {
		t.Errorf("expected HTTPS=5, got %d", totals[classify.ClassHTTPS])
	}
}