package model

// Router maps roles to model names. An unknown role resolves to the configured
// default model. When a request carries image content and a vision model is
// configured, the request routes to the vision model.
type Router struct {
	defaultModel string
	byRole       map[Role]string
	visionModel  string
}

// RouterOption configures a Router.
type RouterOption func(*Router)

// WithRoleModel binds a model to a role.
func WithRoleModel(role Role, model string) RouterOption {
	return func(r *Router) { r.byRole[role] = model }
}

// WithVisionModel configures the dedicated vision model.
func WithVisionModel(model string) RouterOption {
	return func(r *Router) { r.visionModel = model }
}

// NewRouter builds a Router with a default model applied to any role that has
// no explicit binding.
func NewRouter(defaultModel string, opts ...RouterOption) *Router {
	r := &Router{defaultModel: defaultModel, byRole: map[Role]string{}}
	for _, o := range opts {
		o(r)
	}
	return r
}

// Route returns the model for a role, falling back to the default model.
func (r *Router) Route(role Role) string {
	if m, ok := r.byRole[role]; ok && m != "" {
		return m
	}
	return r.defaultModel
}

// RouteRequest picks the model for a request and reports whether the request
// needs a vision model. When the request carries image content and a vision
// model is configured, the vision model is returned with needsVision set to
// true. Otherwise the role of the last message drives routing (user by
// default).
func (r *Router) RouteRequest(req Request) (model string, needsVision bool) {
	role := RoleUser
	if n := len(req.Messages); n > 0 {
		role = req.Messages[n-1].Role
	}
	needsVision = HasImage(req.Messages)
	if needsVision && r.visionModel != "" {
		return r.visionModel, true
	}
	return r.Route(role), needsVision
}

// HasImage reports whether any message carries an image content part.
func HasImage(msgs []Message) bool {
	for _, m := range msgs {
		for _, part := range m.Content {
			if part.Type == ContentTypeImage {
				return true
			}
		}
	}
	return false
}