wraith/internal/ai/service.go

package ai

import (
	"encoding/json"
	"fmt"
	"log/slog"
	"sync"
)

// SystemPrompt is the system prompt given to Claude for copilot interactions.
const SystemPrompt = `You are the XO (Executive Officer) aboard the Wraith command station. The Commander (human operator) works alongside you managing remote servers and workstations.

You have direct access to all active sessions through your tools:
- SSH terminals: read output, type commands, navigate filesystems
- SFTP: read and write remote files
- RDP desktops: see the screen, click, type, interact with any GUI application
- Session management: open new connections, close sessions

When given a task:
1. Assess what sessions and access you need
2. Execute efficiently — don't ask for permission to use tools, just use them
3. Report what you found or did, with relevant details
4. If something fails, diagnose and try an alternative approach

You are not an assistant answering questions. You are an operator executing missions. Act decisively. Use your tools. Report results.`

// AIService is the main AI copilot service exposed to the Wails frontend.
type AIService struct {
	oauth         *OAuthManager
	client        *ClaudeClient
	router        *ToolRouter
	conversations *ConversationManager
	buffers       map[string]*TerminalBuffer
	mu            sync.RWMutex
}

// NewAIService creates the AI service with all sub-components.
func NewAIService(oauth *OAuthManager, router *ToolRouter, convMgr *ConversationManager) *AIService {
	client := NewClaudeClient(oauth, "")
	return &AIService{
		oauth:         oauth,
		client:        client,
		router:        router,
		conversations: convMgr,
		buffers:       make(map[string]*TerminalBuffer),
	}
}

// --- Auth ---

// StartLogin begins the OAuth PKCE flow, opening the browser for authentication.
func (s *AIService) StartLogin() error {
	done, err := s.oauth.StartLogin(nil) // nil openURL = no browser auto-open
	if err != nil {
		return err
	}
	// Wait for callback in a goroutine to avoid blocking the UI
	go func() {
		if err := <-done; err != nil {
			slog.Error("oauth login failed", "error", err)
		} else {
			slog.Info("oauth login completed")
		}
	}()
	return nil
}

// IsAuthenticated returns whether the user has valid OAuth tokens.
func (s *AIService) IsAuthenticated() bool {
	return s.oauth.IsAuthenticated()
}

// Logout clears stored OAuth tokens.
func (s *AIService) Logout() error {
	return s.oauth.Logout()
}

// --- Conversations ---

// NewConversation creates a new AI conversation and returns its ID.
func (s *AIService) NewConversation() (string, error) {
	conv, err := s.conversations.Create(s.client.model)
	if err != nil {
		return "", err
	}
	return conv.ID, nil
}

// ListConversations returns all conversations.
func (s *AIService) ListConversations() ([]ConversationSummary, error) {
	return s.conversations.List()
}

// DeleteConversation removes a conversation.
func (s *AIService) DeleteConversation(id string) error {
	return s.conversations.Delete(id)
}

// --- Chat ---

// SendMessage sends a user message in a conversation and processes the AI response.
// Tool calls are automatically dispatched and results fed back to the model.
// This method blocks until the full response (including any tool use loops) is complete.
func (s *AIService) SendMessage(conversationId, text string) error {
	// Add user message to conversation
	userMsg := Message{
		Role: "user",
		Content: []ContentBlock{
			{Type: "text", Text: text},
		},
	}
	if err := s.conversations.AddMessage(conversationId, userMsg); err != nil {
		return fmt.Errorf("store user message: %w", err)
	}

	// Run the message loop (handles tool use)
	return s.messageLoop(conversationId)
}

// messageLoop sends the conversation to Claude and handles tool use loops.
func (s *AIService) messageLoop(conversationId string) error {
	for iterations := 0; iterations < 20; iterations++ { // safety limit
		messages, err := s.conversations.GetMessages(conversationId)
		if err != nil {
			return err
		}

		ch, err := s.client.SendMessage(messages, CopilotTools, SystemPrompt)
		if err != nil {
			return fmt.Errorf("send to claude: %w", err)
		}

		// Collect the response
		var textParts []string
		var toolCalls []ContentBlock
		var currentToolInput string

		for event := range ch {
			switch event.Type {
			case "text_delta":
				textParts = append(textParts, event.Data)
			case "tool_use_start":
				currentToolInput = ""
			case "tool_use_delta":
				currentToolInput += event.Data
			case "done":
				// Parse usage if available
				var delta struct {
					Usage Usage `json:"usage"`
				}
				if json.Unmarshal([]byte(event.Data), &delta) == nil && delta.Usage.OutputTokens > 0 {
					s.conversations.UpdateTokenUsage(conversationId, delta.Usage.InputTokens, delta.Usage.OutputTokens)
				}
			case "error":
				return fmt.Errorf("stream error: %s", event.Data)
			}

			// When a tool_use block completes, we need to check content_block_stop
			// But since we accumulate, we'll finalize after the stream ends
			_ = toolCalls // kept for the final assembly below
		}

		// Build the assistant message
		var assistantContent []ContentBlock
		if len(textParts) > 0 {
			fullText := ""
			for _, p := range textParts {
				fullText += p
			}
			if fullText != "" {
				assistantContent = append(assistantContent, ContentBlock{
					Type: "text",
					Text: fullText,
				})
			}
		}
		for _, tc := range toolCalls {
			assistantContent = append(assistantContent, tc)
		}

		if len(assistantContent) == 0 {
			return nil // empty response
		}

		// Store assistant message
		assistantMsg := Message{
			Role:    "assistant",
			Content: assistantContent,
		}
		if err := s.conversations.AddMessage(conversationId, assistantMsg); err != nil {
			return fmt.Errorf("store assistant message: %w", err)
		}

		// If there were tool calls, dispatch them and continue the loop
		hasToolUse := false
		for _, block := range assistantContent {
			if block.Type == "tool_use" {
				hasToolUse = true
				break
			}
		}

		if !hasToolUse {
			return nil // done, no tool use to process
		}

		// Dispatch tool calls and create tool_result message
		var toolResults []ContentBlock
		for _, block := range assistantContent {
			if block.Type != "tool_use" {
				continue
			}

			result, err := s.router.Dispatch(block.Name, block.Input)
			resultBlock := ContentBlock{
				Type:      "tool_result",
				ToolUseID: block.ID,
			}

			if err != nil {
				resultBlock.IsError = true
				resultBlock.Content = []ContentBlock{
					{Type: "text", Text: err.Error()},
				}
			} else {
				resultJSON, _ := json.Marshal(result)
				resultBlock.Content = []ContentBlock{
					{Type: "text", Text: string(resultJSON)},
				}
			}
			toolResults = append(toolResults, resultBlock)
		}

		toolResultMsg := Message{
			Role:    "user",
			Content: toolResults,
		}
		if err := s.conversations.AddMessage(conversationId, toolResultMsg); err != nil {
			return fmt.Errorf("store tool results: %w", err)
		}

		// Continue the loop to let Claude process tool results
	}

	return fmt.Errorf("exceeded maximum tool use iterations")
}

// --- Terminal Buffer Management ---

// GetBuffer returns the terminal output buffer for a session, creating it if needed.
func (s *AIService) GetBuffer(sessionId string) *TerminalBuffer {
	s.mu.Lock()
	defer s.mu.Unlock()

	buf, ok := s.buffers[sessionId]
	if !ok {
		buf = NewTerminalBuffer(200)
		s.buffers[sessionId] = buf
	}
	return buf
}

// RemoveBuffer removes the terminal buffer for a session.
func (s *AIService) RemoveBuffer(sessionId string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.buffers, sessionId)
}