package router

import (
	"testing"

	"somegit.dev/Owlibou/gnoma/internal/provider"
	"somegit.dev/Owlibou/gnoma/internal/security"
)

func TestParsePreferPolicy(t *testing.T) {
	cases := []struct {
		in      string
		want    PreferPolicy
		wantErr bool
	}{
		{"", PreferAuto, false},
		{"auto", PreferAuto, false},
		{"AUTO", PreferAuto, false},
		{"  auto  ", PreferAuto, false},
		{"local", PreferLocal, false},
		{"Local", PreferLocal, false},
		{"cloud", PreferCloud, false},
		{"prefer-cloud", PreferAuto, true},
		{"none", PreferAuto, true},
	}
	for _, tc := range cases {
		t.Run(tc.in, func(t *testing.T) {
			got, err := ParsePreferPolicy(tc.in)
			if (err != nil) != tc.wantErr {
				t.Fatalf("err=%v wantErr=%v", err, tc.wantErr)
			}
			if !tc.wantErr && got != tc.want {
				t.Errorf("got %v, want %v", got, tc.want)
			}
		})
	}
}

func TestPreferPolicy_String(t *testing.T) {
	cases := map[PreferPolicy]string{
		PreferAuto:  "auto",
		PreferLocal: "local",
		PreferCloud: "cloud",
	}
	for in, want := range cases {
		if got := in.String(); got != want {
			t.Errorf("%d.String() = %q, want %q", in, got, want)
		}
	}
}

func TestPolicyMultiplier(t *testing.T) {
	localArm := &Arm{IsLocal: true}
	cloudArm := &Arm{IsLocal: false}

	cases := []struct {
		name   string
		arm    *Arm
		policy PreferPolicy
		want   float64
	}{
		{"auto/local", localArm, PreferAuto, 1.0},
		{"auto/cloud", cloudArm, PreferAuto, 1.0},
		{"local/local", localArm, PreferLocal, 1.0},
		{"local/cloud", cloudArm, PreferLocal, 0.3},
		{"cloud/local", localArm, PreferCloud, 0.5},
		{"cloud/cloud", cloudArm, PreferCloud, 1.0},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if got := policyMultiplier(tc.arm, tc.policy); got != tc.want {
				t.Errorf("policyMultiplier(%+v, %v) = %v, want %v", tc.arm, tc.policy, got, tc.want)
			}
		})
	}
}

// TestPreferPolicy_RouterAcceptanceScenarios is the user-facing payoff:
// the prefer knob shifts arm tiers so the dispreferred camp is walked
// last. The test uses a task type that neither arm has in its Strengths
// list so the tier walk actually runs (the Strengths-promoted path
// bypasses tier ordering entirely).
//
// Arms are chosen to be in adjacent base tiers — a general-purpose
// local arm at tier 2 (no MaxComplexity, no family-defaults match) and
// a cloud arm at tier 3. The +2 tier shift then puts the dispreferred
// arm at tier 4 (local) or 5 (cloud), behind the preferred camp.
//
// The Strengths-promoted case (cost-amplification can overwhelm the
// within-tier multiplier) is covered separately by
// TestPreferPolicy_StrengthsBeatsMultiplier, which validates that a
// strongly-tagged arm wins regardless of prefer.
func TestPreferPolicy_RouterAcceptanceScenarios(t *testing.T) {
	makeRouter := func(policy PreferPolicy) *Router {
		r := New(Config{})
		r.SetPreferPolicy(policy)

		// Local arm: family doesn't match any defaults entry, so no
		// Strengths or MaxComplexity get attached — clean tier-2 arm.
		r.RegisterArm(&Arm{
			ID:        NewArmID("ollama", "novel-local-llm:7b"),
			ModelName: "novel-local-llm:7b",
			Provider:  security.WrapProvider(&stubProvider{name: "ollama", model: "novel-local-llm:7b"}, nil),
			IsLocal:   true,
			Capabilities: provider.Capabilities{
				ToolUse:       true,
				ContextWindow: 200000,
			},
		})

		// Cloud arm: also no family match (we use a deliberately
		// non-matching ID so Strengths defaults don't kick in).
		r.RegisterArm(&Arm{
			ID:        NewArmID("anthropic", "novel-cloud-model"),
			ModelName: "novel-cloud-model",
			Provider:  security.WrapProvider(&stubProvider{name: "anthropic", model: "novel-cloud-model"}, nil),
			IsLocal:   false,
			Capabilities: provider.Capabilities{
				ToolUse:       true,
				ContextWindow: 1_000_000,
				ThinkingModes: []provider.EffortLevel{provider.EffortMedium},
			},
		})
		return r
	}

	task := Task{
		Type:            TaskExplain,
		ComplexityScore: 0.5,
		Priority:        PriorityNormal,
		RequiresTools:   true,
		EstimatedTokens: 1500,
	}

	t.Run("prefer=local picks the local arm", func(t *testing.T) {
		r := makeRouter(PreferLocal)
		decision := r.Select(task)
		if decision.Error != nil {
			t.Fatalf("Select error: %v", decision.Error)
		}
		if !decision.Arm.IsLocal {
			t.Errorf("PreferLocal should pick local; got %s (IsLocal=%v)", decision.Arm.ID, decision.Arm.IsLocal)
		}
		decision.Rollback()
	})

	t.Run("prefer=cloud picks the cloud arm", func(t *testing.T) {
		r := makeRouter(PreferCloud)
		decision := r.Select(task)
		if decision.Error != nil {
			t.Fatalf("Select error: %v", decision.Error)
		}
		if decision.Arm.IsLocal {
			t.Errorf("PreferCloud should pick cloud; got %s (IsLocal=%v)", decision.Arm.ID, decision.Arm.IsLocal)
		}
		decision.Rollback()
	})

	t.Run("prefer=auto preserves tier order (local tier 2 < cloud tier 3)", func(t *testing.T) {
		r := makeRouter(PreferAuto)
		decision := r.Select(task)
		if decision.Error != nil {
			t.Fatalf("Select error: %v", decision.Error)
		}
		if !decision.Arm.IsLocal {
			t.Errorf("PreferAuto should preserve tier order (local wins); got %s", decision.Arm.ID)
		}
		decision.Rollback()
	})
}

