wraith/internal/ai/router.go

package ai

import (
	"encoding/json"
	"fmt"
)

// ToolRouter dispatches tool calls to the appropriate service.
// Services are stored as interface{} to avoid circular imports between packages.
type ToolRouter struct {
	ssh         interface{} // *ssh.SSHService
	sftp        interface{} // *sftp.SFTPService
	rdp         interface{} // *rdp.RDPService
	sessions    interface{} // *session.Manager
	connections interface{} // *connections.ConnectionService
	aiService   interface{} // *AIService — for terminal buffer access
}

// NewToolRouter creates an empty ToolRouter. Call SetServices to wire in backends.
func NewToolRouter() *ToolRouter {
	return &ToolRouter{}
}

// SetServices wires the router to actual service implementations.
func (r *ToolRouter) SetServices(ssh, sftp, rdp, sessions, connections interface{}) {
	r.ssh = ssh
	r.sftp = sftp
	r.rdp = rdp
	r.sessions = sessions
	r.connections = connections
}

// SetAIService wires the router to the AI service for terminal buffer access.
func (r *ToolRouter) SetAIService(aiService interface{}) {
	r.aiService = aiService
}

// sshWriter is the interface we need from SSHService for terminal_write.
type sshWriter interface {
	Write(sessionID string, data string) error
}

// sshSessionLister is the interface for listing SSH sessions.
type sshSessionLister interface {
	ListSessions() interface{}
}

// sftpLister is the interface for SFTP list operations.
type sftpLister interface {
	List(sessionID string, path string) (interface{}, error)
}

// sftpReader is the interface for SFTP read operations.
type sftpReader interface {
	ReadFile(sessionID string, path string) (string, error)
}

// sftpWriter is the interface for SFTP write operations.
type sftpWriter interface {
	WriteFile(sessionID string, path string, content string) error
}

// rdpFrameGetter is the interface for getting RDP screenshots.
type rdpFrameGetter interface {
	GetFrame(sessionID string) ([]byte, error)
	GetSessionInfo(sessionID string) (interface{}, error)
}

// rdpMouseSender is the interface for RDP mouse events.
type rdpMouseSender interface {
	SendMouse(sessionID string, x, y int, flags uint32) error
}

// rdpKeySender is the interface for RDP key events.
type rdpKeySender interface {
	SendKey(sessionID string, scancode uint32, pressed bool) error
}

// rdpClipboardSender is the interface for RDP clipboard.
type rdpClipboardSender interface {
	SendClipboard(sessionID string, data string) error
}

// rdpSessionLister is the interface for listing RDP sessions.
type rdpSessionLister interface {
	ListSessions() interface{}
}

// sessionLister is the interface for listing sessions.
type sessionLister interface {
	List() interface{}
}

// bufferProvider provides terminal output buffers.
type bufferProvider interface {
	GetBuffer(sessionId string) *TerminalBuffer
}

// Dispatch routes a tool call to the appropriate handler.
func (r *ToolRouter) Dispatch(toolName string, input json.RawMessage) (interface{}, error) {
	switch toolName {
	case "terminal_write":
		return r.handleTerminalWrite(input)
	case "terminal_read":
		return r.handleTerminalRead(input)
	case "terminal_cwd":
		return r.handleTerminalCwd(input)
	case "sftp_list":
		return r.handleSFTPList(input)
	case "sftp_read":
		return r.handleSFTPRead(input)
	case "sftp_write":
		return r.handleSFTPWrite(input)
	case "rdp_screenshot":
		return r.handleRDPScreenshot(input)
	case "rdp_click":
		return r.handleRDPClick(input)
	case "rdp_doubleclick":
		return r.handleRDPDoubleClick(input)
	case "rdp_type":
		return r.handleRDPType(input)
	case "rdp_keypress":
		return r.handleRDPKeypress(input)
	case "rdp_scroll":
		return r.handleRDPScroll(input)
	case "rdp_move":
		return r.handleRDPMove(input)
	case "list_sessions":
		return r.handleListSessions(input)
	case "connect_ssh":
		return r.handleConnectSSH(input)
	case "disconnect":
		return r.handleDisconnect(input)
	default:
		return nil, fmt.Errorf("unknown tool: %s", toolName)
	}
}

func (r *ToolRouter) handleTerminalWrite(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		Text      string `json:"text"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	w, ok := r.ssh.(sshWriter)
	if !ok || r.ssh == nil {
		return nil, fmt.Errorf("SSH service not available")
	}

	if err := w.Write(params.SessionID, params.Text); err != nil {
		return nil, err
	}
	return map[string]string{"status": "ok"}, nil
}

func (r *ToolRouter) handleTerminalRead(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		Lines     int    `json:"lines"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}
	if params.Lines <= 0 {
		params.Lines = 50
	}

	bp, ok := r.aiService.(bufferProvider)
	if !ok || r.aiService == nil {
		return nil, fmt.Errorf("terminal buffer not available")
	}

	buf := bp.GetBuffer(params.SessionID)
	lines := buf.ReadLast(params.Lines)
	return map[string]interface{}{
		"lines": lines,
		"count": len(lines),
	}, nil
}

