server/task/auto_assignment.go

package task

import (
	"fmt"
	"sort"
	"strings"

	"github.com/iomodo/staff/config"
	"github.com/iomodo/staff/tm"
)

// AutoAssigner handles intelligent task assignment to agents
type AutoAssigner struct {
	agents []config.AgentConfig
}

// NewAutoAssigner creates a new auto-assignment system
func NewAutoAssigner(agents []config.AgentConfig) *AutoAssigner {
	return &AutoAssigner{
		agents: agents,
	}
}

// AssignmentScore represents how well an agent matches a task
type AssignmentScore struct {
	AgentName string
	Score     float64
	Reasons   []string
}

// AssignTask automatically assigns a task to the best matching agent
func (a *AutoAssigner) AssignTask(task *tm.Task) (string, error) {
	if len(a.agents) == 0 {
		return "", fmt.Errorf("no agents available for assignment")
	}

	scores := a.calculateScores(task)
	if len(scores) == 0 {
		return "", fmt.Errorf("no suitable agent found for task")
	}

	// Sort by score (highest first)
	sort.Slice(scores, func(i, j int) bool {
		return scores[i].Score > scores[j].Score
	})

	bestMatch := scores[0]
	if bestMatch.Score == 0 {
		// No good match, assign to CEO as fallback
		return "ceo", nil
	}

	return bestMatch.AgentName, nil
}

// GetAssignmentRecommendations returns ranked recommendations for task assignment
func (a *AutoAssigner) GetAssignmentRecommendations(task *tm.Task) []AssignmentScore {
	scores := a.calculateScores(task)

	// Sort by score (highest first)
	sort.Slice(scores, func(i, j int) bool {
		return scores[i].Score > scores[j].Score
	})

	return scores
}

// calculateScores calculates assignment scores for all agents
func (a *AutoAssigner) calculateScores(task *tm.Task) []AssignmentScore {
	scores := make([]AssignmentScore, 0, len(a.agents))

	taskText := strings.ToLower(task.Title + " " + task.Description)
	taskKeywords := extractKeywords(taskText)

	for _, agent := range a.agents {
		score := &AssignmentScore{
			AgentName: agent.Name,
			Score:     0,
			Reasons:   make([]string, 0),
		}

		// Score based on task type keywords
		score.Score += a.scoreTaskTypes(agent.TaskTypes, taskKeywords, score)

		// Score based on capabilities
		score.Score += a.scoreCapabilities(agent.Capabilities, taskKeywords, score)

		// Score based on task priority and agent model
		score.Score += a.scorePriorityModelMatch(task.Priority, agent.Model, score)

		// Score based on explicit agent mention in task
		score.Score += a.scoreExplicitMention(agent.Name, agent.Role, taskText, score)

		scores = append(scores, *score)
	}

	return scores
}

// scoreTaskTypes scores based on task type matching
func (a *AutoAssigner) scoreTaskTypes(agentTypes []string, taskKeywords []string, score *AssignmentScore) float64 {
	typeScore := 0.0

	for _, agentType := range agentTypes {
		for _, keyword := range taskKeywords {
			if strings.Contains(keyword, agentType) || strings.Contains(agentType, keyword) {
				typeScore += 3.0
				score.Reasons = append(score.Reasons, fmt.Sprintf("matches task type: %s", agentType))
			}
		}
	}

	return typeScore
}

// scoreCapabilities scores based on capability matching
func (a *AutoAssigner) scoreCapabilities(capabilities []string, taskKeywords []string, score *AssignmentScore) float64 {
	capScore := 0.0

	for _, capability := range capabilities {
		for _, keyword := range taskKeywords {
			if strings.Contains(keyword, capability) || strings.Contains(capability, keyword) {
				capScore += 2.0
				score.Reasons = append(score.Reasons, fmt.Sprintf("has capability: %s", capability))
			}
		}
	}

	return capScore
}

