3x-ui/web/service/ai_service.go

package service

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/mhsanaei/3x-ui/v2/database"
	"github.com/mhsanaei/3x-ui/v2/logger"
	"github.com/mhsanaei/3x-ui/v2/util/common"

	"github.com/google/generative-ai-go/genai"
	"google.golang.org/api/option"
)

// AIService provides Gemini AI integration for intelligent features
type AIService struct {
	client       *genai.Client
	model        *genai.GenerativeModel
	settingService SettingService
	inboundService InboundService
	serverService  ServerService
	
	// Cache and rate limiting
	cache          sync.Map // Cache for recent AI responses
	rateLimiter    map[int64]*rateLimitState
	rateLimiterMu  sync.RWMutex
	
	// Configuration
	enabled        bool
	apiKey         string
	maxTokens      int32
	temperature    float32
	cacheDuration  time.Duration
}

type rateLimitState struct {
	requests  int
	resetTime time.Time
	mu        sync.Mutex
}

type cacheEntry struct {
	response  string
	timestamp time.Time
}

// AIIntent represents the detected user intent from natural language
type AIIntent struct {
	Action      string                 `json:"action"`       // status, usage, inbound_list, client_add, etc.
	Parameters  map[string]interface{} `json:"parameters"`   // Extracted parameters
	Confidence  float64                `json:"confidence"`   // Confidence score 0-1
	NeedsAction bool                   `json:"needs_action"` // Whether this requires bot action
	Response    string                 `json:"response"`     // AI-generated response text
}

const (
	// Rate limiting
	maxRequestsPerMinute = 20
	maxRequestsPerHour   = 100
	
	// Cache settings
	defaultCacheDuration = 5 * time.Minute
	
	// AI Model settings
	defaultModel       = "gemini-1.5-flash"
	defaultMaxTokens   = 1024
	defaultTemperature = 0.7
	
	// System prompt for the AI
	systemPrompt = `You are an intelligent assistant for a VPN/Proxy management panel called 3X-UI.

Your role is to understand user commands in natural language and help manage their VPN server.

Available actions:
- server_status: Show CPU, memory, disk usage, uptime, Xray status
- server_usage: Display traffic statistics (total/upload/download)
- inbound_list: List all inbound configurations
- inbound_info: Get details about a specific inbound (by ID or remark)
- client_list: List clients for an inbound
- client_add: Add a new client to an inbound
- client_reset: Reset client traffic
- client_delete: Delete a client
- settings_backup: Create a backup
- settings_restore: Restore from backup
- help: Show available commands

When analyzing user messages:
1. Detect the intent/action they want to perform
2. Extract relevant parameters (inbound ID, client email, etc.)
3. Provide a confidence score (0-1) for your interpretation
4. Generate a helpful response

If the user's request is unclear or ambiguous, ask clarifying questions.
Always be concise, professional, and helpful.

Respond ONLY with valid JSON in this exact format:
{
  "action": "detected_action",
  "parameters": {"key": "value"},
  "confidence": 0.95,
  "needs_action": true,
  "response": "Your helpful response text"
}`
)

// NewAIService initializes the AI service with Gemini API
func NewAIService() *AIService {
	service := &AIService{
		rateLimiter:   make(map[int64]*rateLimitState),
		maxTokens:     defaultMaxTokens,
		temperature:   defaultTemperature,
		cacheDuration: defaultCacheDuration,
	}
	
	// Load settings from database
	if err := service.loadSettings(); err != nil {
		logger.Warning("AI Service: Failed to load settings:", err)
		return service
	}
	
	// Initialize client if enabled
	if service.enabled && service.apiKey != "" {
		if err := service.initClient(); err != nil {
			logger.Warning("AI Service: Failed to initialize Gemini client:", err)
			service.enabled = false
		}
	}
	
	return service
}

