CHORUS/pkg/ai/resetdata.go

package ai

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

// ResetDataProvider implements ModelProvider for ResetData LaaS API
type ResetDataProvider struct {
	config     ProviderConfig
	httpClient *http.Client
}

// ResetDataRequest represents a request to ResetData LaaS API
type ResetDataRequest struct {
	Model       string                 `json:"model"`
	Messages    []ResetDataMessage     `json:"messages"`
	Stream      bool                   `json:"stream"`
	Temperature float32                `json:"temperature,omitempty"`
	MaxTokens   int                    `json:"max_tokens,omitempty"`
	Stop        []string               `json:"stop,omitempty"`
	TopP        float32                `json:"top_p,omitempty"`
}

// ResetDataMessage represents a message in the ResetData format
type ResetDataMessage struct {
	Role    string `json:"role"`    // system, user, assistant
	Content string `json:"content"`
}

// ResetDataResponse represents a response from ResetData LaaS API
type ResetDataResponse struct {
	ID      string                 `json:"id"`
	Object  string                 `json:"object"`
	Created int64                  `json:"created"`
	Model   string                 `json:"model"`
	Choices []ResetDataChoice      `json:"choices"`
	Usage   ResetDataUsage         `json:"usage"`
}

// ResetDataChoice represents a choice in the response
type ResetDataChoice struct {
	Index        int                `json:"index"`
	Message      ResetDataMessage   `json:"message"`
	FinishReason string             `json:"finish_reason"`
}

// ResetDataUsage represents token usage information
type ResetDataUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// ResetDataModelsResponse represents available models response
type ResetDataModelsResponse struct {
	Object string             `json:"object"`
	Data   []ResetDataModel   `json:"data"`
}

// ResetDataModel represents a model in ResetData
type ResetDataModel struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	OwnedBy string `json:"owned_by"`
}

// NewResetDataProvider creates a new ResetData provider instance
func NewResetDataProvider(config ProviderConfig) *ResetDataProvider {
	timeout := config.Timeout
	if timeout == 0 {
		timeout = 300 * time.Second // 5 minutes default for task execution
	}

	return &ResetDataProvider{
		config: config,
		httpClient: &http.Client{
			Timeout: timeout,
		},
	}
}

// ExecuteTask implements the ModelProvider interface for ResetData
func (p *ResetDataProvider) ExecuteTask(ctx context.Context, request *TaskRequest) (*TaskResponse, error) {
	startTime := time.Now()

	// Build messages for the chat completion
	messages, err := p.buildChatMessages(request)
	if err != nil {
		return nil, NewProviderError(ErrTaskExecutionFailed, fmt.Sprintf("failed to build messages: %v", err))
	}

	// Prepare the ResetData request
	resetDataReq := ResetDataRequest{
		Model:       p.selectModel(request.ModelName),
		Messages:    messages,
		Stream:      false,
		Temperature: p.getTemperature(request.Temperature),
		MaxTokens:   p.getMaxTokens(request.MaxTokens),
	}

	// Execute the request
	response, err := p.makeRequest(ctx, "/v1/chat/completions", resetDataReq)
	if err != nil {
		return nil, err
	}

	endTime := time.Now()

	// Process the response
	if len(response.Choices) == 0 {
		return nil, NewProviderError(ErrTaskExecutionFailed, "no response choices returned from ResetData")
	}

	choice := response.Choices[0]
	responseText := choice.Message.Content

	// Parse response for actions and artifacts
	actions, artifacts := p.parseResponseForActions(responseText, request)

	return &TaskResponse{
		Success:   true,
		TaskID:    request.TaskID,
		AgentID:   request.AgentID,
		ModelUsed: response.Model,
		Provider:  "resetdata",
		Response:  responseText,
		Actions:   actions,
		Artifacts: artifacts,
		StartTime: startTime,
		EndTime:   endTime,
		Duration:  endTime.Sub(startTime),
		TokensUsed: TokenUsage{
			PromptTokens:     response.Usage.PromptTokens,
			CompletionTokens: response.Usage.CompletionTokens,
			TotalTokens:      response.Usage.TotalTokens,
		},
	}, nil
}

