package proxy

import (
	"testing"

	"github.com/golang-jwt/jwt/v5"

	"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
)

func handlerWith(allow bool, ids ...string) *Handler {
	return &Handler{config: &config.Config{Auth: config.AuthConfig{
		AllowListedMachineIdentities: allow,
		ZitadelClientIDs:             ids,
	}}}
}

func TestAllowListedMachineIdentityAdmitsANamedClientWithoutRoles(t *testing.T) {
	h := handlerWith(true, "poimen-llm")
	claims := jwt.MapClaims{"client_id": "poimen-llm", "sub": "397071051344839054"}
	if !h.allowListedMachineIdentity(claims) {
		t.Fatalf("an allow-listed machine client must be admitted when the token has no roles")
	}
}

func TestAllowListedMachineIdentityRefusesAnUnknownClient(t *testing.T) {
	h := handlerWith(true, "poimen-llm")
	claims := jwt.MapClaims{"client_id": "someone-else"}
	if h.allowListedMachineIdentity(claims) {
		t.Fatalf("a client that is not on the list must not be admitted")
	}
}

func TestAllowListedMachineIdentityRefusesWhenTheTokenStatesRoles(t *testing.T) {
	h := handlerWith(true, "poimen-llm")
	claims := jwt.MapClaims{
		"client_id": "poimen-llm",
		"urn:zitadel:iam:org:project:roles": map[string]interface{}{"memory:read": nil},
	}
	if h.allowListedMachineIdentity(claims) {
		t.Fatalf("a token with roles must be judged by the role check, not by the allow-list")
	}
}

func TestAllowListedMachineIdentityIsOffByDefault(t *testing.T) {
	h := handlerWith(false, "poimen-llm")
	claims := jwt.MapClaims{"client_id": "poimen-llm"}
	if h.allowListedMachineIdentity(claims) {
		t.Fatalf("without the setting the gateway must fall back to the role check only")
	}
}

func TestRolesClaimReadsBothShapes(t *testing.T) {
	asMap := jwt.MapClaims{"urn:zitadel:iam:org:project:roles": map[string]interface{}{"llm:inference": nil}}
	if got := rolesClaim(asMap); len(got) != 1 || got[0] != "llm:inference" {
		t.Fatalf("map shape must yield the role names, got %v", got)
	}
	asList := jwt.MapClaims{"roles": []interface{}{"llm:inference", "memory:read"}}
	if got := rolesClaim(asList); len(got) != 2 {
		t.Fatalf("list shape must yield every role, got %v", got)
	}
	if got := rolesClaim(jwt.MapClaims{}); len(got) != 0 {
		t.Fatalf("a token with no roles must yield nothing, got %v", got)
	}
}