// scorePriorityModelMatch scores based on priority and model sophistication
func (a *AutoAssigner) scorePriorityModelMatch(priority tm.TaskPriority, model string, score *AssignmentScore) float64 {
	priorityScore := 0.0

	// High priority tasks prefer more capable models
	if priority == tm.PriorityHigh && strings.Contains(model, "gpt-4") {
		priorityScore += 1.0
		score.Reasons = append(score.Reasons, "high priority task matches advanced model")
	}

	// Medium/low priority can use efficient models
	if priority != tm.PriorityHigh && strings.Contains(model, "gpt-3.5") {
		priorityScore += 0.5
		score.Reasons = append(score.Reasons, "priority matches model efficiency")
	}

	return priorityScore
}

// scoreExplicitMention scores based on explicit agent/role mentions
func (a *AutoAssigner) scoreExplicitMention(agentName, agentRole, taskText string, score *AssignmentScore) float64 {
	mentionScore := 0.0

	// Check for explicit agent name mention
	if strings.Contains(taskText, agentName) {
		mentionScore += 5.0
		score.Reasons = append(score.Reasons, "explicitly mentioned by name")
	}

	// Check for role mention
	roleLower := strings.ToLower(agentRole)
	if strings.Contains(taskText, roleLower) {
		mentionScore += 4.0
		score.Reasons = append(score.Reasons, "role mentioned in task")
	}

	return mentionScore
}

// extractKeywords extracts relevant keywords from task text
func extractKeywords(text string) []string {
	// Simple keyword extraction - split by common delimiters
	words := strings.FieldsFunc(text, func(c rune) bool {
		return c == ' ' || c == ',' || c == '.' || c == ':' || c == ';' || c == '\n' || c == '\t'
	})

	// Filter out common stop words and short words
	stopWords := map[string]bool{
		"the": true, "a": true, "an": true, "and": true, "or": true, "but": true,
		"in": true, "on": true, "at": true, "to": true, "for": true, "of": true,
		"with": true, "by": true, "is": true, "are": true, "was": true, "were": true,
		"be": true, "been": true, "have": true, "has": true, "had": true, "do": true,
		"does": true, "did": true, "will": true, "would": true, "could": true, "should": true,
	}

	keywords := make([]string, 0)
	for _, word := range words {
		word = strings.ToLower(strings.TrimSpace(word))
		if len(word) > 2 && !stopWords[word] {
			keywords = append(keywords, word)
		}
	}

	return keywords
}

// ValidateAssignment checks if an assignment is valid
func (a *AutoAssigner) ValidateAssignment(agentName string, task *tm.Task) error {
	// Check if agent exists
	for _, agent := range a.agents {
		if agent.Name == agentName {
			return nil
		}
	}

	return fmt.Errorf("agent '%s' not found", agentName)
}

// GetAgentCapabilities returns the capabilities of a specific agent
func (a *AutoAssigner) GetAgentCapabilities(agentName string) ([]string, error) {
	for _, agent := range a.agents {
		if agent.Name == agentName {
			return agent.Capabilities, nil
		}
	}

	return nil, fmt.Errorf("agent '%s' not found", agentName)
}

// GetRecommendationExplanation returns a human-readable explanation for assignment
func (a *AutoAssigner) GetRecommendationExplanation(task *tm.Task, agentName string) string {
	recommendations := a.GetAssignmentRecommendations(task)

	for _, rec := range recommendations {
		if rec.AgentName == agentName {
			if len(rec.Reasons) == 0 {
				return fmt.Sprintf("Agent %s assigned (score: %.1f)", agentName, rec.Score)
			}

			reasons := strings.Join(rec.Reasons, ", ")
			return fmt.Sprintf("Agent %s assigned (score: %.1f) because: %s", agentName, rec.Score, reasons)
		}
	}

	return fmt.Sprintf("Agent %s assigned (manual override)", agentName)
}