package mcp

import (
	"context"
	"encoding/json"
	"fmt"
)

const protocolVersion = "2025-06-18"

type Tool struct {
	Name        string
	Description string
	InputSchema json.RawMessage
	Handler     func(ctx context.Context, args json.RawMessage) (string, error)
}

type Server struct {
	Name    string
	Version string
	Tools   []Tool
}

type request struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}

type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

type response struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"`
	Result  any             `json:"result,omitempty"`
	Error   *rpcError       `json:"error,omitempty"`
}

type toolDesc struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	InputSchema json.RawMessage `json:"inputSchema"`
}

type content struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

type callResult struct {
	Content []content `json:"content"`
	IsError bool      `json:"isError,omitempty"`
}

func (s *Server) Handle(ctx context.Context, raw json.RawMessage) (json.RawMessage, error) {
	req := request{}
	err := json.Unmarshal(raw, &req)
	if err != nil {
		return nil, fmt.Errorf("parse jsonrpc: %w", err)
	}

	if req.ID == nil {
		return json.Marshal(response{JSONRPC: "2.0", Result: map[string]any{}})
	}

	result, rerr := s.dispatch(ctx, req)
	resp := response{JSONRPC: "2.0", ID: req.ID, Result: result, Error: rerr}
	return json.Marshal(resp)
}

func (s *Server) dispatch(ctx context.Context, req request) (any, *rpcError) {
	switch req.Method {
	case "initialize":
		return map[string]any{
			"protocolVersion": protocolVersion,
			"capabilities":    map[string]any{"tools": map[string]any{}},
			"serverInfo":      map[string]any{"name": s.Name, "version": s.Version},
		}, nil

	case "ping":
		return map[string]any{}, nil

	case "tools/list":
		tools := []toolDesc{}
		for _, t := range s.Tools {
			tools = append(tools, toolDesc{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema})
		}
		return map[string]any{"tools": tools}, nil

	case "tools/call":
		return s.call(ctx, req.Params)
	}

	return nil, &rpcError{Code: -32601, Message: fmt.Sprintf("method not found: %s", req.Method)}
}

func (s *Server) call(ctx context.Context, params json.RawMessage) (any, *rpcError) {
	p := struct {
		Name      string          `json:"name"`
		Arguments json.RawMessage `json:"arguments"`
	}{}
	err := json.Unmarshal(params, &p)
	if err != nil {
		return nil, &rpcError{Code: -32602, Message: err.Error()}
	}

	for _, t := range s.Tools {
		if t.Name != p.Name {
			continue
		}

		text, err := t.Handler(ctx, p.Arguments)
		if err != nil {
			return callResult{Content: []content{{Type: "text", Text: err.Error()}}, IsError: true}, nil
		}

		return callResult{Content: []content{{Type: "text", Text: text}}}, nil
	}

	return nil, &rpcError{Code: -32602, Message: fmt.Sprintf("unknown tool: %s", p.Name)}
}