// TestPreferPolicy_SLMStillWinsUnderPreferCloud documents the
// SLM-protection behavior: under PreferCloud, a tier-0 SLM (an arm
// with MaxComplexity > 0 that fits the task) still wins because the
// +2 tier shift only moves it from tier 0 to tier 2, which is still
// below the cloud arm's tier 3. This matches the plan's intent: "the
// SLM does small stuff" survives PreferCloud — that's exactly what
// the SLM is for.
func TestPreferPolicy_SLMStillWinsUnderPreferCloud(t *testing.T) {
	r := New(Config{})
	r.SetPreferPolicy(PreferCloud)

	// Tier-0 SLM (low MaxComplexity, fits the trivial task).
	r.RegisterArm(&Arm{
		ID:            NewArmID("ollama", "tiny-slm:1.5b"),
		ModelName:     "tiny-slm:1.5b",
		Provider:      security.WrapProvider(&stubProvider{name: "ollama", model: "tiny-slm:1.5b"}, nil),
		IsLocal:       true,
		MaxComplexity: 0.30,
		Strengths:     []TaskType{TaskBoilerplate},
		Capabilities: provider.Capabilities{
			ToolUse:       true,
			ContextWindow: 32768,
		},
	})
	r.RegisterArm(&Arm{
		ID:        NewArmID("anthropic", "claude-sonnet-4-6"),
		ModelName: "claude-sonnet-4-6",
		Provider:  security.WrapProvider(&stubProvider{name: "anthropic", model: "claude-sonnet-4-6"}, nil),
		IsLocal:   false,
		Capabilities: provider.Capabilities{
			ToolUse:       true,
			ContextWindow: 1_000_000,
		},
	})

	decision := r.Select(Task{
		Type:            TaskBoilerplate,
		ComplexityScore: 0.1,
		Priority:        PriorityLow,
		RequiresTools:   true,
		EstimatedTokens: 200,
	})
	if decision.Error != nil {
		t.Fatalf("Select error: %v", decision.Error)
	}
	if decision.Arm.ID != NewArmID("ollama", "tiny-slm:1.5b") {
		t.Errorf("SLM should win trivial task even under PreferCloud (tier 0+2=2 < cloud 3); got %s", decision.Arm.ID)
	}
	decision.Rollback()
}

// TestPreferPolicy_StrengthsBeatsMultiplier: a cloud arm with a strong
// task-type tag still wins over a local arm without that tag, even
// under PreferLocal. Strengths is the primary signal; prefer is a
// secondary multiplier within the promoted/tier set.
func TestPreferPolicy_StrengthsBeatsMultiplier(t *testing.T) {
	r := New(Config{})
	r.SetPreferPolicy(PreferLocal)

	// Local arm has no Strengths for SecurityReview.
	localArm := &Arm{
		ID:            NewArmID("ollama", "qwen3:14b"),
		ModelName:     "qwen3:14b",
		Provider:      security.WrapProvider(&stubProvider{name: "ollama", model: "qwen3:14b"}, nil),
		IsLocal:       true,
		Strengths:     []TaskType{TaskGeneration},
		MaxComplexity: 0.75,
		Capabilities: provider.Capabilities{
			ToolUse:       true,
			ContextWindow: 32768,
		},
	}
	cloudArm := &Arm{
		ID:        NewArmID("anthropic", "claude-opus-4-7"),
		ModelName: "claude-opus-4-7",
		Provider:  security.WrapProvider(&stubProvider{name: "anthropic", model: "claude-opus-4-7"}, nil),
		IsLocal:   false,
		Strengths: []TaskType{TaskSecurityReview, TaskPlanning},
		Capabilities: provider.Capabilities{
			ToolUse:       true,
			ContextWindow: 1_000_000,
			ThinkingModes: []provider.EffortLevel{provider.EffortHigh},
		},
	}
	r.RegisterArm(localArm)
	r.RegisterArm(cloudArm)

	decision := r.Select(Task{
		Type:            TaskSecurityReview,
		ComplexityScore: 0.8,
		Priority:        PriorityCritical,
		RequiresTools:   true,
		EstimatedTokens: 3000,
	})
	if decision.Error != nil {
		t.Fatalf("Select error: %v", decision.Error)
	}
	if decision.Arm.ID != cloudArm.ID {
		t.Errorf("Strengths-tagged cloud arm should beat PreferLocal multiplier; got %s", decision.Arm.ID)
	}
	decision.Rollback()
}

