package engine

import (
	"testing"

	"git.difuse.io/Difuse/Mellaris/analyzer"
	"git.difuse.io/Difuse/Mellaris/io"
	"git.difuse.io/Difuse/Mellaris/ruleset"

	"github.com/bwmarrin/snowflake"
)

type fixedRuleset struct {
	action ruleset.Action
}

func (r fixedRuleset) Analyzers(ruleset.StreamInfo) []analyzer.Analyzer {
	return nil
}

func (r fixedRuleset) Match(ruleset.StreamInfo) ruleset.MatchResult {
	return ruleset.MatchResult{Action: r.action}
}

type analyzerRuleset struct {
	action ruleset.Action
	ans    []analyzer.Analyzer
}

func (r analyzerRuleset) Analyzers(ruleset.StreamInfo) []analyzer.Analyzer {
	return r.ans
}

func (r analyzerRuleset) Match(ruleset.StreamInfo) ruleset.MatchResult {
	return ruleset.MatchResult{Action: r.action}
}

type countingTCPAnalyzer struct {
	newCalls  *int
	feedCalls *int
}

func (a countingTCPAnalyzer) Name() string { return "tls" }
func (a countingTCPAnalyzer) Limit() int   { return 0 }
func (a countingTCPAnalyzer) NewTCP(analyzer.TCPInfo, analyzer.Logger) analyzer.TCPStream {
	(*a.newCalls)++
	return countingTCPStream{feedCalls: a.feedCalls}
}

type countingTCPStream struct {
	feedCalls *int
}

func (s countingTCPStream) Feed(bool, bool, bool, int, []byte) (*analyzer.PropUpdate, bool) {
	(*s.feedCalls)++
	return nil, false
}

func (s countingTCPStream) Close(bool) *analyzer.PropUpdate {
	return nil
}

type logFinalizingRuleset struct {
	ans []analyzer.Analyzer
}

func (r logFinalizingRuleset) Analyzers(ruleset.StreamInfo) []analyzer.Analyzer {
	return r.ans
}

func (r logFinalizingRuleset) Match(info ruleset.StreamInfo) ruleset.MatchResult {
	if _, ok := info.Props["tls"]; ok {
		return ruleset.MatchResult{Action: ruleset.ActionMaybe, Logged: true}
	}
	return ruleset.MatchResult{Action: ruleset.ActionMaybe}
}

func (r logFinalizingRuleset) CanFinalizeAfterLog(ruleset.StreamInfo, []string) bool {
	return true
}

type requestPropTCPAnalyzer struct {
	closeCalls *int
}

func (a requestPropTCPAnalyzer) Name() string { return "tls" }
func (a requestPropTCPAnalyzer) Limit() int   { return 0 }
func (a requestPropTCPAnalyzer) NewTCP(analyzer.TCPInfo, analyzer.Logger) analyzer.TCPStream {
	return requestPropTCPStream{closeCalls: a.closeCalls}
}

type requestPropTCPStream struct {
	closeCalls *int
}

func (s requestPropTCPStream) Feed(bool, bool, bool, int, []byte) (*analyzer.PropUpdate, bool) {
	return &analyzer.PropUpdate{
		Type: analyzer.PropUpdateMerge,
		M:    analyzer.PropMap{"req": analyzer.PropMap{"sni": "good.example"}},
	}, false
}

func (s requestPropTCPStream) Close(bool) *analyzer.PropUpdate {
	(*s.closeCalls)++
	return nil
}

type noopTestLogger struct{}

func (noopTestLogger) WorkerStart(int) {}
func (noopTestLogger) WorkerStop(int)  {}

func (noopTestLogger) TCPStreamNew(int, ruleset.StreamInfo) {}
func (noopTestLogger) TCPStreamPropUpdate(ruleset.StreamInfo, bool) {
}
func (noopTestLogger) TCPStreamAction(ruleset.StreamInfo, ruleset.Action, bool) {
}

func (noopTestLogger) UDPStreamNew(int, ruleset.StreamInfo) {}
func (noopTestLogger) UDPStreamPropUpdate(ruleset.StreamInfo, bool) {
}
func (noopTestLogger) UDPStreamAction(ruleset.StreamInfo, ruleset.Action, bool) {
}

func (noopTestLogger) ModifyError(ruleset.StreamInfo, error) {}

func (noopTestLogger) AnalyzerDebugf(int64, string, string, ...interface{}) {}
func (noopTestLogger) AnalyzerInfof(int64, string, string, ...interface{})  {}
func (noopTestLogger) AnalyzerErrorf(int64, string, string, ...interface{}) {}