// GetCapabilities returns ResetData provider capabilities
func (p *ResetDataProvider) GetCapabilities() ProviderCapabilities {
	return ProviderCapabilities{
		SupportsMCP:       p.config.EnableMCP,
		SupportsTools:     p.config.EnableTools,
		SupportsStreaming: true,
		SupportsFunctions: false, // ResetData LaaS doesn't support function calling
		MaxTokens:         p.config.MaxTokens,
		SupportedModels:   p.getSupportedModels(),
		SupportsImages:    false, // Most ResetData models don't support images
		SupportsFiles:     true,
	}
}

// ValidateConfig validates the ResetData provider configuration
func (p *ResetDataProvider) ValidateConfig() error {
	if p.config.APIKey == "" {
		return NewProviderError(ErrAPIKeyRequired, "API key is required for ResetData provider")
	}

	if p.config.Endpoint == "" {
		return NewProviderError(ErrInvalidConfiguration, "endpoint is required for ResetData provider")
	}

	if p.config.DefaultModel == "" {
		return NewProviderError(ErrInvalidConfiguration, "default_model is required for ResetData provider")
	}

	// Test the API connection
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if err := p.testConnection(ctx); err != nil {
		return NewProviderError(ErrProviderUnavailable, fmt.Sprintf("failed to connect to ResetData: %v", err))
	}

	return nil
}

// GetProviderInfo returns information about the ResetData provider
func (p *ResetDataProvider) GetProviderInfo() ProviderInfo {
	return ProviderInfo{
		Name:           "ResetData",
		Type:           "resetdata",
		Version:        "1.0.0",
		Endpoint:       p.config.Endpoint,
		DefaultModel:   p.config.DefaultModel,
		RequiresAPIKey: true,
		RateLimit:      600, // 10 requests per second typical limit
	}
}

// buildChatMessages constructs messages for the ResetData chat completion
func (p *ResetDataProvider) buildChatMessages(request *TaskRequest) ([]ResetDataMessage, error) {
	var messages []ResetDataMessage

	// System message
	systemPrompt := p.getSystemPrompt(request)
	if systemPrompt != "" {
		messages = append(messages, ResetDataMessage{
			Role:    "system",
			Content: systemPrompt,
		})
	}

	// User message with task details
	userPrompt, err := p.buildTaskPrompt(request)
	if err != nil {
		return nil, err
	}

	messages = append(messages, ResetDataMessage{
		Role:    "user",
		Content: userPrompt,
	})

	return messages, nil
}

// buildTaskPrompt constructs a comprehensive prompt for task execution
func (p *ResetDataProvider) buildTaskPrompt(request *TaskRequest) (string, error) {
	var prompt strings.Builder

	prompt.WriteString(fmt.Sprintf("Acting as a %s agent, analyze and work on this task:\n\n",
		request.AgentRole))

	prompt.WriteString(fmt.Sprintf("**Repository:** %s\n", request.Repository))
	prompt.WriteString(fmt.Sprintf("**Task Title:** %s\n", request.TaskTitle))
	prompt.WriteString(fmt.Sprintf("**Description:**\n%s\n\n", request.TaskDescription))

	if len(request.TaskLabels) > 0 {
		prompt.WriteString(fmt.Sprintf("**Labels:** %s\n", strings.Join(request.TaskLabels, ", ")))
	}

	prompt.WriteString(fmt.Sprintf("**Priority:** %d/10 | **Complexity:** %d/10\n\n",
		request.Priority, request.Complexity))

	if request.WorkingDirectory != "" {
		prompt.WriteString(fmt.Sprintf("**Working Directory:** %s\n", request.WorkingDirectory))
	}

	if len(request.RepositoryFiles) > 0 {
		prompt.WriteString("**Relevant Files:**\n")
		for _, file := range request.RepositoryFiles {
			prompt.WriteString(fmt.Sprintf("- %s\n", file))
		}
		prompt.WriteString("\n")
	}

	// Add role-specific instructions
	prompt.WriteString(p.getRoleSpecificInstructions(request.AgentRole))

	prompt.WriteString("\nProvide a detailed analysis and implementation plan. ")
	prompt.WriteString("Include specific steps, code changes, and any commands that need to be executed. ")
	prompt.WriteString("Focus on delivering actionable results that address the task requirements completely.")

	return prompt.String(), nil
}

