server/agent/manager.go

package agent

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/iomodo/staff/assignment"
	"github.com/iomodo/staff/config"
	"github.com/iomodo/staff/git"
	"github.com/iomodo/staff/llm"
	_ "github.com/iomodo/staff/llm/providers" // Auto-register all providers
	"github.com/iomodo/staff/tm"
)

// Manager manages multiple AI agents with Git operations and task processing
type Manager struct {
	config       *config.Config
	agents       map[string]*Agent
	taskManager  tm.TaskManager
	autoAssigner *assignment.AutoAssigner
	prProvider   git.PullRequestProvider
	cloneManager *git.CloneManager
	isRunning    map[string]bool
	stopChannels map[string]chan struct{}
}

// NewManager creates a new agent manager
func NewManager(cfg *config.Config, taskManager tm.TaskManager) (*Manager, error) {
	// Create auto-assigner
	autoAssigner := assignment.NewAutoAssigner(cfg.Agents)

	// Create GitHub PR provider
	githubConfig := git.GitHubConfig{
		Token: cfg.GitHub.Token,
	}
	prProvider := git.NewGitHubPullRequestProvider(cfg.GitHub.Owner, cfg.GitHub.Repo, githubConfig)

	// Create clone manager for per-agent Git repositories
	repoURL := fmt.Sprintf("https://github.com/%s/%s.git", cfg.GitHub.Owner, cfg.GitHub.Repo)
	workspacePath := filepath.Join(".", "workspace")
	cloneManager := git.NewCloneManager(repoURL, workspacePath)

	manager := &Manager{
		config:       cfg,
		agents:       make(map[string]*Agent),
		taskManager:  taskManager,
		autoAssigner: autoAssigner,
		prProvider:   prProvider,
		cloneManager: cloneManager,
		isRunning:    make(map[string]bool),
		stopChannels: make(map[string]chan struct{}),
	}

	// Initialize agents
	if err := manager.initializeAgents(); err != nil {
		return nil, fmt.Errorf("failed to initialize agents: %w", err)
	}

	return manager, nil
}

// initializeAgents creates agent instances from configuration
func (m *Manager) initializeAgents() error {
	for _, agentConfig := range m.config.Agents {
		agent, err := m.createAgent(agentConfig)
		if err != nil {
			return fmt.Errorf("failed to create agent %s: %w", agentConfig.Name, err)
		}
		m.agents[agentConfig.Name] = agent
	}
	return nil
}

// createAgent creates a single agent instance
func (m *Manager) createAgent(agentConfig config.AgentConfig) (*Agent, error) {
	// Load system prompt
	systemPrompt, err := m.loadSystemPrompt(agentConfig.SystemPromptFile)
	if err != nil {
		return nil, fmt.Errorf("failed to load system prompt: %w", err)
	}

	// Create LLM provider
	llmConfig := llm.Config{
		Provider: llm.ProviderOpenAI,
		APIKey:   m.config.OpenAI.APIKey,
		BaseURL:  m.config.OpenAI.BaseURL,
		Timeout:  m.config.OpenAI.Timeout,
	}

	provider, err := llm.CreateProvider(llmConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create LLM provider: %w", err)
	}

	agent := &Agent{
		Name:         agentConfig.Name,
		Role:         agentConfig.Role,
		Model:        agentConfig.Model,
		SystemPrompt: systemPrompt,
		Provider:     provider,
		MaxTokens:    agentConfig.MaxTokens,
		Temperature:  agentConfig.Temperature,
		Stats:        AgentStats{},
	}

	return agent, nil
}

// loadSystemPrompt loads the system prompt from file
func (m *Manager) loadSystemPrompt(filePath string) (string, error) {
	content, err := os.ReadFile(filePath)
	if err != nil {
		return "", fmt.Errorf("failed to read system prompt file %s: %w", filePath, err)
	}
	return string(content), nil
}

// StartAgent starts an agent to process tasks in a loop
func (m *Manager) StartAgent(agentName string, loopInterval time.Duration) error {
	agent, exists := m.agents[agentName]
	if !exists {
		return fmt.Errorf("agent %s not found", agentName)
	}

	if m.isRunning[agentName] {
		return fmt.Errorf("agent %s is already running", agentName)
	}

	stopChan := make(chan struct{})
	m.stopChannels[agentName] = stopChan
	m.isRunning[agentName] = true

	go m.runAgentLoop(agent, loopInterval, stopChan)

	log.Printf("Started agent %s (%s) with %s model", agentName, agent.Role, agent.Model)
	return nil
}