func TestUDPStreamUsesUpdatedRuleset(t *testing.T) {
	node, err := snowflake.NewNode(0)
	if err != nil {
		t.Fatalf("create node: %v", err)
	}
	f := &udpStreamFactory{
		WorkerID: 0,
		Logger:   noopTestLogger{},
		Node:     node,
		Ruleset:  fixedRuleset{action: ruleset.ActionAllow},
	}

	tuple := udpTupleKey{AIP: [16]byte{10, 0, 0, 1}, BIP: [16]byte{10, 0, 0, 2}, ALen: 4, BLen: 4, APort: 12345, BPort: 53}
	payload := []byte("query")
	ctx := &udpContext{Verdict: udpVerdictAccept}
	s := f.New(tuple, payload, ctx)

	if err := f.UpdateRuleset(fixedRuleset{action: ruleset.ActionBlock}); err != nil {
		t.Fatalf("update ruleset: %v", err)
	}

	if !s.Accept(false, ctx) {
		t.Fatalf("unexpected Accept=false for virgin stream")
	}
	s.Feed(false, payload, ctx)
	if ctx.Verdict != udpVerdictDropStream {
		t.Fatalf("verdict=%v want=%v", ctx.Verdict, udpVerdictDropStream)
	}
}

func TestUDPStreamReevaluatesAfterRulesetVersionChange(t *testing.T) {
	node, err := snowflake.NewNode(0)
	if err != nil {
		t.Fatalf("create node: %v", err)
	}
	f := &udpStreamFactory{
		WorkerID: 0,
		Logger:   noopTestLogger{},
		Node:     node,
		Ruleset:  fixedRuleset{action: ruleset.ActionAllow},
	}

	tuple := udpTupleKey{AIP: [16]byte{10, 0, 0, 1}, BIP: [16]byte{10, 0, 0, 2}, ALen: 4, BLen: 4, APort: 12345, BPort: 53}
	payload := []byte("query")

	ctx1 := &udpContext{Verdict: udpVerdictAccept}
	s := f.New(tuple, payload, ctx1)
	if !s.Accept(false, ctx1) {
		t.Fatalf("unexpected Accept=false before first feed")
	}
	s.Feed(false, payload, ctx1)
	if ctx1.Verdict != udpVerdictAcceptStream {
		t.Fatalf("verdict=%v want=%v", ctx1.Verdict, udpVerdictAcceptStream)
	}

	if err := f.UpdateRuleset(fixedRuleset{action: ruleset.ActionBlock}); err != nil {
		t.Fatalf("update ruleset: %v", err)
	}

	ctx2 := &udpContext{Verdict: udpVerdictAccept}
	if !s.Accept(false, ctx2) {
		t.Fatalf("expected Accept=true after ruleset update")
	}
	s.Feed(false, payload, ctx2)
	if ctx2.Verdict != udpVerdictDropStream {
		t.Fatalf("verdict=%v want=%v", ctx2.Verdict, udpVerdictDropStream)
	}

	ctx3 := &udpContext{Verdict: udpVerdictAccept}
	if s.Accept(false, ctx3) {
		t.Fatalf("expected Accept=false with unchanged ruleset and no active entries")
	}
	if ctx3.Verdict != udpVerdictDropStream {
		t.Fatalf("verdict=%v want=%v", ctx3.Verdict, udpVerdictDropStream)
	}
}

func TestTCPFlowUsesUpdatedRuleset(t *testing.T) {
	node, err := snowflake.NewNode(0)
	if err != nil {
		t.Fatalf("create node: %v", err)
	}
	mgr := newTCPFlowManager(0, noopTestLogger{}, nil, node, nil)
	mgr.updateRuleset(fixedRuleset{action: ruleset.ActionAllow}, 0)

	l3 := L3Info{
		Version:  4,
		Protocol: 6,
		SrcIP:    [4]byte{10, 0, 0, 1},
		DstIP:    [4]byte{10, 0, 0, 2},
	}
	tcp := TCPInfo{
		SrcPort: 12345,
		DstPort: 443,
		Seq:     100,
	}

	v := mgr.handle(1, l3, tcp, nil, nil, nil)
	if v != io.VerdictAcceptStream {
		t.Fatalf("first verdict=%v want=%v", v, io.VerdictAcceptStream)
	}

	mgr.updateRuleset(fixedRuleset{action: ruleset.ActionBlock}, 1)

	tcp2 := TCPInfo{
		SrcPort: 12345,
		DstPort: 443,
		Seq:     100,
	}
	v = mgr.handle(2, l3, tcp2, []byte("data"), nil, nil)
	if v != io.VerdictDropStream {
		t.Fatalf("verdict after update=%v want=%v", v, io.VerdictDropStream)
	}
}

