package model

import "testing"

func TestRouter_Route_PerRole(t *testing.T) {
	t.Parallel()
	r := NewRouter("default",
		WithRoleModel(RoleUser, "user-model"),
		WithRoleModel(RoleAssistant, "assistant-model"),
		WithVisionModel("vision-model"),
	)
	cases := []struct {
		role Role
		want string
	}{
		{RoleUser, "user-model"},
		{RoleAssistant, "assistant-model"},
		{RoleSystem, "default"},
		{Role("unknown"), "default"},
	}
	for _, c := range cases {
		c := c
		t.Run(string(c.role), func(t *testing.T) {
			t.Parallel()
			if got := r.Route(c.role); got != c.want {
				t.Fatalf("Route(%q) = %q, want %q", c.role, got, c.want)
			}
		})
	}
}

func TestRouter_RouteRequest_Text(t *testing.T) {
	t.Parallel()
	r := NewRouter("default",
		WithRoleModel(RoleUser, "user-model"),
		WithRoleModel(RoleAssistant, "assistant-model"),
		WithVisionModel("vision-model"),
	)
	req := Request{Messages: []Message{TextMessage(RoleAssistant, "hi")}}
	model, needsVision := r.RouteRequest(req)
	if needsVision {
		t.Errorf("needsVision = true, want false for text-only")
	}
	if model != "assistant-model" {
		t.Errorf("model = %q, want assistant-model (last-message role)", model)
	}
}

func TestRouter_RouteRequest_Vision(t *testing.T) {
	t.Parallel()
	r := NewRouter("default",
		WithRoleModel(RoleUser, "user-model"),
		WithVisionModel("vision-model"),
	)
	req := Request{Messages: []Message{ImageMessage(RoleUser, "", "data:image/png;base64,abc")}}
	model, needsVision := r.RouteRequest(req)
	if !needsVision {
		t.Errorf("needsVision = false, want true for image content")
	}
	if model != "vision-model" {
		t.Errorf("model = %q, want vision-model", model)
	}
}

func TestRouter_RouteRequest_VisionUnconfigured(t *testing.T) {
	t.Parallel()
	r := NewRouter("default", WithRoleModel(RoleUser, "user-model"))
	req := Request{Messages: []Message{ImageMessage(RoleUser, "", "data:image/png;base64,abc")}}
	model, needsVision := r.RouteRequest(req)
	if !needsVision {
		t.Errorf("needsVision = false, want true when image present even without a vision model")
	}
	if model != "user-model" {
		t.Errorf("model = %q, want user-model (role fallback)", model)
	}
}

func TestRouter_RouteRequest_EmptyMessages(t *testing.T) {
	t.Parallel()
	r := NewRouter("default", WithRoleModel(RoleUser, "user-model"))
	model, needsVision := r.RouteRequest(Request{})
	if needsVision {
		t.Errorf("needsVision = true, want false")
	}
	if model != "user-model" {
		t.Errorf("model = %q, want user-model (default role for empty messages)", model)
	}
}