package sidekick

import (
  "context"
  "encoding/json"
  "strings"
  "testing"

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

func TestMCPServerQueryHandler(t *testing.T) {
  // 1. Setup Agent with mock LLM
  ac := AgentConfig{ID: "test-agent", Role: RoleSpecialist, MaxIterations: 2}
  mCfg := ModelConfig{}
  agent := NewSidekick(ac, mCfg, nil)

  mockLLM := &mockLLMClient{
    responses: []Message{
      {Content: `{"action": "RESPOND", "params": {"response": "I am Sidekick. How can I help?"}}`},
    },
  }
  SetLLMClient(mockLLM)

  pool := map[string]*Sidekick{"test-agent": agent}

  // 2. Define the handler (mirrors logic in cmd/sidekick-mcp/main.go)
  handler := func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
    msg, err := request.RequireString("message")
    if err != nil {
      return nil, err
    }

    var inCtx []Message
    args := request.GetArguments()
    if histInter, ok := args["history"]; ok {
      if histList, ok := histInter.([]interface{}); ok {
        for _, h := range histList {
          if hMap, ok := h.(map[string]interface{}); ok {
            role, _ := hMap["role"].(string)
            content, _ := hMap["content"].(string)
            if role != "" && content != "" {
              inCtx = append(inCtx, Message{
                Role:    MessageRole(role),
                Content: content,
              })
            }
          }
        }
      }
    }

    inCtx = append(inCtx, Message{
      Role:    MessageRoleUser,
      Content: msg,
    })

    responseStr, err := agent.Run(ctx, inCtx, pool)
    if err != nil {
      return nil, err
    }

    outHistory := append(inCtx, Message{
      Role:    MessageRole("assistant"),
      Content: responseStr,
    })

    resultObj := map[string]interface{}{
      "response": responseStr,
      "history":  outHistory,
    }

    resultBytes, _ := json.Marshal(resultObj)

    return &mcp.CallToolResult{
      Content: []mcp.Content{
        mcp.NewTextContent(string(resultBytes)),
      },
    }, nil
  }

  // 3. Test with simple message
  req := mcp.CallToolRequest{}
  req.Params.Name = "query"
  req.Params.Arguments = map[string]interface{}{
    "message": "Hello",
  }

  res, err := handler(context.Background(), req)
  if err != nil {
    t.Fatalf("handler failed: %v", err)
  }

  if len(res.Content) != 1 {
    t.Fatalf("expected 1 content item, got %d", len(res.Content))
  }

  textContent, ok := res.Content[0].(mcp.TextContent)
  if !ok {
    t.Fatalf("expected TextContent")
  }

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

  if !strings.Contains(result["response"].(string), "How can I help?") {
    t.Errorf("unexpected response: %v", result["response"])
  }

  history, ok := result["history"].([]interface{})
  if !ok || len(history) != 2 {
    t.Errorf("expected history of length 2, got %v", len(history))
  }

  // 4. Test with history
  reqWithHist := mcp.CallToolRequest{}
  reqWithHist.Params.Name = "query"
  reqWithHist.Params.Arguments = map[string]interface{}{
    "message": "What did I say?",
    "history": []interface{}{
      map[string]interface{}{"role": "user", "content": "I like cats"},
      map[string]interface{}{"role": "assistant", "content": "Me too"},
    },
  }

  // Reset mock for next run
  mockLLM.calls = 0
  mockLLM.responses = []Message{
    {Content: `{"action": "RESPOND", "params": {"response": "You said you like cats."}}`},
  }

  resHist, err := handler(context.Background(), reqWithHist)
  if err != nil {
    t.Fatalf("handler with history failed: %v", err)
  }

  var resultHist map[string]interface{}
  textContentHist := resHist.Content[0].(mcp.TextContent)
  json.Unmarshal([]byte(textContentHist.Text), &resultHist)

  historyFinal, ok := resultHist["history"].([]interface{})
  if !ok || len(historyFinal) != 4 { // 2 old + 1 new user + 1 assistant response
    t.Errorf("expected history of length 4, got %v", len(historyFinal))
  }

  lastMsg := historyFinal[3].(map[string]interface{})
  if lastMsg["content"] != "You said you like cats." {
    t.Errorf("unexpected last message: %v", lastMsg["content"])
  }
}