// loadSettings loads AI configuration from database
func (s *AIService) loadSettings() error {
	db := database.GetDB()
	
	// Check if AI is enabled
	enabledStr, err := s.settingService.GetAISetting("AIEnabled")
	if err == nil && enabledStr == "true" {
		s.enabled = true
	}
	
	// Load API key
	apiKey, err := s.settingService.GetAISetting("AIApiKey")
	if err == nil && apiKey != "" {
		s.apiKey = apiKey
	}
	
	// Load optional settings
	if maxTokensStr, err := s.settingService.GetAISetting("AIMaxTokens"); err == nil {
		var maxTokens int
		if err := json.Unmarshal([]byte(maxTokensStr), &maxTokens); err == nil {
			s.maxTokens = int32(maxTokens)
		}
	}
	
	if tempStr, err := s.settingService.GetAISetting("AITemperature"); err == nil {
		var temp float64
		if err := json.Unmarshal([]byte(tempStr), &temp); err == nil {
			s.temperature = float32(temp)
		}
	}
	
	logger.Debug("AI Service settings loaded - Enabled:", s.enabled, "API Key present:", s.apiKey != "")
	
	return db.Error
}

// initClient initializes the Gemini AI client
func (s *AIService) initClient() error {
	ctx := context.Background()
	
	client, err := genai.NewClient(ctx, option.WithAPIKey(s.apiKey))
	if err != nil {
		return fmt.Errorf("failed to create Gemini client: %w", err)
	}
	
	s.client = client
	s.model = client.GenerativeModel(defaultModel)
	
	// Configure model parameters
	s.model.SetMaxOutputTokens(s.maxTokens)
	s.model.SetTemperature(s.temperature)
	s.model.SystemInstruction = &genai.Content{
		Parts: []genai.Part{genai.Text(systemPrompt)},
	}
	
	// Configure safety settings to be less restrictive for technical content
	s.model.SafetySettings = []*genai.SafetySetting{
		{
			Category:  genai.HarmCategoryHarassment,
			Threshold: genai.HarmBlockMediumAndAbove,
		},
		{
			Category:  genai.HarmCategoryHateSpeech,
			Threshold: genai.HarmBlockMediumAndAbove,
		},
		{
			Category:  genai.HarmCategoryDangerousContent,
			Threshold: genai.HarmBlockMediumAndAbove,
		},
	}
	
	logger.Info("AI Service: Gemini client initialized successfully")
	return nil
}

// IsEnabled checks if AI service is currently enabled
func (s *AIService) IsEnabled() bool {
	return s.enabled && s.client != nil
}

// ProcessMessage processes a natural language message and returns detected intent
func (s *AIService) ProcessMessage(ctx context.Context, userID int64, message string) (*AIIntent, error) {
	if !s.IsEnabled() {
		return nil, fmt.Errorf("AI service is not enabled")
	}
	
	// Check rate limiting
	if !s.checkRateLimit(userID) {
		return nil, fmt.Errorf("rate limit exceeded, please try again later")
	}
	
	// Check cache first
	cacheKey := fmt.Sprintf("%d:%s", userID, strings.ToLower(strings.TrimSpace(message)))
	if cached, ok := s.cache.Load(cacheKey); ok {
		entry := cached.(cacheEntry)
		if time.Since(entry.timestamp) < s.cacheDuration {
			logger.Debug("AI Service: Cache hit for user", userID)
			var intent AIIntent
			if err := json.Unmarshal([]byte(entry.response), &intent); err == nil {
				return &intent, nil
			}
		}
	}
	
	// Generate AI response
	intent, err := s.generateIntent(ctx, message)
	if err != nil {
		return nil, fmt.Errorf("failed to generate intent: %w", err)
	}
	
	// Cache the response
	responseJSON, _ := json.Marshal(intent)
	s.cache.Store(cacheKey, cacheEntry{
		response:  string(responseJSON),
		timestamp: time.Now(),
	})
	
	logger.Debug("AI Service: Processed message for user", userID, "- Action:", intent.Action, "Confidence:", intent.Confidence)
	
	return intent, nil
}