// getRoleSpecificInstructions returns instructions specific to the agent role
func (p *ResetDataProvider) getRoleSpecificInstructions(role string) string {
	switch strings.ToLower(role) {
	case "developer":
		return `**Developer Focus Areas:**
- Implement robust, well-tested code solutions
- Follow coding standards and best practices
- Ensure proper error handling and edge case coverage
- Write clear documentation and comments
- Consider performance, security, and maintainability`

	case "reviewer":
		return `**Code Review Focus Areas:**
- Evaluate code quality, style, and best practices
- Identify potential bugs, security issues, and performance bottlenecks
- Check test coverage and test quality
- Verify documentation completeness and accuracy
- Suggest refactoring and improvement opportunities`

	case "architect":
		return `**Architecture Focus Areas:**
- Design scalable and maintainable system components
- Make informed decisions about technologies and patterns
- Define clear interfaces and integration points
- Consider scalability, security, and performance requirements
- Document architectural decisions and trade-offs`

	case "tester":
		return `**Testing Focus Areas:**
- Design comprehensive test strategies and test cases
- Implement automated tests at multiple levels
- Identify edge cases and failure scenarios
- Set up continuous testing and quality assurance
- Validate requirements and acceptance criteria`

	default:
		return `**General Focus Areas:**
- Understand requirements and constraints thoroughly
- Apply software engineering best practices
- Provide clear, actionable recommendations
- Consider long-term maintainability and extensibility`
	}
}

// selectModel chooses the appropriate ResetData model
func (p *ResetDataProvider) selectModel(requestedModel string) string {
	if requestedModel != "" {
		return requestedModel
	}
	return p.config.DefaultModel
}

// getTemperature returns the temperature setting
func (p *ResetDataProvider) getTemperature(requestTemp float32) float32 {
	if requestTemp > 0 {
		return requestTemp
	}
	if p.config.Temperature > 0 {
		return p.config.Temperature
	}
	return 0.7 // Default temperature
}

// getMaxTokens returns the max tokens setting
func (p *ResetDataProvider) getMaxTokens(requestTokens int) int {
	if requestTokens > 0 {
		return requestTokens
	}
	if p.config.MaxTokens > 0 {
		return p.config.MaxTokens
	}
	return 4096 // Default max tokens
}

// getSystemPrompt constructs the system prompt
func (p *ResetDataProvider) getSystemPrompt(request *TaskRequest) string {
	if request.SystemPrompt != "" {
		return request.SystemPrompt
	}

	return fmt.Sprintf(`You are an expert software development AI assistant working as a %s agent
in the CHORUS autonomous development system.

Your expertise includes:
- Software architecture and design patterns
- Code implementation across multiple programming languages
- Testing strategies and quality assurance
- DevOps and deployment practices
- Security and performance optimization

Provide detailed, practical solutions with specific implementation steps.
Focus on delivering high-quality, production-ready results.`, request.AgentRole)
}

