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)
	Rich        func(ctx context.Context, args json.RawMessage) Result
}

type Content struct {
	Type     string `json:"type"`
	Text     string `json:"text"`
	Data     string `json:"data"`
	MimeType string `json:"mimeType"`
}

func (c Content) MarshalJSON() ([]byte, error) {
	if c.Type == "image" {
		return json.Marshal(map[string]string{"type": c.Type, "data": c.Data, "mimeType": c.MimeType})
	}
	return json.Marshal(map[string]string{"type": c.Type, "text": c.Text})
}

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

func TextResult(text string, isError bool) Result {
	return Result{Content: []Content{{Type: "text", Text: text}}, IsError: isError}
}

func (t Tool) Call(ctx context.Context, args json.RawMessage) Result {
	if t.Rich != nil {
		return t.Rich(ctx, args)
	}
	text, err := t.Handler(ctx, args)
	if err != nil {
		return TextResult(err.Error(), true)
	}
	return TextResult(text, false)
}

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

type Stdio struct {
	Name    string
	Command string
	Args    []string
}

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"`
}

func (s *Server) Tool(name string) (Tool, bool) {
	for _, t := range s.Tools {
		if t.Name == name {
			return t, true
		}
	}
	return Tool{}, false
}

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()}
	}

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

	return t.Call(ctx, p.Arguments), nil
}