// StopAgent stops a running agent
func (m *Manager) StopAgent(agentName string) error {
	if !m.isRunning[agentName] {
		return fmt.Errorf("agent %s is not running", agentName)
	}

	close(m.stopChannels[agentName])
	delete(m.stopChannels, agentName)
	m.isRunning[agentName] = false

	log.Printf("Stopped agent %s", agentName)
	return nil
}

// runAgentLoop runs the main processing loop for an agent
func (m *Manager) runAgentLoop(agent *Agent, interval time.Duration, stopChan <-chan struct{}) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-stopChan:
			log.Printf("Agent %s stopping", agent.Name)
			return
		case <-ticker.C:
			if err := m.processAgentTasks(agent); err != nil {
				log.Printf("Error processing tasks for agent %s: %v", agent.Name, err)
			}
		}
	}
}

// processAgentTasks processes all assigned tasks for an agent
func (m *Manager) processAgentTasks(agent *Agent) error {
	if agent.CurrentTask != nil {
		return nil
	}

	// Get tasks assigned to this agent
	tasks, err := m.taskManager.GetTasksByAssignee(agent.Name)
	if err != nil {
		return fmt.Errorf("failed to get tasks for agent %s: %w", agent.Name, err)
	}

	log.Printf("Processing %d tasks for agent %s", len(tasks), agent.Name)

	for _, task := range tasks {
		if task.Status == tm.StatusToDo || task.Status == tm.StatusPending {
			if err := m.processTask(agent, task); err != nil {
				log.Printf("Error processing task %s: %v", task.ID, err)
				// Mark task as failed
				task.Status = tm.StatusFailed
				if err := m.taskManager.UpdateTask(task); err != nil {
					log.Printf("Error updating failed task %s: %v", task.ID, err)
				}
				agent.Stats.TasksFailed++
			} else {
				agent.Stats.TasksCompleted++
			}
			// Update success rate
			total := agent.Stats.TasksCompleted + agent.Stats.TasksFailed
			if total > 0 {
				agent.Stats.SuccessRate = float64(agent.Stats.TasksCompleted) / float64(total) * 100
			}
		}
	}

	return nil
}

// processTask processes a single task with an agent
func (m *Manager) processTask(agent *Agent, task *tm.Task) error {
	ctx := context.Background()
	startTime := time.Now()

	log.Printf("Agent %s processing task %s: %s", agent.Name, task.ID, task.Title)

	// Mark task as in progress
	task.Status = tm.StatusInProgress
	agent.CurrentTask = &task.ID
	if err := m.taskManager.UpdateTask(task); err != nil {
		return fmt.Errorf("failed to update task status: %w", err)
	}

	// Generate solution using LLM
	solution, err := m.generateSolution(ctx, agent, task)
	if err != nil {
		return fmt.Errorf("failed to generate solution: %w", err)
	}

	// Create Git branch and commit solution
	branchName := m.generateBranchName(task)
	if err := m.createAndCommitSolution(branchName, task, solution, agent); err != nil {
		return fmt.Errorf("failed to commit solution: %w", err)
	}

	// Create pull request
	prURL, err := m.createPullRequest(ctx, task, solution, agent, branchName)
	if err != nil {
		return fmt.Errorf("failed to create pull request: %w", err)
	}

	// Update task as completed
	task.Status = tm.StatusCompleted
	task.Solution = solution
	task.PullRequestURL = prURL
	completedAt := time.Now()
	task.CompletedAt = &completedAt
	agent.CurrentTask = nil

	if err := m.taskManager.UpdateTask(task); err != nil {
		return fmt.Errorf("failed to update completed task: %w", err)
	}

	// Update agent stats
	duration := time.Since(startTime)
	if agent.Stats.AvgTime == 0 {
		agent.Stats.AvgTime = duration.Milliseconds()
	} else {
		agent.Stats.AvgTime = (agent.Stats.AvgTime + duration.Milliseconds()) / 2
	}

	log.Printf("Task %s completed by agent %s in %v. PR: %s", task.ID, agent.Name, duration, prURL)
	return nil
}

// generateSolution uses the agent's LLM to generate a solution
func (m *Manager) generateSolution(ctx context.Context, agent *Agent, task *tm.Task) (string, error) {
	prompt := m.buildTaskPrompt(task)

	req := llm.ChatCompletionRequest{
		Model: agent.Model,
		Messages: []llm.Message{
			{
				Role:    llm.RoleSystem,
				Content: agent.SystemPrompt,
			},
			{
				Role:    llm.RoleUser,
				Content: prompt,
			},
		},
		MaxTokens:   agent.MaxTokens,
		Temperature: agent.Temperature,
	}

	resp, err := agent.Provider.ChatCompletion(ctx, req)
	if err != nil {
		return "", fmt.Errorf("LLM request failed: %w", err)
	}

	if len(resp.Choices) == 0 {
		return "", fmt.Errorf("no response from LLM")
	}

	return resp.Choices[0].Message.Content, nil
}