func TestTCPFlowReevaluatesAfterRulesetVersionChange(t *testing.T) {
	node, err := snowflake.NewNode(0)
	if err != nil {
		t.Fatalf("create node: %v", err)
	}
	mgr := newTCPFlowManager(0, noopTestLogger{}, nil, node, nil)
	mgr.updateRuleset(fixedRuleset{action: ruleset.ActionAllow}, 0)

	l3 := L3Info{
		Version:  4,
		Protocol: 6,
		SrcIP:    [4]byte{10, 0, 0, 1},
		DstIP:    [4]byte{10, 0, 0, 2},
	}
	tcp := TCPInfo{
		SrcPort: 12345,
		DstPort: 443,
		Seq:     100,
	}

	v := mgr.handle(1, l3, tcp, nil, nil, nil)
	if v != io.VerdictAcceptStream {
		t.Fatalf("first verdict=%v want=%v", v, io.VerdictAcceptStream)
	}

	mgr.updateRuleset(fixedRuleset{action: ruleset.ActionBlock}, 1)

	tcp2 := TCPInfo{
		SrcPort: 12345,
		DstPort: 443,
		Seq:     100,
	}
	v = mgr.handle(2, l3, tcp2, []byte("data"), nil, nil)
	if v != io.VerdictDropStream {
		t.Fatalf("verdict after update=%v want=%v", v, io.VerdictDropStream)
	}

	tcp3 := TCPInfo{
		SrcPort: 12345,
		DstPort: 443,
		Seq:     104,
	}
	v = mgr.handle(1, l3, tcp3, nil, nil, nil)
	if v != io.VerdictDropStream {
		t.Fatalf("cached verdict after update=%v want=%v", v, io.VerdictDropStream)
	}
}

func TestTCPFlowDelaysAnalyzerCreationUntilPayload(t *testing.T) {
	node, err := snowflake.NewNode(0)
	if err != nil {
		t.Fatalf("create node: %v", err)
	}
	newCalls := 0
	feedCalls := 0
	mgr := newTCPFlowManager(0, noopTestLogger{}, nil, node, newAnalyzerSelector(AnalyzerSelectionModeSignature, &statsCounters{}))
	mgr.updateRuleset(analyzerRuleset{
		action: ruleset.ActionMaybe,
		ans: []analyzer.Analyzer{countingTCPAnalyzer{
			newCalls:  &newCalls,
			feedCalls: &feedCalls,
		}},
	}, 0)

	l3 := L3Info{
		Version:  4,
		Protocol: 6,
		SrcIP:    [4]byte{10, 0, 0, 1},
		DstIP:    [4]byte{10, 0, 0, 2},
	}
	tcp := TCPInfo{
		SrcPort: 12345,
		DstPort: 443,
		Seq:     100,
	}

	v := mgr.handle(1, l3, tcp, nil, nil, nil)
	if v != io.VerdictAccept {
		t.Fatalf("empty packet verdict=%v want=%v", v, io.VerdictAccept)
	}
	if newCalls != 0 || feedCalls != 0 {
		t.Fatalf("empty packet created/feed analyzer: new=%d feed=%d", newCalls, feedCalls)
	}

	tcp.Seq = 101
	v = mgr.handle(1, l3, tcp, []byte{0x16, 0x03, 0x01}, nil, nil)
	if v != io.VerdictAccept {
		t.Fatalf("payload verdict=%v want=%v", v, io.VerdictAccept)
	}
	if newCalls != 1 || feedCalls != 1 {
		t.Fatalf("payload should create/feed analyzer once: new=%d feed=%d", newCalls, feedCalls)
	}
}

func TestTCPFlowFinalizesAfterLogClassification(t *testing.T) {
	node, err := snowflake.NewNode(0)
	if err != nil {
		t.Fatalf("create node: %v", err)
	}
	closeCalls := 0
	mgr := newTCPFlowManager(0, noopTestLogger{}, nil, node, newAnalyzerSelector(AnalyzerSelectionModeSignature, &statsCounters{}))
	mgr.updateRuleset(logFinalizingRuleset{
		ans: []analyzer.Analyzer{requestPropTCPAnalyzer{closeCalls: &closeCalls}},
	}, 0)

	l3 := L3Info{
		Version:  4,
		Protocol: 6,
		SrcIP:    [4]byte{10, 0, 0, 1},
		DstIP:    [4]byte{10, 0, 0, 2},
	}
	tcp := TCPInfo{
		SrcPort: 12345,
		DstPort: 443,
		Seq:     100,
	}

	v := mgr.handle(1, l3, tcp, nil, nil, nil)
	if v != io.VerdictAccept {
		t.Fatalf("empty packet verdict=%v want=%v", v, io.VerdictAccept)
	}

	tcp.Seq = 101
	v = mgr.handle(1, l3, tcp, []byte{0x16, 0x03, 0x01}, nil, nil)
	if v != io.VerdictAcceptStream {
		t.Fatalf("payload verdict=%v want=%v", v, io.VerdictAcceptStream)
	}
	if closeCalls != 1 {
		t.Fatalf("expected analyzer to be closed once after finalization, got %d", closeCalls)
	}
	if _, ok := mgr.flows[1]; ok {
		t.Fatal("expected finalized TCP flow to be removed from manager")
	}
}