package ruleset

import (
	"reflect"
	"strings"
	"testing"

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

	"github.com/expr-lang/expr/ast"
	"github.com/expr-lang/expr/parser"
)

func TestExtractGeoSiteConditions(t *testing.T) {
	expression := `
		(geosite(tls.req.sni, "openai") || geosite(quic.req.sni, "OpenAI")) &&
		geosite(http.req.headers.host, "google@ads")
	`
	got := extractGeoSiteConditions(expression)
	want := []string{"openai", "google@ads"}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("extractGeoSiteConditions() = %v, want %v", got, want)
	}
}

func TestExtractGeoSiteHostCandidates(t *testing.T) {
	info := StreamInfo{
		Props: analyzer.CombinedPropMap{
			"quic": analyzer.PropMap{
				"req": analyzer.PropMap{
					"sni": "ChatGPT.com",
				},
			},
			"http": analyzer.PropMap{
				"req": analyzer.PropMap{
					"headers": analyzer.PropMap{
						"host": "api.openai.com:443",
					},
				},
			},
			"dns": analyzer.PropMap{
				"questions": []analyzer.PropMap{
					{"name": "chatgpt.com."},
					{"name": "8.8.8.8"},
				},
			},
		},
	}
	got := extractGeoSiteHostCandidates(info)
	want := []string{"chatgpt.com", "api.openai.com"}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("extractGeoSiteHostCandidates() = %v, want %v", got, want)
	}
}

func TestMatchGeoSiteConditions(t *testing.T) {
	hosts := []string{"chatgpt.com", "api.openai.com"}
	conditions := []string{" openai ", "google", "OPENAI"}
	got := matchGeoSiteConditions(hosts, conditions, func(site, condition string) bool {
		if condition != "openai" {
			return false
		}
		return site == "chatgpt.com" || site == "api.openai.com"
	})
	want := []string{"openai"}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("matchGeoSiteConditions() = %v, want %v", got, want)
	}
}

func TestIDPatcher_PatchesGeoSiteORChainToGeoSiteSet(t *testing.T) {
	tree, err := parser.Parse(`geosite(tls.req.sni, "google") || geosite(tls.req.sni, "youtube") || geosite(tls.req.sni, "openai")`)
	if err != nil {
		t.Fatalf("parse expression: %v", err)
	}
	root := tree.Node
	patcher := &idPatcher{GeoMatcher: geo.NewGeoMatcher("", "")}
	ast.Walk(&root, patcher)
	if patcher.Err != nil {
		t.Fatalf("patch error: %v", patcher.Err)
	}
	got := root.String()
	if !strings.Contains(got, "geosite_set(") {
		t.Fatalf("expected geosite_set rewrite, got %q", got)
	}
	if strings.Contains(got, "||") || strings.Contains(got, " or ") {
		t.Fatalf("expected OR chain to be collapsed, got %q", got)
	}
}