package tools

import (
	"context"
	"fmt"

	"github.com/ClareAI/astra-insight-mcp/internal/gateway"
	"github.com/ClareAI/astra-insight-mcp/internal/mcp"
)

func registerSendTestMessage(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "send_test_message",
		Description: "Send a test message to the agent and get its response (via SSE streaming). Use this to verify the agent's behavior after updating instructions. You can continue a conversation by passing conversation_id, or start a new one by omitting it.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"agent_id":        {Type: "string", Description: "Agent ID to test"},
				"user_id":         {Type: "string", Description: "User ID for the test session (any unique string)"},
				"message":         {Type: "string", Description: "Test message to send to the agent"},
				"conversation_id": {Type: "string", Description: "Existing conversation ID to continue (omit for new conversation)"},
			},
			Required: []string{"agent_id", "message"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		agentID, _ := args["agent_id"].(string)
		message, _ := args["message"].(string)
		if apiKey == "" || agentID == "" || message == "" {
			return nil, fmt.Errorf("api_key, agent_id, and message are required")
		}

		userID, _ := args["user_id"].(string)
		if userID == "" {
			userID = "insight-mcp-tester"
		}

		req := gateway.ChatRequest{
			AgentID:   agentID,
			UserQuery: message,
			UserID:    userID,
			Stream:    true,
		}
		if cid, ok := args["conversation_id"].(string); ok && cid != "" {
			req.ConversationID = cid
		}

		return gw.SendTestMessage(ctx, apiKey, req)
	})
}

func registerRetrieveKnowledge(s *mcp.Server, gw *gateway.Client) {
	tool := mcp.Tool{
		Name:        "retrieve_knowledge",
		Description: "Search the agent's knowledge base with a query and return matching document content as text. Use this to verify what knowledge the agent has access to.",
		InputSchema: mcp.InputSchema{
			Type: "object",
			Properties: map[string]mcp.Property{
				"agent_id": {Type: "string", Description: "Agent ID"},
				"query":    {Type: "string", Description: "Search query to find relevant knowledge"},
			},
			Required: []string{"agent_id", "query"},
		},
	}

	s.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (interface{}, error) {
		apiKey := apiKeyFromArgs(args)
		agentID, _ := args["agent_id"].(string)
		query, _ := args["query"].(string)
		if apiKey == "" || agentID == "" || query == "" {
			return nil, fmt.Errorf("api_key, agent_id, and query are required")
		}

		req := gateway.RetrieveRequest{
			AgentID: agentID,
			Query:   query,
		}

		return gw.RetrieveKnowledge(ctx, apiKey, req)
	})
}