// buildTaskPrompt creates a detailed prompt for the LLM
func (m *Manager) buildTaskPrompt(task *tm.Task) string {
	return fmt.Sprintf(`Task: %s

Priority: %s
Description: %s

Please provide a complete solution for this task. Include:
1. Detailed implementation plan
2. Code changes needed (if applicable)
3. Files to be created or modified
4. Testing considerations
5. Any dependencies or prerequisites

Your response should be comprehensive and actionable.`,
		task.Title,
		task.Priority,
		task.Description)
}

// generateBranchName creates a Git branch name for the task
func (m *Manager) generateBranchName(task *tm.Task) string {
	// Clean title for use in branch name
	cleanTitle := strings.ToLower(task.Title)
	cleanTitle = strings.ReplaceAll(cleanTitle, " ", "-")
	cleanTitle = strings.ReplaceAll(cleanTitle, "/", "-")
	// Remove special characters
	var result strings.Builder
	for _, r := range cleanTitle {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-' {
			result.WriteRune(r)
		}
	}
	cleanTitle = result.String()

	// Limit length
	if len(cleanTitle) > 40 {
		cleanTitle = cleanTitle[:40]
	}

	return fmt.Sprintf("%s%s-%s", m.config.Git.BranchPrefix, task.ID, cleanTitle)
}

// createAndCommitSolution creates a Git branch and commits the solution using per-agent clones
func (m *Manager) createAndCommitSolution(branchName string, task *tm.Task, solution string, agent *Agent) error {
	ctx := context.Background()

	// Get agent's dedicated Git clone
	clonePath, err := m.cloneManager.GetAgentClonePath(agent.Name)
	if err != nil {
		return fmt.Errorf("failed to get agent clone: %w", err)
	}

	log.Printf("Agent %s working in clone: %s", agent.Name, clonePath)

	// Refresh the clone with latest changes
	if err := m.cloneManager.RefreshAgentClone(agent.Name); err != nil {
		log.Printf("Warning: Failed to refresh clone for agent %s: %v", agent.Name, err)
	}

	// All Git operations use the agent's clone directory
	gitCmd := func(args ...string) *exec.Cmd {
		return exec.CommandContext(ctx, "git", append([]string{"-C", clonePath}, args...)...)
	}

	// Ensure we're on main branch before creating new branch
	cmd := gitCmd("checkout", "main")
	if err := cmd.Run(); err != nil {
		// Try master branch if main doesn't exist
		cmd = gitCmd("checkout", "master")
		if err := cmd.Run(); err != nil {
			return fmt.Errorf("failed to checkout main/master branch: %w", err)
		}
	}

	// Create branch
	cmd = gitCmd("checkout", "-b", branchName)
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to create branch: %w", err)
	}

	// Create solution file in agent's clone
	solutionDir := filepath.Join(clonePath, "tasks", "solutions")
	if err := os.MkdirAll(solutionDir, 0755); err != nil {
		return fmt.Errorf("failed to create solution directory: %w", err)
	}

	solutionFile := filepath.Join(solutionDir, fmt.Sprintf("%s-solution.md", task.ID))
	solutionContent := fmt.Sprintf(`# Solution for Task: %s

**Agent:** %s (%s)  
**Model:** %s  
**Completed:** %s

## Task Description
%s

## Solution
%s

---
*Generated by Staff AI Agent System*
`, task.Title, agent.Name, agent.Role, agent.Model, time.Now().Format(time.RFC3339), task.Description, solution)

	if err := os.WriteFile(solutionFile, []byte(solutionContent), 0644); err != nil {
		return fmt.Errorf("failed to write solution file: %w", err)
	}

	// Stage files
	relativeSolutionFile := filepath.Join("tasks", "solutions", fmt.Sprintf("%s-solution.md", task.ID))
	cmd = gitCmd("add", relativeSolutionFile)
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to stage files: %w", err)
	}

	// Commit changes
	commitMsg := m.buildCommitMessage(task, agent)
	cmd = gitCmd("commit", "-m", commitMsg)
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to commit: %w", err)
	}

	// Push branch
	cmd = gitCmd("push", "-u", "origin", branchName)
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to push branch: %w", err)
	}

	log.Printf("Agent %s successfully pushed branch %s", agent.Name, branchName)
	return nil
}

