server/subtasks/service.go

package subtasks

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

	"github.com/iomodo/staff/git"
	"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
	prProvider     git.PullRequestProvider // GitHub PR provider
	githubOwner    string
	githubRepo     string
	cloneManager   *git.CloneManager
}

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

// 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) {
	if s.prProvider == nil {
		return "", fmt.Errorf("PR provider not configured")
	}

	// Generate branch name for subtask proposal
	branchName := fmt.Sprintf("subtasks/%s-proposal", analysis.ParentTaskID)

	// Create Git branch and commit subtask proposal
	if err := s.createSubtaskBranch(ctx, analysis, branchName); err != nil {
		return "", fmt.Errorf("failed to create subtask branch: %w", err)
	}

	// Generate PR content
	prContent := s.generateSubtaskPRContent(analysis)
	title := fmt.Sprintf("Subtask Proposal: %s", analysis.ParentTaskID)

	// Create the pull request
	options := git.PullRequestOptions{
		Title:       title,
		Description: prContent,
		HeadBranch:  branchName,
		BaseBranch:  "main",
		Labels:      []string{"subtasks", "proposal", "ai-generated"},
		Draft:       false,
	}

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

	prURL := fmt.Sprintf("https://github.com/%s/%s/pull/%d", s.githubOwner, s.githubRepo, pr.Number)
	log.Printf("Generated subtask proposal PR: %s", prURL)

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

// createSubtaskBranch creates a Git branch with subtask proposal content
func (s *SubtaskService) createSubtaskBranch(ctx context.Context, analysis *tm.SubtaskAnalysis, branchName string) error {
	if s.cloneManager == nil {
		return fmt.Errorf("clone manager not configured")
	}

	// Get a temporary clone for creating the subtask branch
	clonePath, err := s.cloneManager.GetAgentClonePath("subtask-service")
	if err != nil {
		return fmt.Errorf("failed to get clone path: %w", err)
	}

	// All Git operations use the 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)
		}
	}

	// Pull latest changes
	cmd = gitCmd("pull", "origin")
	if err := cmd.Run(); err != nil {
		log.Printf("Warning: failed to pull latest changes: %v", err)
	}

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

	// Create subtask proposal file
	proposalDir := filepath.Join(clonePath, "tasks", "subtasks")
	if err := os.MkdirAll(proposalDir, 0755); err != nil {
		return fmt.Errorf("failed to create proposal directory: %w", err)
	}

	proposalFile := filepath.Join(proposalDir, fmt.Sprintf("%s-proposal.md", analysis.ParentTaskID))
	proposalContent := s.generateSubtaskProposalFile(analysis)

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

	// Stage the file
	relativeFile := filepath.Join("tasks", "subtasks", fmt.Sprintf("%s-proposal.md", analysis.ParentTaskID))
	cmd = gitCmd("add", relativeFile)
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to stage files: %w", err)
	}

	// Commit changes
	commitMsg := fmt.Sprintf("Subtask proposal for task %s\n\nGenerated by Staff AI Agent System\nProposed %d subtasks with %d new agents", 
		analysis.ParentTaskID, len(analysis.Subtasks), len(analysis.AgentCreations))
	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("Created subtask proposal branch: %s", branchName)
	return nil
}

// generateSubtaskProposalFile creates the content for the subtask proposal file
func (s *SubtaskService) generateSubtaskProposalFile(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("**Generated:** %s\n\n", time.Now().Format(time.RFC3339)))
	content.WriteString(fmt.Sprintf("## Analysis Summary\n%s\n\n", analysis.AnalysisSummary))

	if len(analysis.AgentCreations) > 0 {
		content.WriteString("## Proposed New Agents\n\n")
		for i, agent := range analysis.AgentCreations {
			content.WriteString(fmt.Sprintf("### %d. %s Agent\n", i+1, strings.Title(agent.Role)))
			content.WriteString(fmt.Sprintf("- **Skills:** %s\n", strings.Join(agent.Skills, ", ")))
			content.WriteString(fmt.Sprintf("- **Description:** %s\n", agent.Description))
			content.WriteString(fmt.Sprintf("- **Justification:** %s\n\n", agent.Justification))
		}
	}

	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.RequiredSkills) > 0 {
			content.WriteString(fmt.Sprintf("- **Required Skills:** %s\n", strings.Join(subtask.RequiredSkills, ", ")))
		}
		if len(subtask.Dependencies) > 0 {
			content.WriteString(fmt.Sprintf("- **Dependencies:** %s\n", strings.Join(subtask.Dependencies, ", ")))
		}
		content.WriteString(fmt.Sprintf("- **Description:** %s\n\n", subtask.Description))
	}

	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("---\n\n")
	content.WriteString("*This proposal was generated by the Staff AI Agent System. Review and approve to proceed with subtask creation.*\n")
	
	return content.String()
}

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