goclaw/internal/tools/heartbeat.go

package tools

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"

	"github.com/nextlevelbuilder/goclaw/internal/store"
)

// HeartbeatTool lets agents manage their own heartbeat configuration.
type HeartbeatTool struct {
	hbStore    store.HeartbeatStore
	permStore  store.ConfigPermissionStore
	agentStore store.AgentStore
	wakeFn     func(uuid.UUID) // triggers immediate heartbeat run
}

// NewHeartbeatTool creates a heartbeat management tool.
func NewHeartbeatTool(hb store.HeartbeatStore, perms store.ConfigPermissionStore) *HeartbeatTool {
	return &HeartbeatTool{hbStore: hb, permStore: perms}
}

// SetAgentStore sets the agent store for HEARTBEAT.md read/write.
func (t *HeartbeatTool) SetAgentStore(as store.AgentStore) {
	t.agentStore = as
}

// SetWakeFn sets the function called when the "test" action triggers an immediate run.
func (t *HeartbeatTool) SetWakeFn(fn func(uuid.UUID)) {
	t.wakeFn = fn
}

func (t *HeartbeatTool) Name() string { return "heartbeat" }

func (t *HeartbeatTool) Description() string {
	return `Manage agent heartbeat — periodic proactive check-in.

ACTIONS:
- status: Show heartbeat status (enabled, last/next run, counts)
- get: Get full heartbeat configuration
- set: Create or update heartbeat config (interval, active_hours, ack_max_chars, etc.)
- toggle: Enable or disable heartbeat (enabled: true/false)
- set_checklist: Set HEARTBEAT.md content (the checklist the agent follows each run)
- get_checklist: Read current HEARTBEAT.md content
- test: Trigger an immediate heartbeat run (background)
- logs: View heartbeat run history (limit, offset)

EXAMPLES:
  {"action":"status"}
  {"action":"set","interval":1800,"channel":"telegram","chat_id":"-100123456"}
  {"action":"toggle","enabled":true}
  {"action":"set_checklist","content":"# Heartbeat Checklist\n- Check server status\n- Report any alerts"}
  {"action":"test"}
  {"action":"logs","limit":10}`
}

func (t *HeartbeatTool) Parameters() map[string]any {
	return map[string]any{
		"type": "object",
		"properties": map[string]any{
			"action": map[string]any{
				"type":        "string",
				"enum":        []string{"status", "get", "set", "toggle", "set_checklist", "get_checklist", "test", "logs"},
				"description": "Action to perform",
			},
			"enabled":          map[string]any{"type": "boolean", "description": "For toggle: enable/disable heartbeat"},
			"interval":         map[string]any{"type": "integer", "description": "Heartbeat interval in seconds (min 300)"},
			"prompt":           map[string]any{"type": "string", "description": "Custom heartbeat prompt (empty = default)"},
			"ack_max_chars":    map[string]any{"type": "integer", "description": "Max chars for suppression threshold (default 300)"},
			"max_retries":      map[string]any{"type": "integer", "description": "Max retry attempts on failure (default 2)"},
			"isolated_session": map[string]any{"type": "boolean", "description": "Use isolated session per run (default true)"},
			"light_context":    map[string]any{"type": "boolean", "description": "Skip loading context files, only inject heartbeat checklist (default false)"},
			"active_hours":     map[string]any{"type": "string", "description": "Active hours range, e.g. '08:00-22:00'"},
			"timezone":         map[string]any{"type": "string", "description": "IANA timezone for active hours, e.g. 'Asia/Ho_Chi_Minh'"},
			"channel":          map[string]any{"type": "string", "description": "Delivery channel name (auto-filled from current context if empty)"},
			"chat_id":          map[string]any{"type": "string", "description": "Delivery target chat ID (auto-filled from current context if empty)"},
			"content":          map[string]any{"type": "string", "description": "For set_checklist: HEARTBEAT.md content (the checklist the agent follows each run)"},
			"limit":            map[string]any{"type": "integer", "description": "For logs: max entries to return"},
			"offset":           map[string]any{"type": "integer", "description": "For logs: pagination offset"},
		},
		"required": []string{"action"},
	}
}

func (t *HeartbeatTool) Execute(ctx context.Context, args map[string]any) *Result {
	action, _ := args["action"].(string)
	if action == "" {
		return ErrorResult("action parameter is required")
	}

	agentID := store.AgentIDFromContext(ctx)
	if agentID == uuid.Nil {
		return ErrorResult("no agent context")
	}

	// Permission check for mutation actions.
	switch action {
	case "set", "toggle", "set_checklist":
		if err := t.checkPermission(ctx, agentID); err != nil {
			return ErrorResult(err.Error())
		}
	}

	switch action {
	case "status":
		return t.handleStatus(ctx, agentID)
	case "get":
		return t.handleGet(ctx, agentID)
	case "set":
		return t.handleSet(ctx, agentID, args)
	case "toggle":
		enabled, ok := args["enabled"].(bool)
		if !ok {
			return ErrorResult("toggle requires 'enabled' parameter (boolean)")
		}
		return t.handleToggle(ctx, agentID, enabled)
	case "set_checklist":
		content, _ := args["content"].(string)
		return t.handleSetChecklist(ctx, agentID, content)
	case "get_checklist":
		return t.handleGetChecklist(ctx, agentID)
	case "test":
		return t.handleTest(agentID)
	case "logs":
		limit := intArg(args, "limit", 10)
		offset := intArg(args, "offset", 0)
		return t.handleLogs(ctx, agentID, limit, offset)
	default:
		return ErrorResult(fmt.Sprintf("unknown action %q — use status/get/set/toggle/set_checklist/get_checklist/test/logs", action))
	}
}