func (r *ToolRouter) handleTerminalCwd(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	// terminal_cwd works by writing "pwd" to the terminal and reading output
	w, ok := r.ssh.(sshWriter)
	if !ok || r.ssh == nil {
		return nil, fmt.Errorf("SSH service not available")
	}

	if err := w.Write(params.SessionID, "pwd\n"); err != nil {
		return nil, fmt.Errorf("send pwd: %w", err)
	}

	return map[string]string{
		"status": "pwd command sent — read terminal output for result",
	}, nil
}

func (r *ToolRouter) handleSFTPList(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		Path      string `json:"path"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	l, ok := r.sftp.(sftpLister)
	if !ok || r.sftp == nil {
		return nil, fmt.Errorf("SFTP service not available")
	}

	return l.List(params.SessionID, params.Path)
}

func (r *ToolRouter) handleSFTPRead(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		Path      string `json:"path"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	reader, ok := r.sftp.(sftpReader)
	if !ok || r.sftp == nil {
		return nil, fmt.Errorf("SFTP service not available")
	}

	content, err := reader.ReadFile(params.SessionID, params.Path)
	if err != nil {
		return nil, err
	}
	return map[string]string{"content": content}, nil
}

func (r *ToolRouter) handleSFTPWrite(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		Path      string `json:"path"`
		Content   string `json:"content"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	w, ok := r.sftp.(sftpWriter)
	if !ok || r.sftp == nil {
		return nil, fmt.Errorf("SFTP service not available")
	}

	if err := w.WriteFile(params.SessionID, params.Path, params.Content); err != nil {
		return nil, err
	}
	return map[string]string{"status": "ok"}, nil
}

func (r *ToolRouter) handleRDPScreenshot(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	fg, ok := r.rdp.(rdpFrameGetter)
	if !ok || r.rdp == nil {
		return nil, fmt.Errorf("RDP service not available")
	}

	frame, err := fg.GetFrame(params.SessionID)
	if err != nil {
		return nil, err
	}

	// Get session info for dimensions
	info, err := fg.GetSessionInfo(params.SessionID)
	if err != nil {
		return nil, err
	}

	// Try to get dimensions from session config
	type configGetter interface {
		GetConfig() (int, int)
	}
	width, height := 1920, 1080
	if cg, ok := info.(configGetter); ok {
		width, height = cg.GetConfig()
	}

	// Encode as JPEG
	jpeg, err := EncodeScreenshot(frame, width, height, 1280, 720, 75)
	if err != nil {
		return nil, fmt.Errorf("encode screenshot: %w", err)
	}

	return map[string]interface{}{
		"image":  jpeg,
		"width":  width,
		"height": height,
	}, nil
}

func (r *ToolRouter) handleRDPClick(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		X         int    `json:"x"`
		Y         int    `json:"y"`
		Button    string `json:"button"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	ms, ok := r.rdp.(rdpMouseSender)
	if !ok || r.rdp == nil {
		return nil, fmt.Errorf("RDP service not available")
	}

	var buttonFlag uint32 = 0x1000 // left
	switch params.Button {
	case "right":
		buttonFlag = 0x2000
	case "middle":
		buttonFlag = 0x4000
	}

	// Press
	if err := ms.SendMouse(params.SessionID, params.X, params.Y, buttonFlag|0x8000); err != nil {
		return nil, err
	}
	// Release
	if err := ms.SendMouse(params.SessionID, params.X, params.Y, buttonFlag); err != nil {
		return nil, err
	}
	return map[string]string{"status": "ok"}, nil
}

func (r *ToolRouter) handleRDPDoubleClick(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		X         int    `json:"x"`
		Y         int    `json:"y"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	ms, ok := r.rdp.(rdpMouseSender)
	if !ok || r.rdp == nil {
		return nil, fmt.Errorf("RDP service not available")
	}

	// Two clicks
	for i := 0; i < 2; i++ {
		if err := ms.SendMouse(params.SessionID, params.X, params.Y, 0x1000|0x8000); err != nil {
			return nil, err
		}
		if err := ms.SendMouse(params.SessionID, params.X, params.Y, 0x1000); err != nil {
			return nil, err
		}
	}
	return map[string]string{"status": "ok"}, nil
}

func (r *ToolRouter) handleRDPType(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		Text      string `json:"text"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	cs, ok := r.rdp.(rdpClipboardSender)
	if !ok || r.rdp == nil {
		return nil, fmt.Errorf("RDP service not available")
	}

	// Send text via clipboard, then simulate Ctrl+V
	if err := cs.SendClipboard(params.SessionID, params.Text); err != nil {
		return nil, err
	}

	ks, ok := r.rdp.(rdpKeySender)
	if !ok {
		return nil, fmt.Errorf("RDP key service not available")
	}

	// Ctrl down, V down, V up, Ctrl up
	if err := ks.SendKey(params.SessionID, 0x001D, true); err != nil {
		return nil, err
	}
	if err := ks.SendKey(params.SessionID, 0x002F, true); err != nil {
		return nil, err
	}
	if err := ks.SendKey(params.SessionID, 0x002F, false); err != nil {
		return nil, err
	}
	if err := ks.SendKey(params.SessionID, 0x001D, false); err != nil {
		return nil, err
	}

	return map[string]string{"status": "ok"}, nil
}

