server/subtasks/service.go

package subtasks

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"strings"

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

// SubtaskService handles subtask generation and management
type SubtaskService struct {
	llmProvider    llm.LLMProvider
	taskManager    tm.TaskManager
	agentRoles     []string // Available agent roles for assignment
}

// NewSubtaskService creates a new subtask service
func NewSubtaskService(provider llm.LLMProvider, taskManager tm.TaskManager, agentRoles []string) *SubtaskService {
	return &SubtaskService{
		llmProvider: provider,
		taskManager: taskManager,
		agentRoles:  agentRoles,
	}
}

// ShouldGenerateSubtasks asks LLM whether a task needs subtasks based on existing agents
func (s *SubtaskService) ShouldGenerateSubtasks(ctx context.Context, task *tm.Task) (*tm.SubtaskDecision, error) {
	prompt := s.buildSubtaskDecisionPrompt(task)
	
	req := llm.ChatCompletionRequest{
		Model: "gpt-4",
		Messages: []llm.Message{
			{
				Role:    llm.RoleSystem,
				Content: s.getSubtaskDecisionSystemPrompt(),
			},
			{
				Role:    llm.RoleUser,
				Content: prompt,
			},
		},
		MaxTokens:   &[]int{1000}[0],
		Temperature: &[]float64{0.3}[0],
	}

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

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

	// Parse the LLM response
	decision, err := s.parseSubtaskDecision(resp.Choices[0].Message.Content)
	if err != nil {
		return nil, fmt.Errorf("failed to parse LLM decision: %w", err)
	}

	return decision, nil
}

// AnalyzeTaskForSubtasks uses LLM to analyze a task and propose subtasks
func (s *SubtaskService) AnalyzeTaskForSubtasks(ctx context.Context, task *tm.Task) (*tm.SubtaskAnalysis, error) {
	prompt := s.buildSubtaskAnalysisPrompt(task)
	
	req := llm.ChatCompletionRequest{
		Model: "gpt-4",
		Messages: []llm.Message{
			{
				Role:    llm.RoleSystem,
				Content: s.getSubtaskAnalysisSystemPrompt(),
			},
			{
				Role:    llm.RoleUser,
				Content: prompt,
			},
		},
		MaxTokens:   &[]int{4000}[0],
		Temperature: &[]float64{0.3}[0],
	}

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

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

	// Parse the LLM response
	analysis, err := s.parseSubtaskAnalysis(resp.Choices[0].Message.Content, task.ID)
	if err != nil {
		return nil, fmt.Errorf("failed to parse LLM response: %w", err)
	}

	return analysis, nil
}

// getSubtaskDecisionSystemPrompt returns the system prompt for subtask decision
func (s *SubtaskService) getSubtaskDecisionSystemPrompt() string {
	availableRoles := strings.Join(s.agentRoles, ", ")
	
	return fmt.Sprintf(`You are an expert project manager and task analyst. Your job is to determine whether a task needs to be broken down into subtasks.

Currently available team roles and their capabilities: %s

When evaluating a task, consider:
1. Task complexity and scope
2. Whether multiple specialized skills are needed
3. If the task can be completed by a single agent with current capabilities
4. Whether new agent roles might be needed for specialized skills

Respond with a JSON object in this exact format:
{
  "needs_subtasks": true/false,
  "reasoning": "Clear explanation of why subtasks are or aren't needed",
  "complexity_score": 5,
  "required_skills": ["skill1", "skill2", "skill3"]
}

Complexity score should be 1-10 where:
- 1-3: Simple tasks that can be handled by one agent
- 4-6: Moderate complexity, might benefit from subtasks
- 7-10: Complex tasks that definitely need breaking down

Required skills should list all technical/domain skills needed to complete the task.`, availableRoles)
}