func (t *HeartbeatTool) handleStatus(ctx context.Context, agentID uuid.UUID) *Result {
	hb, err := t.hbStore.Get(ctx, agentID)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return NewResult("Heartbeat not configured for this agent.")
		}
		return ErrorResult(err.Error())
	}
	return NewResult(fmt.Sprintf(
		"Heartbeat: %s | interval: %ds | runs: %d | suppressed: %d | last: %s (%s) | next: %s",
		boolLabel(hb.Enabled, "enabled", "disabled"),
		hb.IntervalSec, hb.RunCount, hb.SuppressCount,
		fmtTimePtr(hb.LastRunAt), derefStr(hb.LastStatus, "never"),
		fmtTimePtr(hb.NextRunAt),
	))
}

func (t *HeartbeatTool) handleGet(ctx context.Context, agentID uuid.UUID) *Result {
	hb, err := t.hbStore.Get(ctx, agentID)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return NewResult("Heartbeat not configured for this agent.")
		}
		return ErrorResult(err.Error())
	}
	data, _ := json.MarshalIndent(hb, "", "  ")
	return NewResult(string(data))
}

func (t *HeartbeatTool) handleSet(ctx context.Context, agentID uuid.UUID, args map[string]any) *Result {
	hb, err := t.hbStore.Get(ctx, agentID)
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return ErrorResult(err.Error())
	}
	if hb == nil {
		hb = &store.AgentHeartbeat{
			AgentID:         agentID,
			IntervalSec:     1800,
			IsolatedSession: true,
			AckMaxChars:     300,
			MaxRetries:      2,
		}
	}

	if v, ok := args["enabled"].(bool); ok {
		hb.Enabled = v
	}
	if v := intArg(args, "interval", 0); v > 0 {
		if v < 300 {
			return ErrorResult("minimum interval is 300 seconds (5 minutes)")
		}
		hb.IntervalSec = v
	}
	if v, ok := args["prompt"].(string); ok {
		hb.Prompt = &v
	}
	if v := intArg(args, "ack_max_chars", 0); v > 0 {
		hb.AckMaxChars = v
	}
	if v := intArg(args, "max_retries", -1); v >= 0 {
		hb.MaxRetries = v
	}
	if v, ok := args["isolated_session"].(bool); ok {
		hb.IsolatedSession = v
	}
	if v, ok := args["light_context"].(bool); ok {
		hb.LightContext = v
	}
	if v, ok := args["timezone"].(string); ok {
		hb.Timezone = &v
	}
	if v, ok := args["active_hours"].(string); ok {
		start, end := parseActiveHoursRange(v)
		hb.ActiveHoursStart = &start
		hb.ActiveHoursEnd = &end
	}

	// Auto-fill delivery from context.
	if v, ok := args["channel"].(string); ok {
		hb.Channel = &v
	} else if hb.Channel == nil {
		if ch := ToolChannelFromCtx(ctx); ch != "" {
			hb.Channel = &ch
		}
	}
	if v, ok := args["chat_id"].(string); ok {
		hb.ChatID = &v
	} else if hb.ChatID == nil {
		if cid := ToolChatIDFromCtx(ctx); cid != "" {
			hb.ChatID = &cid
		}
	}

	if hb.Enabled && hb.NextRunAt == nil {
		nextRun := time.Now().Add(time.Duration(hb.IntervalSec)*time.Second + store.StaggerOffset(hb.AgentID, hb.IntervalSec))
		hb.NextRunAt = &nextRun
	}

	if err := t.hbStore.Upsert(ctx, hb); err != nil {
		return ErrorResult(fmt.Sprintf("failed to save heartbeat config: %v", err))
	}

	return NewResult(fmt.Sprintf("Heartbeat config saved. Enabled: %v, interval: %ds, delivery: %s/%s",
		hb.Enabled, hb.IntervalSec, derefStr(hb.Channel, "none"), derefStr(hb.ChatID, "none")))
}

func (t *HeartbeatTool) handleToggle(ctx context.Context, agentID uuid.UUID, enabled bool) *Result {
	hb, err := t.hbStore.Get(ctx, agentID)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return ErrorResult("heartbeat not configured — use set action first")
		}
		return ErrorResult(err.Error())
	}
	hb.Enabled = enabled
	if enabled && hb.NextRunAt == nil {
		nextRun := time.Now().Add(time.Duration(hb.IntervalSec) * time.Second)
		hb.NextRunAt = &nextRun
	}
	if err := t.hbStore.Upsert(ctx, hb); err != nil {
		return ErrorResult(err.Error())
	}
	return NewResult(fmt.Sprintf("Heartbeat %s.", boolLabel(enabled, "enabled", "disabled")))
}