func (r *ToolRouter) handleRDPKeypress(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		Key       string `json:"key"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	ks, ok := r.rdp.(rdpKeySender)
	if !ok || r.rdp == nil {
		return nil, fmt.Errorf("RDP service not available")
	}

	// Simple key name to scancode mapping for common keys
	keyMap := map[string]uint32{
		"Enter":     0x001C,
		"Tab":       0x000F,
		"Escape":    0x0001,
		"Backspace": 0x000E,
		"Delete":    0xE053,
		"Space":     0x0039,
		"Up":        0xE048,
		"Down":      0xE050,
		"Left":      0xE04B,
		"Right":     0xE04D,
	}

	scancode, ok := keyMap[params.Key]
	if !ok {
		return nil, fmt.Errorf("unknown key: %s", params.Key)
	}

	if err := ks.SendKey(params.SessionID, scancode, true); err != nil {
		return nil, err
	}
	if err := ks.SendKey(params.SessionID, scancode, false); err != nil {
		return nil, err
	}
	return map[string]string{"status": "ok"}, nil
}

func (r *ToolRouter) handleRDPScroll(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		X         int    `json:"x"`
		Y         int    `json:"y"`
		Direction string `json:"direction"`
		Clicks    int    `json:"clicks"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}
	if params.Clicks <= 0 {
		params.Clicks = 3
	}

	ms, ok := r.rdp.(rdpMouseSender)
	if !ok || r.rdp == nil {
		return nil, fmt.Errorf("RDP service not available")
	}

	var flags uint32 = 0x0200 // wheel flag
	if params.Direction == "down" {
		flags |= 0x0100 // negative flag
	}

	for i := 0; i < params.Clicks; i++ {
		if err := ms.SendMouse(params.SessionID, params.X, params.Y, flags); err != nil {
			return nil, err
		}
	}
	return map[string]string{"status": "ok"}, nil
}

func (r *ToolRouter) handleRDPMove(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
		X         int    `json:"x"`
		Y         int    `json:"y"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	ms, ok := r.rdp.(rdpMouseSender)
	if !ok || r.rdp == nil {
		return nil, fmt.Errorf("RDP service not available")
	}

	if err := ms.SendMouse(params.SessionID, params.X, params.Y, 0x0800); err != nil {
		return nil, err
	}
	return map[string]string{"status": "ok"}, nil
}

func (r *ToolRouter) handleListSessions(_ json.RawMessage) (interface{}, error) {
	result := map[string]interface{}{
		"ssh": []interface{}{},
		"rdp": []interface{}{},
	}

	if r.sessions != nil {
		if sl, ok := r.sessions.(sessionLister); ok {
			result["all"] = sl.List()
		}
	}

	return result, nil
}

func (r *ToolRouter) handleConnectSSH(input json.RawMessage) (interface{}, error) {
	var params struct {
		ConnectionID int64 `json:"connectionId"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	// This will be wired to the app-level connect logic later
	return nil, fmt.Errorf("connect_ssh requires app-level wiring — not yet available via tool dispatch")
}

func (r *ToolRouter) handleDisconnect(input json.RawMessage) (interface{}, error) {
	var params struct {
		SessionID string `json:"sessionId"`
	}
	if err := json.Unmarshal(input, &params); err != nil {
		return nil, fmt.Errorf("parse input: %w", err)
	}

	// Try SSH first
	type disconnecter interface {
		Disconnect(sessionID string) error
	}

	if d, ok := r.ssh.(disconnecter); ok {
		if err := d.Disconnect(params.SessionID); err == nil {
			return map[string]string{"status": "disconnected", "protocol": "ssh"}, nil
		}
	}

	if d, ok := r.rdp.(disconnecter); ok {
		if err := d.Disconnect(params.SessionID); err == nil {
			return map[string]string{"status": "disconnected", "protocol": "rdp"}, nil
		}
	}

	return nil, fmt.Errorf("session %s not found in SSH or RDP", params.SessionID)
}