// getSubtaskAnalysisSystemPrompt returns the system prompt for subtask analysis
func (s *SubtaskService) getSubtaskAnalysisSystemPrompt() string {
	availableRoles := strings.Join(s.agentRoles, ", ")
	
	return fmt.Sprintf(`You are an expert project manager and technical architect. Your job is to analyze complex tasks and break them down into well-defined subtasks that can be assigned to specialized team members.

Currently available team roles: %s

When analyzing a task, you should:
1. Understand the task requirements and scope
2. Break it down into logical, manageable subtasks
3. Assign each subtask to the most appropriate team role OR propose creating new agents
4. Estimate effort and identify dependencies
5. Provide a clear execution strategy

If you need specialized skills not covered by existing roles, propose new agent creation.

Respond with a JSON object in this exact format:
{
  "analysis_summary": "Brief analysis of the task and approach",
  "subtasks": [
    {
      "title": "Subtask title",
      "description": "Detailed description of what needs to be done",
      "priority": "high|medium|low",
      "assigned_to": "role_name",
      "estimated_hours": 8,
      "dependencies": ["subtask_index_1", "subtask_index_2"],
      "required_skills": ["skill1", "skill2"]
    }
  ],
  "agent_creations": [
    {
      "role": "new_role_name",
      "skills": ["specialized_skill1", "specialized_skill2"],
      "description": "Description of what this agent does",
      "justification": "Why this new agent is needed"
    }
  ],
  "recommended_approach": "High-level strategy for executing these subtasks",
  "estimated_total_hours": 40,
  "risk_assessment": "Potential risks and mitigation strategies"
}

For existing roles, use: %s
For new agents, propose appropriate role names and skill sets.
Dependencies should reference subtask indices (e.g., ["0", "1"] means depends on first and second subtasks).`, availableRoles, availableRoles)
}

// buildSubtaskDecisionPrompt creates the user prompt for subtask decision
func (s *SubtaskService) buildSubtaskDecisionPrompt(task *tm.Task) string {
	return fmt.Sprintf(`Please evaluate whether the following task needs to be broken down into subtasks:

**Task Title:** %s

**Description:** %s

**Priority:** %s

**Current Status:** %s

Consider:
- Can this be completed by a single agent with existing capabilities?
- Does it require multiple specialized skills?
- Is the scope too large for one person?
- Are there logical components that could be parallelized?

Provide your decision in the JSON format specified in the system prompt.`, 
		task.Title, 
		task.Description, 
		task.Priority, 
		task.Status)
}

// buildSubtaskAnalysisPrompt creates the user prompt for LLM analysis
func (s *SubtaskService) buildSubtaskAnalysisPrompt(task *tm.Task) string {
	return fmt.Sprintf(`Please analyze the following task and break it down into subtasks:

**Task Title:** %s

**Description:** %s

**Priority:** %s

**Current Status:** %s

Please analyze this task and provide a detailed breakdown into subtasks. Consider:
- Technical complexity and requirements
- Logical task dependencies 
- Appropriate skill sets needed for each subtask
- Risk factors and potential blockers
- Estimated effort for each component

Provide the analysis in the JSON format specified in the system prompt.`, 
		task.Title, 
		task.Description, 
		task.Priority, 
		task.Status)
}

// parseSubtaskDecision parses the LLM response into a SubtaskDecision struct
func (s *SubtaskService) parseSubtaskDecision(response string) (*tm.SubtaskDecision, error) {
	// Try to extract JSON from the response
	jsonStart := strings.Index(response, "{")
	jsonEnd := strings.LastIndex(response, "}")
	
	if jsonStart == -1 || jsonEnd == -1 {
		return nil, fmt.Errorf("no JSON found in LLM response")
	}
	
	jsonStr := response[jsonStart : jsonEnd+1]
	
	var decision tm.SubtaskDecision
	if err := json.Unmarshal([]byte(jsonStr), &decision); err != nil {
		return nil, fmt.Errorf("failed to unmarshal JSON: %w", err)
	}
	
	return &decision, nil
}