// generateIntent calls Gemini API to analyze the message
func (s *AIService) generateIntent(ctx context.Context, message string) (*AIIntent, error) {
	// Create prompt with user message
	prompt := fmt.Sprintf("User message: %s\n\nAnalyze this message and respond with the JSON format specified in the system prompt.", message)
	
	// Set timeout for API call
	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()
	
	// Generate response
	resp, err := s.model.GenerateContent(ctx, genai.Text(prompt))
	if err != nil {
		return nil, fmt.Errorf("Gemini API error: %w", err)
	}
	
	// Extract text response
	if len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 {
		return nil, fmt.Errorf("no response from Gemini API")
	}
	
	responseText := fmt.Sprintf("%v", resp.Candidates[0].Content.Parts[0])
	
	// Parse JSON response
	intent, err := s.parseIntentResponse(responseText)
	if err != nil {
		// If parsing fails, try to extract JSON from markdown code block
		if cleaned := extractJSONFromMarkdown(responseText); cleaned != "" {
			intent, err = s.parseIntentResponse(cleaned)
		}
		
		if err != nil {
			logger.Warning("AI Service: Failed to parse response:", err, "Raw:", responseText)
			// Return a fallback intent
			return &AIIntent{
				Action:      "unknown",
				Parameters:  make(map[string]interface{}),
				Confidence:  0.0,
				NeedsAction: false,
				Response:    "I couldn't understand your request. Please try rephrasing or use /help to see available commands.",
			}, nil
		}
	}
	
	return intent, nil
}

// parseIntentResponse parses the JSON response from Gemini
func (s *AIService) parseIntentResponse(responseText string) (*AIIntent, error) {
	var intent AIIntent
	
	// Try to parse as JSON
	if err := json.Unmarshal([]byte(responseText), &intent); err != nil {
		return nil, fmt.Errorf("invalid JSON response: %w", err)
	}
	
	// Validate required fields
	if intent.Action == "" {
		intent.Action = "unknown"
	}
	if intent.Parameters == nil {
		intent.Parameters = make(map[string]interface{})
	}
	if intent.Response == "" {
		intent.Response = "Processing your request..."
	}
	
	return &intent, nil
}

// checkRateLimit checks if user has exceeded rate limits
func (s *AIService) checkRateLimit(userID int64) bool {
	now := time.Now()
	
	s.rateLimiterMu.Lock()
	defer s.rateLimiterMu.Unlock()
	
	state, exists := s.rateLimiter[userID]
	if !exists {
		state = &rateLimitState{
			requests:  1,
			resetTime: now.Add(time.Minute),
		}
		s.rateLimiter[userID] = state
		return true
	}
	
	state.mu.Lock()
	defer state.mu.Unlock()
	
	// Reset if time window passed
	if now.After(state.resetTime) {
		state.requests = 1
		state.resetTime = now.Add(time.Minute)
		return true
	}
	
	// Check limit
	if state.requests >= maxRequestsPerMinute {
		return false
	}
	
	state.requests++
	return true
}

// GetContextForUser generates context information to enhance AI responses
func (s *AIService) GetContextForUser(userID int64) string {
	var context strings.Builder
	
	// Add server status
	if serverInfo, err := s.serverService.GetStatus(true); err == nil {
		context.WriteString(fmt.Sprintf("Server CPU: %.1f%%, Memory: %.1f%%, ", 
			serverInfo.Cpu, serverInfo.Mem))
	}
	
	// Add inbound count
	if db := database.GetDB(); db != nil {
		var count int64
		db.Model(&struct{}{}).Count(&count)
		context.WriteString(fmt.Sprintf("Total inbounds: %d. ", count))
	}
	
	return context.String()
}

// Close gracefully shuts down the AI service
func (s *AIService) Close() error {
	if s.client != nil {
		return s.client.Close()
	}
	return nil
}

// extractJSONFromMarkdown extracts JSON from markdown code blocks
func extractJSONFromMarkdown(text string) string {
	// Try to find JSON in ```json or ``` blocks
	patterns := []string{
		"```json\n(.+?)\n```",
		"```\n(.+?)\n```",
	}
	
	for _, pattern := range patterns {
		if idx := strings.Index(text, "```"); idx != -1 {
			// Find closing ```
			if endIdx := strings.Index(text[idx+3:], "```"); endIdx != -1 {
				extracted := text[idx+3 : idx+3+endIdx]
				// Remove "json" if present at start
				extracted = strings.TrimPrefix(extracted, "json")
				extracted = strings.TrimSpace(extracted)
				return extracted
			}
		}
	}
	
	// Try to find JSON by looking for { and }
	if start := strings.Index(text, "{"); start != -1 {
		if end := strings.LastIndex(text, "}"); end != -1 && end > start {
			return text[start : end+1]
		}
	}
	
	return ""
}