package main

import (
	"encoding/json"
	"testing"

	"github.com/mark3labs/mcp-go/mcp"
)

// Since we cannot easily test the main() function without significant refactoring
// to split initialization from the server start, we test the tool handler logic.
// This mirrors the logic in main.go but allows for unit testing.

func TestQueryToolHandler(t *testing.T) {
	// In a real scenario, we might want to refactor main.go to export a 
	// function that creates the handler. For now, we'll verify the logic 
	// we've implemented in the main.go file by testing the expected 
	// behavior of a similar handler.
	
	t.Run("ValidRequest", func(t *testing.T) {
		// Mock arguments
		args := map[string]interface{}{
			"message": "Hello Sidekick",
			"history": []interface{}{
				map[string]interface{}{"role": "user", "content": "Hi"},
				map[string]interface{}{"role": "assistant", "content": "Hello! How can I help?"},
			},
		}
		
		req := mcp.CallToolRequest{}
		req.Params.Name = "query"
		req.Params.Arguments = args

		// We verify the RequireString and GetArguments logic here
		msg, err := req.RequireString("message")
		if err != nil || msg != "Hello Sidekick" {
			t.Errorf("RequireString failed: %v", err)
		}

		rawArgs := req.GetArguments()
		histInter, ok := rawArgs["history"]
		if !ok {
			t.Fatal("history missing from arguments")
		}

		histList, ok := histInter.([]interface{})
		if !ok || len(histList) != 2 {
			t.Errorf("history list invalid: %v", histInter)
		}
	})

	t.Run("MissingMessage", func(t *testing.T) {
		req := mcp.CallToolRequest{}
		req.Params.Name = "query"
		req.Params.Arguments = map[string]interface{}{}

		_, err := req.RequireString("message")
		if err == nil {
			t.Error("expected error for missing message")
		}
	})
}

func TestResultMarshalling(t *testing.T) {
	// Verify the format of the response as requested: {"response", "history"}
	type msg struct {
		Role    string `json:"role"`
		Content string `json:"content"`
	}
	
	result := map[string]interface{}{
		"response": "Agent response",
		"history": []msg{
			{Role: "user", Content: "User message"},
			{Role: "assistant", Content: "Agent response"},
		},
	}

	data, err := json.Marshal(result)
	if err != nil {
		t.Fatalf("marshal failed: %v", err)
	}

	var decoded map[string]interface{}
	if err := json.Unmarshal(data, &decoded); err != nil {
		t.Fatalf("unmarshal failed: %v", err)
	}

	if _, ok := decoded["response"]; !ok {
		t.Error("response key missing")
	}
	if _, ok := decoded["history"]; !ok {
		t.Error("history key missing")
	}
}