package main

import (
	"context"
	"encoding/json"
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/creachadair/jrpc2"
	"github.com/creachadair/jrpc2/channel"
	"github.com/creachadair/jrpc2/handler"
)

// EchoRequest is the request structure for the echo method
type EchoRequest struct {
	Message string `json:"message"`
}

// EchoResponse is the response structure for the echo method
type EchoResponse struct {
	Message string `json:"message"`
	Status  string `json:"status"`
}

// Echo is a simple method that echoes back the message
func Echo(ctx context.Context, req *EchoRequest) (*EchoResponse, error) {
	log.Printf("Echo method called with message: %s", req.Message)
	return &EchoResponse{
		Message: req.Message,
		Status:  "success",
	}, nil
}

// ErrorRequest is the request structure for the error method
type ErrorRequest struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// Error is a method that always returns an error
func Error(ctx context.Context, req *ErrorRequest) (interface{}, error) {
	log.Printf("Error method called with code: %d, message: %s", req.Code, req.Message)
	return nil, &jrpc2.Error{
		Code:    jrpc2.Code(req.Code),
		Message: req.Message,
	}
}

// Sum calculates the sum of an array of numbers
func Sum(ctx context.Context, numbers []float64) (float64, error) {
	log.Printf("Sum method called with numbers: %v", numbers)
	var sum float64
	for _, num := range numbers {
		sum += num
	}
	return sum, nil
}

// ToolsCallRequest handles a tools/call request
type ToolsCallRequest struct {
	Name       string          `json:"name"`
	Parameters json.RawMessage `json:"parameters"`
}

type ToolsCallResponse struct {
	Result json.RawMessage `json:"result"`
}

// ToolsCall handles a tools/call request
func ToolsCall(ctx context.Context, req *ToolsCallRequest) (*ToolsCallResponse, error) {
	log.Printf("tools/call method called with name: %s, parameters: %s", req.Name, string(req.Parameters))

	// Echo back the parameters as the result
	return &ToolsCallResponse{
		Result: req.Parameters,
	}, nil
}

func main() {
	// Set up signal handling for graceful shutdown
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigs
		log.Printf("Received signal %v, shutting down...", sig)
		cancel()
	}()

	// Create a JSON-RPC server that communicates over stdin/stdout
	log.Println("Starting mock MCP adapter")

	// Define the methods
	methods := map[string]handler.Func{
		"echo":       handler.New(Echo),
		"error":      handler.New(Error),
		"sum":        handler.New(Sum),
		"tools/call": handler.New(ToolsCall),
	}

	// Create a server with the defined methods
	srv := jrpc2.NewServer(handler.Map(methods), &jrpc2.ServerOptions{
		AllowPush: true,
		Logger:    jrpc2.StdLogger(log.New(os.Stderr, "[jrpc2] ", log.LstdFlags)),
	})

	// Create a channel for communication over stdin/stdout
	ch := channel.Line(os.Stdin, os.Stdout)

	// Start the server and wait for it to exit
	log.Println("Server ready to handle requests")
	if err := srv.Start(ch); err != nil {
		log.Fatalf("Server failed to start: %v", err)
	}

	// Wait for the server to exit or for context cancellation
	select {
	case <-ctx.Done():
		log.Println("Context canceled, shutting down server")
		srv.Stop()
	}

	log.Println("Server exited gracefully")
}