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/subtasks"
	"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
	subtaskService  *subtasks.SubtaskService
	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)
	}

	// Initialize subtask service after agents are created
	if err := manager.initializeSubtaskService(); err != nil {
		return nil, fmt.Errorf("failed to initialize subtask service: %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
}

// initializeSubtaskService creates the subtask service with available agent roles
func (m *Manager) initializeSubtaskService() error {
	// Get agent roles from configuration
	agentRoles := make([]string, 0, len(m.config.Agents))
	for _, agentConfig := range m.config.Agents {
		agentRoles = append(agentRoles, agentConfig.Name)
	}

	// Use the first agent's LLM provider for subtask analysis
	if len(m.agents) == 0 {
		return fmt.Errorf("no agents available for subtask service")
	}

	var firstAgent *Agent
	for _, agent := range m.agents {
		firstAgent = agent
		break
	}

	m.subtaskService = subtasks.NewSubtaskService(
		firstAgent.Provider,
		m.taskManager,
		agentRoles,
	)

	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.ProviderFake, // Use fake provider for testing
		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)
	}

	// Check if this task should generate subtasks (with LLM decision)
	if m.shouldGenerateSubtasks(task) {
		log.Printf("LLM determined task %s should generate subtasks", task.ID)
		if err := m.generateSubtasksForTask(ctx, task); err != nil {
			log.Printf("Warning: Failed to generate subtasks for task %s: %v", task.ID, err)
			// Continue with normal processing if subtask generation fails
		} else {
			// Task has been converted to subtask management, mark as completed
			task.Status = tm.StatusCompleted
			task.Solution = "Task analyzed by LLM and broken down into subtasks with potential new agent creation. See subtasks PR for details."
			completedAt := time.Now()
			task.CompletedAt = &completedAt
			agent.CurrentTask = nil

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

			log.Printf("Task %s converted to subtasks by agent %s using LLM analysis", task.ID, agent.Name)
			return nil
		}
	}

	// 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
}

// shouldGenerateSubtasks determines if a task should be broken down into subtasks using LLM
func (m *Manager) shouldGenerateSubtasks(task *tm.Task) bool {
	// Don't generate subtasks for subtasks
	if task.ParentTaskID != "" {
		return false
	}

	// Don't generate if already evaluated
	if task.SubtasksEvaluated {
		return false
	}

	// Ask LLM to decide
	ctx := context.Background()
	decision, err := m.subtaskService.ShouldGenerateSubtasks(ctx, task)
	if err != nil {
		log.Printf("Warning: Failed to get LLM subtask decision for task %s: %v", task.ID, err)
		// Fallback to simple heuristics
		return task.Priority == tm.PriorityHigh || len(task.Description) > 200
	}

	// Update task to mark as evaluated
	task.SubtasksEvaluated = true
	if err := m.taskManager.UpdateTask(task); err != nil {
		log.Printf("Warning: Failed to update task evaluation status: %v", err)
	}

	log.Printf("LLM subtask decision for task %s: needs_subtasks=%v, complexity=%d, reasoning=%s", 
		task.ID, decision.NeedsSubtasks, decision.ComplexityScore, decision.Reasoning)

	return decision.NeedsSubtasks
}

// generateSubtasksForTask analyzes a task and creates a PR with proposed subtasks
func (m *Manager) generateSubtasksForTask(ctx context.Context, task *tm.Task) error {
	if m.subtaskService == nil {
		return fmt.Errorf("subtask service not initialized")
	}

	// Analyze the task for subtasks
	analysis, err := m.subtaskService.AnalyzeTaskForSubtasks(ctx, task)
	if err != nil {
		return fmt.Errorf("failed to analyze task for subtasks: %w", err)
	}

	// Generate a PR with the subtask proposals
	prURL, err := m.subtaskService.GenerateSubtaskPR(ctx, analysis)
	if err != nil {
		return fmt.Errorf("failed to generate subtask PR: %w", err)
	}

	// Update the task with subtask information
	task.SubtasksPRURL = prURL
	task.SubtasksGenerated = true

	log.Printf("Generated subtask PR for task %s: %s", task.ID, prURL)
	log.Printf("Proposed %d subtasks and %d new agents for task %s", len(analysis.Subtasks), len(analysis.AgentCreations), task.ID)
	
	// Log proposed new agents if any
	if len(analysis.AgentCreations) > 0 {
		for _, agent := range analysis.AgentCreations {
			log.Printf("Proposed new agent: %s with skills: %v", agent.Role, agent.Skills)
		}
	}

	return nil
}

// 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)
	}

	// Cleanup subtask service
	if m.subtaskService != nil {
		if err := m.subtaskService.Close(); err != nil {
			log.Printf("Error closing subtask service: %v", err)
		}
	}

	return nil
}