package engine

import (
	"net"
	"testing"

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

	"github.com/bwmarrin/snowflake"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
)

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 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},
	}

	ipFlow := gopacket.NewFlow(layers.EndpointIPv4, net.IPv4(10, 0, 0, 1).To4(), net.IPv4(10, 0, 0, 2).To4())
	udp := &layers.UDP{
		SrcPort: 12345,
		DstPort: 53,
		BaseLayer: layers.BaseLayer{
			Payload: []byte("query"),
		},
	}
	ctx := &udpContext{Verdict: udpVerdictAccept}
	s := f.New(ipFlow, udp.TransportFlow(), udp, ctx)

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

	if !s.Accept(udp, false, ctx) {
		t.Fatalf("unexpected Accept=false for virgin stream")
	}
	s.Feed(udp, false, 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},
	}

	ipFlow := gopacket.NewFlow(layers.EndpointIPv4, net.IPv4(10, 0, 0, 1).To4(), net.IPv4(10, 0, 0, 2).To4())
	udp := &layers.UDP{
		SrcPort: 12345,
		DstPort: 53,
		BaseLayer: layers.BaseLayer{
			Payload: []byte("query"),
		},
	}

	ctx1 := &udpContext{Verdict: udpVerdictAccept}
	s := f.New(ipFlow, udp.TransportFlow(), udp, ctx1)
	if !s.Accept(udp, false, ctx1) {
		t.Fatalf("unexpected Accept=false before first feed")
	}
	s.Feed(udp, false, 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(udp, false, ctx2) {
		t.Fatalf("expected Accept=true after ruleset update")
	}
	s.Feed(udp, false, ctx2)
	if ctx2.Verdict != udpVerdictDropStream {
		t.Fatalf("verdict=%v want=%v", ctx2.Verdict, udpVerdictDropStream)
	}

	ctx3 := &udpContext{Verdict: udpVerdictAccept}
	if s.Accept(udp, 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)
	}
}