package ctxmgr

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
)

// UpdateMessageMetadata 更新已持久化助手消息的 metadata 与任务状态。
func (m *ContextManager) UpdateMessageMetadata(ctx context.Context, messageID int64, meta map[string]any, taskStatus string) error {
	if messageID <= 0 {
		return errors.New("invalid message_id")
	}
	if meta == nil {
		meta = map[string]any{}
	}
	if taskStatus != "" {
		meta["task_status"] = taskStatus
	}

	if m.pool != nil {
		metaJSON, _ := json.Marshal(meta)
		_, err := m.pool.Exec(ctx, `
			UPDATE chat_messages
			SET metadata = COALESCE(metadata, '{}'::jsonb) || $1::jsonb
			WHERE id = $2
		`, metaJSON, messageID)
		if err != nil {
			return fmt.Errorf("update message metadata failed: %w", err)
		}
		return nil
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	for sid, msgs := range m.messages {
		for i := range msgs {
			if msgs[i].ID == messageID {
				if msgs[i].Metadata == nil {
					msgs[i].Metadata = map[string]any{}
				}
				for k, v := range meta {
					msgs[i].Metadata[k] = v
				}
				m.messages[sid][i] = msgs[i]
				return nil
			}
		}
	}
	return errors.New("message not found")
}

// FinalizeAssistantTaskMessage 替换助手任务占位正文并写入完整 metadata（含 task_status）。
func (m *ContextManager) FinalizeAssistantTaskMessage(ctx context.Context, messageID int64, content string, meta map[string]any, taskStatus string) error {
	if messageID <= 0 {
		return errors.New("invalid message_id")
	}
	if meta == nil {
		meta = map[string]any{}
	}
	if taskStatus != "" {
		meta["task_status"] = taskStatus
	}

	if m.pool != nil {
		metaJSON, err := json.Marshal(meta)
		if err != nil {
			return fmt.Errorf("marshal metadata failed: %w", err)
		}
		tag, err := m.pool.Exec(ctx, `
			UPDATE chat_messages
			SET content = $1,
			    metadata = COALESCE(metadata, '{}'::jsonb) || $2::jsonb
			WHERE id = $3
		`, content, metaJSON, messageID)
		if err != nil {
			return fmt.Errorf("finalize assistant task message failed: %w", err)
		}
		if tag.RowsAffected() == 0 {
			return errors.New("message not found")
		}
		return nil
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	for sid, msgs := range m.messages {
		for i := range msgs {
			if msgs[i].ID == messageID {
				msgs[i].Content = content
				if msgs[i].Metadata == nil {
					msgs[i].Metadata = map[string]any{}
				}
				for k, v := range meta {
					msgs[i].Metadata[k] = v
				}
				m.messages[sid][i] = msgs[i]
				return nil
			}
		}
	}
	return errors.New("message not found")
}