// parseSubtaskAnalysis parses the LLM response into a SubtaskAnalysis struct
func (s *SubtaskService) parseSubtaskAnalysis(response string, parentTaskID string) (*tm.SubtaskAnalysis, error) {
	// Try to extract JSON from the response (LLM might wrap it in markdown)
	jsonStart := strings.Index(response, "{")
	jsonEnd := strings.LastIndex(response, "}")
	
	if jsonStart == -1 || jsonEnd == -1 {
		return nil, fmt.Errorf("no JSON found in LLM response")
	}
	
	jsonStr := response[jsonStart : jsonEnd+1]
	
	var rawAnalysis struct {
		AnalysisSummary     string `json:"analysis_summary"`
		Subtasks            []struct {
			Title          string   `json:"title"`
			Description    string   `json:"description"`
			Priority       string   `json:"priority"`
			AssignedTo     string   `json:"assigned_to"`
			EstimatedHours int      `json:"estimated_hours"`
			Dependencies   []string `json:"dependencies"`
			RequiredSkills []string `json:"required_skills"`
		} `json:"subtasks"`
		AgentCreations      []tm.AgentCreationProposal `json:"agent_creations"`
		RecommendedApproach   string `json:"recommended_approach"`
		EstimatedTotalHours   int    `json:"estimated_total_hours"`
		RiskAssessment        string `json:"risk_assessment"`
	}
	
	if err := json.Unmarshal([]byte(jsonStr), &rawAnalysis); err != nil {
		return nil, fmt.Errorf("failed to unmarshal JSON: %w", err)
	}
	
	// Convert to our types
	analysis := &tm.SubtaskAnalysis{
		ParentTaskID:        parentTaskID,
		AnalysisSummary:     rawAnalysis.AnalysisSummary,
		AgentCreations:      rawAnalysis.AgentCreations,
		RecommendedApproach: rawAnalysis.RecommendedApproach,
		EstimatedTotalHours: rawAnalysis.EstimatedTotalHours,
		RiskAssessment:      rawAnalysis.RiskAssessment,
	}
	
	// Convert subtasks
	for _, st := range rawAnalysis.Subtasks {
		priority := tm.PriorityMedium // default
		switch strings.ToLower(st.Priority) {
		case "high":
			priority = tm.PriorityHigh
		case "low":
			priority = tm.PriorityLow
		}
		
		subtask := tm.SubtaskProposal{
			Title:          st.Title,
			Description:    st.Description,
			Priority:       priority,
			AssignedTo:     st.AssignedTo,
			EstimatedHours: st.EstimatedHours,
			Dependencies:   st.Dependencies,
			RequiredSkills: st.RequiredSkills,
		}
		
		analysis.Subtasks = append(analysis.Subtasks, subtask)
	}
	
	// Validate agent assignments and handle new agent creation
	if err := s.validateAndHandleAgentAssignments(analysis); err != nil {
		log.Printf("Warning during agent assignment handling: %v", err)
	}
	
	return analysis, nil
}

// validateAndHandleAgentAssignments validates assignments and creates agent creation subtasks if needed
func (s *SubtaskService) validateAndHandleAgentAssignments(analysis *tm.SubtaskAnalysis) error {
	// Collect all agent roles that will be available (existing + proposed new ones)
	availableRoles := make(map[string]bool)
	for _, role := range s.agentRoles {
		availableRoles[role] = true
	}
	
	// Add proposed new agent roles
	for _, agentCreation := range analysis.AgentCreations {
		availableRoles[agentCreation.Role] = true
		
		// Create a subtask for agent creation
		agentCreationSubtask := tm.SubtaskProposal{
			Title:          fmt.Sprintf("Create %s Agent", strings.Title(agentCreation.Role)),
			Description:    fmt.Sprintf("Create and configure a new %s agent with skills: %s. %s", agentCreation.Role, strings.Join(agentCreation.Skills, ", "), agentCreation.Justification),
			Priority:       tm.PriorityHigh, // Agent creation is high priority
			AssignedTo:     "ceo", // CEO creates new agents
			EstimatedHours: 4,     // Estimated time to set up new agent
			Dependencies:   []string{}, // No dependencies for agent creation
			RequiredSkills: []string{"agent_configuration", "system_design"},
		}
		
		// Insert at the beginning so agent creation happens first
		analysis.Subtasks = append([]tm.SubtaskProposal{agentCreationSubtask}, analysis.Subtasks...)
		
		// Update dependencies to account for the new subtask at index 0
		for i := 1; i < len(analysis.Subtasks); i++ {
			for j, dep := range analysis.Subtasks[i].Dependencies {
				// Convert dependency index and increment by 1
				if depIndex := s.parseDependencyIndex(dep); depIndex >= 0 {
					analysis.Subtasks[i].Dependencies[j] = fmt.Sprintf("%d", depIndex+1)
				}
			}
		}
	}
	
	// Now validate all assignments against available roles
	defaultRole := "ceo" // fallback role
	if len(s.agentRoles) > 0 {
		defaultRole = s.agentRoles[0]
	}
	
	for i := range analysis.Subtasks {
		if !availableRoles[analysis.Subtasks[i].AssignedTo] {
			log.Printf("Warning: Unknown agent role '%s' for subtask '%s', assigning to %s", 
				analysis.Subtasks[i].AssignedTo, analysis.Subtasks[i].Title, defaultRole)
			analysis.Subtasks[i].AssignedTo = defaultRole
		}
	}
	
	return nil
}

