package rhntest

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
)

// Server is an httptest Robinhood MCP mock that speaks JSON-RPC tools/call.
type Server struct {
	URL   string
	Token string // if non-empty, require Authorization: Bearer <Token>

	mu         sync.Mutex
	srv        *httptest.Server
	results    map[string]json.RawMessage
	rpcErrs    map[string]string
	httpStatus int
	httpBody   string
	lastName   string
	lastArgs   map[string]any
}

// New starts an httptest server and registers t.Cleanup(Close).
func New(t *testing.T) *Server {
	t.Helper()
	s := &Server{
		results: make(map[string]json.RawMessage),
		rpcErrs: make(map[string]string),
	}
	s.srv = httptest.NewServer(http.HandlerFunc(s.serve))
	s.URL = s.srv.URL
	t.Cleanup(s.Close)
	return s
}

// Set registers a JSON-RPC result for a tool name.
func (s *Server) Set(name string, result json.RawMessage) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.results[name] = result
}

// SetHTTPError makes every request return this HTTP status and body.
func (s *Server) SetHTTPError(status int, body string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.httpStatus = status
	s.httpBody = body
}

// SetRPCError makes tools/call for name return a JSON-RPC error.
func (s *Server) SetRPCError(name, message string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.rpcErrs[name] = message
}

// LastName is the most recent tools/call params.name.
func (s *Server) LastName() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.lastName
}

// LastArgs is the most recent tools/call params.arguments.
func (s *Server) LastArgs() map[string]any {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.lastArgs
}

// Close shuts down the httptest server.
func (s *Server) Close() {
	s.srv.Close()
}

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

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

func (s *Server) serve(w http.ResponseWriter, r *http.Request) {
	if s.Token != "" && r.Header.Get("Authorization") != "Bearer "+s.Token {
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	s.mu.Lock()
	httpStatus := s.httpStatus
	httpBody := s.httpBody
	s.mu.Unlock()
	if httpStatus != 0 {
		w.WriteHeader(httpStatus)
		_, _ = w.Write([]byte(httpBody))
		return
	}

	var req struct {
		ID     json.RawMessage `json:"id"`
		Params struct {
			Name      string         `json:"name"`
			Arguments map[string]any `json:"arguments"`
		} `json:"params"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	s.mu.Lock()
	s.lastName = req.Params.Name
	s.lastArgs = req.Params.Arguments
	rpcMsg, hasRPC := s.rpcErrs[req.Params.Name]
	result, hasResult := s.results[req.Params.Name]
	s.mu.Unlock()

	reply := rpcReply{JSONRPC: "2.0", ID: req.ID}
	if hasRPC {
		reply.Error = &rpcError{Message: rpcMsg}
	} else if !hasResult {
		reply.Error = &rpcError{Message: "unknown tool"}
	} else {
		reply.Result = result
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(reply)
}