package engine

import (
	"net"
	"testing"
	"time"

	"github.com/apernet/OpenGFW/analyzer"
	"github.com/apernet/OpenGFW/analyzer/tcp"
	"github.com/apernet/OpenGFW/analyzer/udp"
	"github.com/apernet/OpenGFW/io"
	"github.com/apernet/OpenGFW/ruleset"
	"github.com/apernet/OpenGFW/ruleset/builtins/geo"

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

type nullEngineLogger struct{}

func (nullEngineLogger) WorkerStart(int)                                          {}
func (nullEngineLogger) WorkerStop(int)                                           {}
func (nullEngineLogger) TCPStreamNew(int, ruleset.StreamInfo)                     {}
func (nullEngineLogger) TCPStreamPropUpdate(ruleset.StreamInfo, bool)             {}
func (nullEngineLogger) TCPStreamAction(ruleset.StreamInfo, ruleset.Action, bool) {}
func (nullEngineLogger) TCPFlush(int, int, int)                                   {}
func (nullEngineLogger) UDPStreamNew(int, ruleset.StreamInfo)                     {}
func (nullEngineLogger) UDPStreamPropUpdate(ruleset.StreamInfo, bool)             {}
func (nullEngineLogger) UDPStreamAction(ruleset.StreamInfo, ruleset.Action, bool) {}
func (nullEngineLogger) ModifyError(ruleset.StreamInfo, error)                    {}
func (nullEngineLogger) AnalyzerDebugf(int64, string, string, ...interface{})     {}
func (nullEngineLogger) AnalyzerInfof(int64, string, string, ...interface{})      {}
func (nullEngineLogger) AnalyzerErrorf(int64, string, string, ...interface{})     {}

type nullRulesetLogger struct{}

func (nullRulesetLogger) Log(ruleset.StreamInfo, string)               {}
func (nullRulesetLogger) MatchError(ruleset.StreamInfo, string, error) {}

func testRuleset(t *testing.T, rules ...ruleset.ExprRule) ruleset.Ruleset {
	t.Helper()
	ans := []analyzer.Analyzer{
		&tcp.FETAnalyzer{}, &tcp.HTTPAnalyzer{}, &tcp.TLSAnalyzer{},
		&udp.QUICAnalyzer{}, &udp.WireGuardAnalyzer{},
	}
	rs, err := ruleset.CompileExprRules(rules, ans, nil, &ruleset.BuiltinConfig{Logger: nullRulesetLogger{}, GeoMatcher: geo.NewGeoMatcher("", "")})
	if err != nil {
		t.Fatal(err)
	}
	return rs
}

// capturingFactory 记下工厂创建的每条 tcpStream，让测试能检查判定之后它的状态
type capturingFactory struct {
	inner   *tcpStreamFactory
	streams []*tcpStream
}

func (f *capturingFactory) New(ipFlow, tcpFlow gopacket.Flow, tcp *layers.TCP, ac reassembly.AssemblerContext) reassembly.Stream {
	s := f.inner.New(ipFlow, tcpFlow, tcp, ac).(*tcpStream)
	f.streams = append(f.streams, s)
	return s
}

func newTestAssembler(t *testing.T, rs ruleset.Ruleset) (*reassembly.Assembler, *capturingFactory) {
	t.Helper()
	node, err := snowflake.NewNode(1)
	if err != nil {
		t.Fatal(err)
	}
	cf := &capturingFactory{inner: &tcpStreamFactory{WorkerID: 1, Logger: nullEngineLogger{}, Node: node, Ruleset: rs}}
	asm := reassembly.NewAssembler(reassembly.NewStreamPool(cf))
	asm.MaxBufferedPagesTotal = defaultTCPMaxBufferedPagesTotal
	asm.MaxBufferedPagesPerConnection = defaultTCPMaxBufferedPagesPerConnection
	return asm, cf
}

var testIPFlow = gopacket.NewFlow(layers.EndpointIPv4, net.IPv4(10, 0, 0, 1).To4(), net.IPv4(203, 0, 113, 1).To4())

func tcpSegment(seq uint32, syn bool, payload []byte) *layers.TCP {
	seg := &layers.TCP{SrcPort: 40000, DstPort: 80, Seq: seq, SYN: syn, ACK: !syn, PSH: len(payload) > 0, Window: 65535}
	seg.Payload = payload
	return seg
}

func feedTCP(asm *reassembly.Assembler, seg *layers.TCP, ts time.Time) tcpVerdict {
	ctx := &tcpContext{PacketMetadata: &gopacket.PacketMetadata{CaptureInfo: gopacket.CaptureInfo{Timestamp: ts}}, Verdict: tcpVerdictAccept}
	asm.AssembleWithContext(testIPFlow, seg, ctx)
	return ctx.Verdict
}

var httpRequest = []byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")

func TestTCPStreamReleasesStateAfterVerdict(t *testing.T) {
	rs := testRuleset(t, ruleset.ExprRule{Name: "block-http", Action: "block", Expr: "http?.req != nil"})
	asm, cf := newTestAssembler(t, rs)
	ts := time.Now()

	feedTCP(asm, tcpSegment(1000, true, nil), ts)
	if v := feedTCP(asm, tcpSegment(1001, false, httpRequest), ts); v != tcpVerdictDropStream {
		t.Fatalf("verdict = %v, want drop stream", v)
	}
	if len(cf.streams) != 1 {
		t.Fatalf("streams created = %d, want 1", len(cf.streams))
	}
	s := cf.streams[0]
	if !s.released || s.activeEntries != nil || s.info.Props != nil || s.ruleset != nil {
		t.Fatalf("stream still holds state after verdict: released=%v entries=%d props=%v ruleset=%v", s.released, len(s.activeEntries), s.info.Props, s.ruleset != nil)
	}

	// 后续报文沿用判定，且不能碰已经释放的状态
	if v := feedTCP(asm, tcpSegment(1001+uint32(len(httpRequest)), false, []byte("more")), ts); v != tcpVerdictDropStream {
		t.Fatalf("verdict on a later segment = %v, want drop stream", v)
	}
}

func TestTCPStreamReleasedWhenAllAnalyzersFinish(t *testing.T) {
	// 规则只引用 tls；一个 HTTP 请求让 tls 分析器直接放弃，没有裁决，引擎应当发 AcceptStream 并释放
	rs := testRuleset(t, ruleset.ExprRule{Name: "block-tls", Action: "block", Expr: "tls?.req != nil"})
	asm, cf := newTestAssembler(t, rs)
	ts := time.Now()

	feedTCP(asm, tcpSegment(1000, true, nil), ts)
	if v := feedTCP(asm, tcpSegment(1001, false, httpRequest), ts); v != tcpVerdictAcceptStream {
		t.Fatalf("verdict = %v, want accept stream", v)
	}
	if s := cf.streams[0]; !s.released || s.info.Props != nil {
		t.Fatalf("stream not released after all analyzers finished: released=%v props=%v", s.released, s.info.Props)
	}
}

// 判定之前缓冲的乱序页会在超时 flush 时被补送到 ReassembledSG，那时流已经释放，不能崩
func TestTCPStreamSurvivesFlushDeliveringBufferedPageAfterRelease(t *testing.T) {
	rs := testRuleset(t, ruleset.ExprRule{Name: "block-http", Action: "block", Expr: "http?.req != nil"})
	asm, cf := newTestAssembler(t, rs)
	ts := time.Now()

	feedTCP(asm, tcpSegment(1000, true, nil), ts)
	// 乱序段：序号跳过一大截，被缓冲起来等空洞补齐
	feedTCP(asm, tcpSegment(9000, false, []byte("out of order")), ts)
	if v := feedTCP(asm, tcpSegment(1001, false, httpRequest), ts); v != tcpVerdictDropStream {
		t.Fatalf("verdict = %v, want drop stream", v)
	}
	s := cf.streams[0]
	if !s.released {
		t.Fatal("stream not released after verdict")
	}

	// 超时 flush 会带着空洞把那一页送进来
	flushed, closed := asm.FlushCloseOlderThan(ts.Add(time.Hour))
	if flushed == 0 || closed == 0 {
		t.Fatalf("flush did nothing: flushed=%d closed=%d", flushed, closed)
	}
	if s.info.Props != nil || s.activeEntries != nil {
		t.Fatal("flush resurrected released state")
	}
}

func TestUDPStreamReleasesStateAfterVerdict(t *testing.T) {
	rs := testRuleset(t, ruleset.ExprRule{Name: "allow-443", Action: "allow", Expr: "port.dst == 443"})
	w, err := newWorker(workerConfig{ID: 1, Logger: nullEngineLogger{}, Ruleset: rs})
	if err != nil {
		t.Fatal(err)
	}
	pkt := &layers.UDP{SrcPort: 40000, DstPort: 443}
	pkt.Payload = []byte("hello")

	if v, _ := w.handleUDP(7, testIPFlow, pkt); v != io.VerdictAcceptStream {
		t.Fatalf("verdict = %v, want accept stream", v)
	}
	value, ok := w.udpStreamManager.streams.Get(7)
	if !ok {
		t.Fatal("stream missing from the LRU")
	}
	if s := value.Stream; s.activeEntries != nil || s.info.Props != nil || s.ruleset != nil {
		t.Fatalf("udp stream still holds state after verdict: entries=%d props=%v", len(s.activeEntries), s.info.Props)
	}
	if v, _ := w.handleUDP(7, testIPFlow, pkt); v != io.VerdictAcceptStream {
		t.Fatalf("verdict on a later packet = %v, want accept stream", v)
	}
}

func TestTCPFlushIntervalFollowsTimeout(t *testing.T) {
	cases := []struct {
		timeout, want time.Duration
	}{
		{30 * time.Second, 15 * time.Second},
		{10 * time.Minute, time.Minute},
		{time.Second, time.Second},
		{0, time.Second},
	}
	for _, c := range cases {
		if got := tcpFlushIntervalFor(c.timeout); got != c.want {
			t.Errorf("timeout %v: interval = %v, want %v", c.timeout, got, c.want)
		}
	}
}