// makeRequest makes an HTTP request to the ResetData API
func (p *ResetDataProvider) makeRequest(ctx context.Context, endpoint string, request interface{}) (*ResetDataResponse, error) {
	requestJSON, err := json.Marshal(request)
	if err != nil {
		return nil, NewProviderError(ErrTaskExecutionFailed, fmt.Sprintf("failed to marshal request: %v", err))
	}

	url := strings.TrimSuffix(p.config.Endpoint, "/") + endpoint
	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(requestJSON))
	if err != nil {
		return nil, NewProviderError(ErrTaskExecutionFailed, fmt.Sprintf("failed to create request: %v", err))
	}

	// Set required headers
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+p.config.APIKey)

	// Add custom headers if configured
	for key, value := range p.config.CustomHeaders {
		req.Header.Set(key, value)
	}

	resp, err := p.httpClient.Do(req)
	if err != nil {
		return nil, NewProviderError(ErrProviderUnavailable, fmt.Sprintf("request failed: %v", err))
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, NewProviderError(ErrTaskExecutionFailed, fmt.Sprintf("failed to read response: %v", err))
	}

	if resp.StatusCode != http.StatusOK {
		return nil, p.handleHTTPError(resp.StatusCode, body)
	}

	var resetDataResp ResetDataResponse
	if err := json.Unmarshal(body, &resetDataResp); err != nil {
		return nil, NewProviderError(ErrTaskExecutionFailed, fmt.Sprintf("failed to parse response: %v", err))
	}

	return &resetDataResp, nil
}

// testConnection tests the connection to ResetData API
func (p *ResetDataProvider) testConnection(ctx context.Context) error {
	url := strings.TrimSuffix(p.config.Endpoint, "/") + "/v1/models"
	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return err
	}

	req.Header.Set("Authorization", "Bearer "+p.config.APIKey)

	resp, err := p.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("API test failed with status %d: %s", resp.StatusCode, string(body))
	}

	return nil
}

// getSupportedModels returns a list of supported ResetData models
func (p *ResetDataProvider) getSupportedModels() []string {
	// Common models available through ResetData LaaS
	return []string{
		"llama3.1:8b", "llama3.1:70b",
		"mistral:7b", "mixtral:8x7b",
		"qwen2:7b", "qwen2:72b",
		"gemma:7b", "gemma2:9b",
		"codellama:7b", "codellama:13b",
	}
}

// handleHTTPError converts HTTP errors to provider errors
func (p *ResetDataProvider) handleHTTPError(statusCode int, body []byte) *ProviderError {
	bodyStr := string(body)

	switch statusCode {
	case http.StatusUnauthorized:
		return &ProviderError{
			Code:      "UNAUTHORIZED",
			Message:   "Invalid ResetData API key",
			Details:   bodyStr,
			Retryable: false,
		}
	case http.StatusTooManyRequests:
		return &ProviderError{
			Code:      "RATE_LIMIT_EXCEEDED",
			Message:   "ResetData API rate limit exceeded",
			Details:   bodyStr,
			Retryable: true,
		}
	case http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable:
		return &ProviderError{
			Code:      "SERVICE_UNAVAILABLE",
			Message:   "ResetData API service unavailable",
			Details:   bodyStr,
			Retryable: true,
		}
	default:
		return &ProviderError{
			Code:      "API_ERROR",
			Message:   fmt.Sprintf("ResetData API error (status %d)", statusCode),
			Details:   bodyStr,
			Retryable: true,
		}
	}
}

// parseResponseForActions extracts actions from the response text
func (p *ResetDataProvider) parseResponseForActions(response string, request *TaskRequest) ([]TaskAction, []Artifact) {
	// Use the response parser to extract structured actions and artifacts
	parser := NewResponseParser()
	actions, artifacts := parser.ParseResponse(response)

	// If parser found concrete actions, return them
	if len(actions) > 0 {
		return actions, artifacts
	}

	// Otherwise, create a basic task analysis action as fallback
	action := TaskAction{
		Type:      "task_analysis",
		Target:    request.TaskTitle,
		Content:   response,
		Result:    "Task analyzed by ResetData model",
		Success:   true,
		Timestamp: time.Now(),
		Metadata: map[string]interface{}{
			"agent_role": request.AgentRole,
			"repository": request.Repository,
			"model":      p.config.DefaultModel,
		},
	}
	actions = append(actions, action)

	return actions, artifacts
}