// checkPermission verifies the calling user has permission to modify heartbeat config.
// Flow: deny list → allow list → agent owner/shares fallback.
// Read-only actions (status, get, get_checklist, logs) skip this check.
func (t *HeartbeatTool) checkPermission(ctx context.Context, agentID uuid.UUID) error {
	if t.permStore == nil {
		return nil // no permission store = allow all (backward compat)
	}
	// Use individual sender ID for permission check (not group-scoped UserID).
	// This matches the pattern used by file_writer and /addwriter commands.
	senderID := store.SenderIDFromContext(ctx)
	if senderID == "" {
		return nil // system context (cron, subagent) = allow
	}
	numericID := strings.SplitN(senderID, "|", 2)[0]

	// Determine scope from context: "agent" for DM, "group:{channel}:{chatId}" for groups.
	scope := "agent"
	if ch := ToolChannelFromCtx(ctx); ch != "" {
		if cid := ToolChatIDFromCtx(ctx); cid != "" {
			scope = "group:" + ch + ":" + cid
		}
	}

	allowed, err := t.permStore.CheckPermission(ctx, agentID, scope, store.ConfigTypeHeartbeat, numericID)
	if err != nil {
		return fmt.Errorf("permission check failed: %w", err)
	}
	if allowed {
		return nil
	}

	// Fallback: check if user is agent owner (via agent store).
	if t.agentStore != nil {
		ag, agErr := t.agentStore.GetByID(ctx, agentID)
		if agErr == nil && ag.OwnerID != "" && ag.OwnerID == senderID {
			return nil // agent owner = allow
		}
	}

	return fmt.Errorf("permission denied: you are not authorized to modify heartbeat config for this agent")
}

func (t *HeartbeatTool) handleSetChecklist(ctx context.Context, agentID uuid.UUID, content string) *Result {
	if t.agentStore == nil {
		return ErrorResult("agent store not configured")
	}
	if content == "" {
		return ErrorResult("content parameter is required — provide the HEARTBEAT.md checklist")
	}
	if err := t.agentStore.SetAgentContextFile(ctx, agentID, "HEARTBEAT.md", content); err != nil {
		return ErrorResult(fmt.Sprintf("failed to save HEARTBEAT.md: %v", err))
	}
	return NewResult(fmt.Sprintf("HEARTBEAT.md saved (%d chars). The heartbeat ticker will use this checklist on each run.", len([]rune(content))))
}

func (t *HeartbeatTool) handleGetChecklist(ctx context.Context, agentID uuid.UUID) *Result {
	if t.agentStore == nil {
		return ErrorResult("agent store not configured")
	}
	files, err := t.agentStore.GetAgentContextFiles(ctx, agentID)
	if err != nil {
		return ErrorResult(err.Error())
	}
	for _, f := range files {
		if f.FileName == "HEARTBEAT.md" && f.Content != "" {
			return NewResult(f.Content)
		}
	}
	return NewResult("HEARTBEAT.md not found. Use set_checklist to create one.")
}

func (t *HeartbeatTool) handleTest(agentID uuid.UUID) *Result {
	if t.wakeFn == nil {
		return ErrorResult("heartbeat ticker not available")
	}
	t.wakeFn(agentID)
	return NewResult("Heartbeat test triggered (running in background).")
}

func (t *HeartbeatTool) handleLogs(ctx context.Context, agentID uuid.UUID, limit, offset int) *Result {
	logs, total, err := t.hbStore.ListLogs(ctx, agentID, limit, offset)
	if err != nil {
		return ErrorResult(err.Error())
	}
	if len(logs) == 0 {
		return NewResult("No heartbeat run logs found.")
	}
	data, _ := json.MarshalIndent(map[string]any{
		"total": total,
		"logs":  logs,
	}, "", "  ")
	return NewResult(string(data))
}

// parseActiveHoursRange parses "HH:MM-HH:MM" into start and end strings.
func parseActiveHoursRange(s string) (string, string) {
	for _, sep := range []string{"-", " - ", "~"} {
		parts := splitOnce(s, sep)
		if len(parts) == 2 {
			return parts[0], parts[1]
		}
	}
	return s, ""
}

func splitOnce(s, sep string) []string {
	idx := -1
	for i := range len(s) - len(sep) + 1 {
		if s[i:i+len(sep)] == sep {
			idx = i
			break
		}
	}
	if idx < 0 {
		return []string{s}
	}
	return []string{s[:idx], s[idx+len(sep):]}
}

func intArg(args map[string]any, key string, fallback int) int {
	if v, ok := args[key].(float64); ok {
		return int(v)
	}
	return fallback
}

func boolLabel(v bool, trueStr, falseStr string) string {
	if v {
		return trueStr
	}
	return falseStr
}

func derefStr(s *string, fallback string) string {
	if s == nil || *s == "" {
		return fallback
	}
	return *s
}

func fmtTimePtr(t *time.Time) string {
	if t == nil {
		return "never"
	}
	return t.Format(time.RFC3339)
}