package main

import (
  "context"
  "encoding/json"
  "fmt"
  "log"
  "os"

  "sidekick/pkg/sidekick"

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

func main() {
  sidekick.SetupLogging(os.Stderr)

  cfg, err := sidekick.LoadConfig("config.toml")
  if err != nil {
    log.Fatalf("Failed to load config: %v", err)
  }

  // Disable standard logging to stdout if stdio transport is used,
  // so it doesn't corrupt the JSON-RPC stream.
  if cfg.MCPListener.Transport == "stdio" || cfg.MCPListener.Transport == "" {
    sidekick.SetupLogging(os.Stderr)
  }

  sidekick.GlobalLogger = log.Printf

  tmpl, _ := sidekick.LoadPrompts("prompts.toml")

  sidekick.InitMCPServers(cfg.MCPServers)
  extTools := sidekick.LoadExternalTools(context.Background())
  pool := make(map[string]*sidekick.Sidekick)

  for id, aCfg := range cfg.Agents {
    if tmpl != nil && tmpl.Lookup(id) != nil {
      rendered, rErr := sidekick.RenderPrompt(tmpl, id)
      if rErr == nil {
        aCfg.SystemPrompt = rendered
      }
    }
    mCfg := cfg.Models[aCfg.ModelID]
    agent := sidekick.NewSidekick(aCfg, mCfg, extTools)
    if aCfg.LogIntermediate {
      agentID := id
      agent.IntermediateHandler = func(msg string) {
        log.Printf("[Agent %s] %s\n", agentID, msg)
      }
    }
    pool[id] = agent
  }

  entryAgent := "coordinator"
  if _, ok := pool[entryAgent]; !ok {
    for id := range pool {
      entryAgent = id
      break
    }
  }
  agent := pool[entryAgent]

  mcpServer := server.NewMCPServer("Sidekick MCP", "1.0.0")

  mcpServer.AddTool(
    mcp.NewToolWithRawSchema(
      "query",
      "Ask Sidekick a question or provide a message, optionally with conversation history",
      json.RawMessage(`{
        "type": "object",
        "properties": {
          "message": { "type": "string", "description": "The user's message" },
          "history": {
            "type": "array",
            "description": "Optional list of previous messages",
            "items": {
              "type": "object",
              "properties": {
                "role": { "type": "string" },
                "content": { "type": "string" }
              },
              "required": ["role", "content"]
            }
          }
        },
        "required": ["message"]
      }`),
    ),
    func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
      msg, err := request.RequireString("message")
      if err != nil {
        log.Printf("Received invalid 'query' request: missing message argument")
        return nil, fmt.Errorf("message argument is required and must be a string: %v", err)
      }

      log.Printf("Received 'query' request (message length: %d)", len(msg))

      var inCtx []sidekick.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 {
              roleInter, okRole := hMap["role"]
              contentInter, okContent := hMap["content"]
              if okRole && okContent {
                if roleStr, isStr := roleInter.(string); isStr {
                  if contentStr, isStrContent := contentInter.(string); isStrContent {
                    inCtx = append(inCtx, sidekick.Message{
                      Role:    sidekick.MessageRole(roleStr),
                      Content: contentStr,
                    })
                  }
                }
              }
            }
          }
        }
      }

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

      responseStr, err := agent.Run(ctx, inCtx, pool)
      if err != nil {
        log.Printf("Agent run failed: %v", err)
        return nil, fmt.Errorf("agent run failed: %w", err)
      }

      log.Printf("Sending response (length: %d)", len(responseStr))

      // Add the final response to the history structure to be returned
      outHistory := append(inCtx, sidekick.Message{
        Role:    sidekick.MessageRole("assistant"),
        Content: responseStr,
      })

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

      resultBytes, err := json.Marshal(resultObj)
      if err != nil {
        return nil, fmt.Errorf("failed to marshal result: %w", err)
      }

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

  transport := cfg.MCPListener.Transport
  if transport == "stdio" || transport == "" {
    if err := server.ServeStdio(mcpServer); err != nil {
      log.Fatalf("MCP Server (stdio) error: %v", err)
    }
  } else if transport == "http" || transport == "sse" { // mcp-go supports SSE, which is Streamable HTTP
    port := cfg.MCPListener.Port
    if port == 0 {
      port = 8080
    }
    addr := fmt.Sprintf(":%d", port)
    log.Printf("Starting MCP Streamable HTTP server on %s", addr)
    srv := server.NewStreamableHTTPServer(mcpServer)
    if err := srv.Start(addr); err != nil {
      log.Fatalf("MCP Server (http) error: %v", err)
    }
  } else {
    log.Fatalf("Unsupported MCP listener transport: %s", transport)
  }
}