// buildCommitMessage creates a commit message from template
func (m *Manager) buildCommitMessage(task *tm.Task, agent *Agent) string {
	template := m.config.Git.CommitMessageTemplate

	replacements := map[string]string{
		"{task_id}":    task.ID,
		"{task_title}": task.Title,
		"{agent_name}": agent.Name,
		"{solution}":   "See solution file for details",
	}

	result := template
	for placeholder, value := range replacements {
		result = strings.ReplaceAll(result, placeholder, value)
	}

	return result
}

// createPullRequest creates a GitHub pull request
func (m *Manager) createPullRequest(ctx context.Context, task *tm.Task, solution string, agent *Agent, branchName string) (string, error) {
	title := fmt.Sprintf("Task %s: %s", task.ID, task.Title)

	// Build PR description from template
	description := m.buildPRDescription(task, solution, agent)

	options := git.PullRequestOptions{
		Title:       title,
		Description: description,
		HeadBranch:  branchName,
		BaseBranch:  "main",
		Labels:      []string{"ai-generated", "staff-agent", strings.ToLower(agent.Role)},
		Draft:       false,
	}

	pr, err := m.prProvider.CreatePullRequest(ctx, options)
	if err != nil {
		return "", fmt.Errorf("failed to create PR: %w", err)
	}

	return fmt.Sprintf("https://github.com/%s/%s/pull/%d", m.config.GitHub.Owner, m.config.GitHub.Repo, pr.Number), nil
}

// buildPRDescription creates PR description from template
func (m *Manager) buildPRDescription(task *tm.Task, solution string, agent *Agent) string {
	template := m.config.Git.PRTemplate

	// Truncate solution for PR if too long
	truncatedSolution := solution
	if len(solution) > 1000 {
		truncatedSolution = solution[:1000] + "...\n\n*See solution file for complete details*"
	}

	replacements := map[string]string{
		"{task_id}":          task.ID,
		"{task_title}":       task.Title,
		"{task_description}": task.Description,
		"{agent_name}":       fmt.Sprintf("%s (%s)", agent.Name, agent.Role),
		"{priority}":         string(task.Priority),
		"{solution}":         truncatedSolution,
		"{files_changed}":    fmt.Sprintf("- `tasks/solutions/%s-solution.md`", task.ID),
	}

	result := template
	for placeholder, value := range replacements {
		result = strings.ReplaceAll(result, placeholder, value)
	}

	return result
}

// AutoAssignTask automatically assigns a task to the best matching agent
func (m *Manager) AutoAssignTask(taskID string) error {
	task, err := m.taskManager.GetTask(taskID)
	if err != nil {
		return fmt.Errorf("failed to get task: %w", err)
	}

	agentName, err := m.autoAssigner.AssignTask(task)
	if err != nil {
		return fmt.Errorf("failed to auto-assign task: %w", err)
	}

	task.Assignee = agentName
	if err := m.taskManager.UpdateTask(task); err != nil {
		return fmt.Errorf("failed to update task assignment: %w", err)
	}

	explanation := m.autoAssigner.GetRecommendationExplanation(task, agentName)
	log.Printf("Auto-assigned task %s to %s: %s", taskID, agentName, explanation)

	return nil
}

// GetAgentStatus returns the status of all agents
func (m *Manager) GetAgentStatus() map[string]AgentInfo {
	status := make(map[string]AgentInfo)

	for name, agent := range m.agents {
		agentStatus := StatusIdle
		if m.isRunning[name] {
			if agent.CurrentTask != nil {
				agentStatus = StatusRunning
			}
		} else {
			agentStatus = StatusStopped
		}

		status[name] = AgentInfo{
			Name:        agent.Name,
			Role:        agent.Role,
			Model:       agent.Model,
			Status:      agentStatus,
			CurrentTask: agent.CurrentTask,
			Stats:       agent.Stats,
		}
	}

	return status
}

// IsAgentRunning checks if an agent is currently running
func (m *Manager) IsAgentRunning(agentName string) bool {
	return m.isRunning[agentName]
}

// Close shuts down the agent manager
func (m *Manager) Close() error {
	// Stop all running agents
	for agentName := range m.isRunning {
		if m.isRunning[agentName] {
			m.StopAgent(agentName)
		}
	}

	// Close all LLM providers
	for _, agent := range m.agents {
		if err := agent.Provider.Close(); err != nil {
			log.Printf("Error closing provider for agent %s: %v", agent.Name, err)
		}
	}

	// Cleanup all agent Git clones
	if err := m.cloneManager.CleanupAllClones(); err != nil {
		log.Printf("Error cleaning up agent clones: %v", err)
	}

	return nil
}