wraith/internal/ai/service.go

package ai

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

	"github.com/pkg/browser"
)

// 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),
	}
}

// ChatResponse is returned to the frontend after a complete AI turn.
type ChatResponse struct {
	Text      string           `json:"text"`
	ToolCalls []ToolCallResult `json:"toolCalls,omitempty"`
}

// ToolCallResult captures a single tool invocation and its outcome.
type ToolCallResult struct {
	Name   string      `json:"name"`
	Input  interface{} `json:"input"`
	Result interface{} `json:"result"`
	Error  string      `json:"error,omitempty"`
}

// --- Auth ---

// StartLogin begins the OAuth PKCE flow, opening the browser for authentication.
func (s *AIService) StartLogin() error {
	done, err := s.oauth.StartLogin(browser.OpenURL)
	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.
// It returns the aggregated text and tool-call results from all iterations.
func (s *AIService) SendMessage(conversationId, text string) (*ChatResponse, 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 nil, 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.
// It returns the aggregated ChatResponse containing all text and tool-call results.
func (s *AIService) messageLoop(conversationId string) (*ChatResponse, error) {
	resp := &ChatResponse{}

	for iterations := 0; iterations < 20; iterations++ { // safety limit
		messages, err := s.conversations.GetMessages(conversationId)
		if err != nil {
			return nil, err
		}

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

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

		for event := range ch {
			switch event.Type {
			case "text_delta":
				textParts = append(textParts, event.Data)
			case "tool_use_start":
				// Finalize any previous tool call
				if currentToolID != "" {
					toolCalls = append(toolCalls, ContentBlock{
						Type:  "tool_use",
						ID:    currentToolID,
						Name:  currentToolName,
						Input: json.RawMessage(currentToolInput),
					})
				}
				currentToolID = event.ToolID
				currentToolName = event.ToolName
				currentToolInput = ""
			case "tool_use_delta":
				currentToolInput += event.Data
			case "done":
				// Finalize any in-progress tool call
				if currentToolID != "" {
					toolCalls = append(toolCalls, ContentBlock{
						Type:  "tool_use",
						ID:    currentToolID,
						Name:  currentToolName,
						Input: json.RawMessage(currentToolInput),
					})
					currentToolID = ""
					currentToolName = ""
					currentToolInput = ""
				}
				// 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 nil, fmt.Errorf("stream error: %s", event.Data)
			}
		}

		// Finalize any trailing tool call (if stream ended without a "done" event)
		if currentToolID != "" {
			toolCalls = append(toolCalls, ContentBlock{
				Type:  "tool_use",
				ID:    currentToolID,
				Name:  currentToolName,
				Input: json.RawMessage(currentToolInput),
			})
		}

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

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

		// Store assistant message
		assistantMsg := Message{
			Role:    "assistant",
			Content: assistantContent,
		}
		if err := s.conversations.AddMessage(conversationId, assistantMsg); err != nil {
			return nil, 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 resp, 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, dispatchErr := s.router.Dispatch(block.Name, block.Input)
			resultBlock := ContentBlock{
				Type:      "tool_result",
				ToolUseID: block.ID,
			}

			tcResult := ToolCallResult{
				Name:  block.Name,
				Input: block.Input,
			}

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

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

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

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

// --- Model ---

// GetModel returns the current Claude model identifier.
func (s *AIService) GetModel() string {
	return s.client.model
}

// SetModel changes the Claude model used for subsequent requests.
func (s *AIService) SetModel(model string) {
	s.client.model = model
}

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