// TestPreferPolicy_ForcedArmBypassesPolicy: --provider X must always win.
func TestPreferPolicy_ForcedArmBypassesPolicy(t *testing.T) {
	r := New(Config{})
	r.SetPreferPolicy(PreferLocal)

	cloudArmID := NewArmID("anthropic", "claude-sonnet-4-6")
	r.RegisterArm(&Arm{
		ID:        cloudArmID,
		ModelName: "claude-sonnet-4-6",
		Provider:  security.WrapProvider(&stubProvider{name: "anthropic", model: "claude-sonnet-4-6"}, nil),
		IsLocal:   false,
		Capabilities: provider.Capabilities{
			ToolUse:       true,
			ContextWindow: 1_000_000,
		},
	})
	r.ForceArm(cloudArmID)

	decision := r.Select(Task{Type: TaskGeneration, RequiresTools: true})
	if decision.Error != nil {
		t.Fatalf("Select error: %v", decision.Error)
	}
	if decision.Arm.ID != cloudArmID {
		t.Errorf("forced arm should bypass PreferLocal; got %s, want %s", decision.Arm.ID, cloudArmID)
	}
}

// TestPreferPolicy_IncognitoStillWins: incognito's hard filter must
// dominate the soft prefer bias.
func TestPreferPolicy_IncognitoStillWins(t *testing.T) {
	r := New(Config{})
	r.SetPreferPolicy(PreferCloud) // bias toward cloud
	r.SetLocalOnly(true)           // but incognito filters cloud out

	factory := func(name, model string) SecureProvider {
		return security.WrapProvider(&stubProvider{name: name, model: model}, nil)
	}
	RegisterDiscoveredModels(r, []DiscoveredModel{
		{ID: "qwen3:14b", Provider: "ollama", SupportsTools: true, ContextSize: 32768},
	}, factory)
	r.RegisterArm(&Arm{
		ID:        NewArmID("anthropic", "claude-sonnet-4-6"),
		ModelName: "claude-sonnet-4-6",
		Provider:  security.WrapProvider(&stubProvider{name: "anthropic", model: "claude-sonnet-4-6"}, nil),
		IsLocal:   false,
		Capabilities: provider.Capabilities{
			ToolUse:       true,
			ContextWindow: 1_000_000,
		},
	})

	decision := r.Select(Task{
		Type:            TaskExplain,
		ComplexityScore: 0.4,
		Priority:        PriorityNormal,
		RequiresTools:   true,
		EstimatedTokens: 1500,
	})
	if decision.Error != nil {
		t.Fatalf("Select error: %v", decision.Error)
	}
	if !decision.Arm.IsLocal {
		t.Errorf("incognito (LocalOnly=true) must beat PreferCloud; got %s", decision.Arm.ID)
	}
	decision.Rollback()
}

// TestPreferPolicy_LocalArmsExhaustedFallsBackToCloud: PreferLocal must
// not block cloud selection when the local fleet can't handle the task.
func TestPreferPolicy_LocalArmsExhaustedFallsBackToCloud(t *testing.T) {
	r := New(Config{})
	r.SetPreferPolicy(PreferLocal)

	// Only a cloud arm registered.
	r.RegisterArm(&Arm{
		ID:        NewArmID("anthropic", "claude-opus-4-7"),
		ModelName: "claude-opus-4-7",
		Provider:  security.WrapProvider(&stubProvider{name: "anthropic", model: "claude-opus-4-7"}, nil),
		IsLocal:   false,
		Capabilities: provider.Capabilities{
			ToolUse:       true,
			ContextWindow: 1_000_000,
			ThinkingModes: []provider.EffortLevel{provider.EffortHigh},
		},
	})

	decision := r.Select(Task{
		Type:            TaskSecurityReview,
		ComplexityScore: 0.9,
		Priority:        PriorityCritical,
		RequiresTools:   true,
		EstimatedTokens: 5000,
	})
	if decision.Error != nil {
		t.Fatalf("Select error: %v", decision.Error)
	}
	if decision.Arm.ID != NewArmID("anthropic", "claude-opus-4-7") {
		t.Errorf("expected cloud arm to win when no local feasible; got %s", decision.Arm.ID)
	}
	decision.Rollback()
}