// parseDependencyIndex parses a dependency string to an integer index
func (s *SubtaskService) parseDependencyIndex(dep string) int {
	var idx int
	if _, err := fmt.Sscanf(dep, "%d", &idx); err == nil {
		return idx
	}
	return -1 // Invalid dependency format
}

// isValidAgentRole checks if a role is in the available agent roles
func (s *SubtaskService) isValidAgentRole(role string) bool {
	for _, availableRole := range s.agentRoles {
		if availableRole == role {
			return true
		}
	}
	return false
}

// GenerateSubtaskPR creates a PR with the proposed subtasks
func (s *SubtaskService) GenerateSubtaskPR(ctx context.Context, analysis *tm.SubtaskAnalysis) (string, error) {
	// Generate markdown content for the PR
	prContent := s.generateSubtaskPRContent(analysis)
	
	// This would typically create a Git branch and PR
	// For now, we'll return a mock PR URL
	prURL := fmt.Sprintf("https://github.com/example/repo/pull/subtasks-%s", analysis.ParentTaskID)
	
	log.Printf("Generated subtask proposal PR: %s", prURL)
	log.Printf("PR Content:\n%s", prContent)
	
	return prURL, nil
}

// generateSubtaskPRContent creates markdown content for the subtask proposal PR
func (s *SubtaskService) generateSubtaskPRContent(analysis *tm.SubtaskAnalysis) string {
	var content strings.Builder
	
	content.WriteString(fmt.Sprintf("# Subtask Proposal for Task %s\n\n", analysis.ParentTaskID))
	content.WriteString(fmt.Sprintf("## Analysis Summary\n%s\n\n", analysis.AnalysisSummary))
	content.WriteString(fmt.Sprintf("## Recommended Approach\n%s\n\n", analysis.RecommendedApproach))
	content.WriteString(fmt.Sprintf("**Estimated Total Hours:** %d\n\n", analysis.EstimatedTotalHours))
	
	if analysis.RiskAssessment != "" {
		content.WriteString(fmt.Sprintf("## Risk Assessment\n%s\n\n", analysis.RiskAssessment))
	}
	
	content.WriteString("## Proposed Subtasks\n\n")
	
	for i, subtask := range analysis.Subtasks {
		content.WriteString(fmt.Sprintf("### %d. %s\n", i+1, subtask.Title))
		content.WriteString(fmt.Sprintf("- **Assigned to:** %s\n", subtask.AssignedTo))
		content.WriteString(fmt.Sprintf("- **Priority:** %s\n", subtask.Priority))
		content.WriteString(fmt.Sprintf("- **Estimated Hours:** %d\n", subtask.EstimatedHours))
		
		if len(subtask.Dependencies) > 0 {
			deps := strings.Join(subtask.Dependencies, ", ")
			content.WriteString(fmt.Sprintf("- **Dependencies:** %s\n", deps))
		}
		
		content.WriteString(fmt.Sprintf("- **Description:** %s\n\n", subtask.Description))
	}
	
	content.WriteString("---\n")
	content.WriteString("*Generated by Staff AI Agent System*\n\n")
	content.WriteString("**Instructions:**\n")
	content.WriteString("- Review the proposed subtasks\n")
	content.WriteString("- Approve or request changes\n")
	content.WriteString("- When merged, the subtasks will be automatically created and assigned\n")
	
	return content.String()
}

// Close cleans up the service
func (s *SubtaskService) Close() error {
	if s.llmProvider != nil {
		return s.llmProvider.Close()
	}
	return nil
}