package mcp

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"strings"
	"sync"
)

const (
	sessionHeader   = "Mcp-Session-Id"
	protocolVersion = "2025-03-26"
	contentType     = "Content-Type"
	sseContentType  = "text/event-stream"
	jsonContentType = "application/json"
)

// ── Context-based auth token propagation ──

type contextKey string

const authTokenKey contextKey = "mcp_auth_token"

// WithAuthToken stores a bearer token in the context.
func WithAuthToken(ctx context.Context, token string) context.Context {
	return context.WithValue(ctx, authTokenKey, token)
}

// AuthTokenFromContext retrieves the bearer token set by WithAuthToken.
func AuthTokenFromContext(ctx context.Context) string {
	if v, ok := ctx.Value(authTokenKey).(string); ok {
		return v
	}
	return ""
}

// ToolHandler executes a tool and returns the result as JSON-serializable data.
type ToolHandler func(ctx context.Context, args map[string]interface{}) (interface{}, error)

// PromptHandler returns messages for a given prompt invocation.
type PromptHandler func(args map[string]string) PromptGetResult

// Server is a Streamable HTTP MCP server (spec 2025-03-26).
//
// Supported transports:
//   - POST  → JSON-RPC request, responds with JSON or SSE based on Accept header
//   - GET   → Opens an SSE stream for server-initiated notifications
//   - DELETE → Terminates a session
type Server struct {
	name           string
	version        string
	mu             sync.RWMutex
	tools          []Tool
	handlers       map[string]ToolHandler
	prompts        []Prompt
	promptHandlers map[string]PromptHandler
	sessions       *SessionStore
}

func NewServer(name, version string) *Server {
	return &Server{
		name:           name,
		version:        version,
		handlers:       make(map[string]ToolHandler),
		promptHandlers: make(map[string]PromptHandler),
		sessions:       NewSessionStore(),
	}
}

// RegisterTool adds a tool definition and its handler.
func (s *Server) RegisterTool(tool Tool, handler ToolHandler) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.tools = append(s.tools, tool)
	s.handlers[tool.Name] = handler
}

// UnregisterTools removes tools by name.
func (s *Server) UnregisterTools(names []string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	remove := make(map[string]bool, len(names))
	for _, n := range names {
		remove[n] = true
	}
	filtered := s.tools[:0]
	for _, t := range s.tools {
		if !remove[t.Name] {
			filtered = append(filtered, t)
		} else {
			delete(s.handlers, t.Name)
		}
	}
	s.tools = filtered
}

// RegisterPrompt adds a prompt definition and its handler.
func (s *Server) RegisterPrompt(prompt Prompt, handler PromptHandler) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.prompts = append(s.prompts, prompt)
	s.promptHandlers[prompt.Name] = handler
}

// HasTool checks whether a tool with the given name is registered.
func (s *Server) HasTool(name string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	_, ok := s.handlers[name]
	return ok
}

// ExportTools returns tool definitions and their handlers for in-process use (e.g. CLI agent).
func (s *Server) ExportTools() ([]Tool, []ToolHandler) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	handlers := make([]ToolHandler, len(s.tools))
	for i, t := range s.tools {
		handlers[i] = s.handlers[t.Name]
	}
	toolsCopy := make([]Tool, len(s.tools))
	copy(toolsCopy, s.tools)
	return toolsCopy, handlers
}

// Stop shuts down the session cleaner.
func (s *Server) Stop() {
	s.sessions.Stop()
}

// ──────────────────────────── HTTP Handler ────────────────────────────

// ServeHTTP implements http.Handler — Streamable HTTP MCP endpoint.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case http.MethodPost:
		s.handlePost(w, r)
	case http.MethodGet:
		s.handleGet(w, r)
	case http.MethodDelete:
		s.handleDelete(w, r)
	default:
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
	}
}

// handlePost processes a JSON-RPC request. Responds with JSON or SSE depending
// on the Accept header.
func (s *Server) handlePost(w http.ResponseWriter, r *http.Request) {
	var req JSONRPCRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeJSON(w, JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      nil,
			Error:   &JSONRPCErr{Code: -32700, Message: "parse error"},
		})
		return
	}

	// Validate session for non-initialize requests
	if req.Method != "initialize" {
		sid := r.Header.Get(sessionHeader)
		if sid != "" && !s.sessions.Touch(sid) {
			http.Error(w, "invalid or expired session", http.StatusBadRequest)
			return
		}
	}

	// Propagate bearer token from HTTP header into context for tool handlers
	ctx := r.Context()
	if token := extractBearerToken(r); token != "" {
		ctx = WithAuthToken(ctx, token)
	}

	resp := s.dispatch(ctx, &req)

	// On initialize, create a session and set the header
	if req.Method == "initialize" {
		sid := s.sessions.Create()
		w.Header().Set(sessionHeader, sid)
	}

	if wantsSSE(r) {
		writeSSE(w, resp)
	} else {
		writeJSON(w, resp)
	}
}

// handleGet opens a server-initiated SSE notification stream.
// Currently sends an initial comment to keep the connection alive, then blocks
// until the client disconnects. Real notifications (tools/list_changed, etc.)
// can be added here later.
func (s *Server) handleGet(w http.ResponseWriter, r *http.Request) {
	sid := r.Header.Get(sessionHeader)
	if sid == "" || !s.sessions.Touch(sid) {
		http.Error(w, "session required", http.StatusBadRequest)
		return
	}

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming not supported", http.StatusInternalServerError)
		return
	}

	w.Header().Set(contentType, sseContentType)
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.WriteHeader(http.StatusOK)

	// Keep-alive comment
	fmt.Fprint(w, ": mcp stream opened\n\n")
	flusher.Flush()

	// Block until client disconnects
	<-r.Context().Done()
}

// handleDelete terminates a session.
func (s *Server) handleDelete(w http.ResponseWriter, r *http.Request) {
	sid := r.Header.Get(sessionHeader)
	if sid == "" {
		http.Error(w, "session required", http.StatusBadRequest)
		return
	}
	if !s.sessions.Delete(sid) {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

// ──────────────────────────── Dispatch ────────────────────────────

func (s *Server) dispatch(ctx context.Context, req *JSONRPCRequest) JSONRPCResponse {
	switch req.Method {
	case "initialize":
		return s.handleInitialize(req)
	case "tools/list":
		return s.handleToolsList(req)
	case "tools/call":
		return s.handleToolsCall(ctx, req)
	case "prompts/list":
		return s.handlePromptsList(req)
	case "prompts/get":
		return s.handlePromptsGet(req)
	case "notifications/initialized":
		return JSONRPCResponse{JSONRPC: "2.0", ID: req.ID, Result: map[string]interface{}{}}
	default:
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error:   &JSONRPCErr{Code: -32601, Message: fmt.Sprintf("method not found: %s", req.Method)},
		}
	}
}

func (s *Server) handleInitialize(req *JSONRPCRequest) JSONRPCResponse {
	caps := Capabilities{Tools: &ToolsCap{ListChanged: true}}
	s.mu.RLock()
	hasPrompts := len(s.prompts) > 0
	s.mu.RUnlock()
	if hasPrompts {
		caps.Prompts = &PromptsCap{}
	}
	return JSONRPCResponse{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result: ServerInfo{
			ProtocolVersion: protocolVersion,
			Capabilities:    caps,
			ServerInfo:      ServerMeta{Name: s.name, Version: s.version},
		},
	}
}

func (s *Server) handleToolsList(req *JSONRPCRequest) JSONRPCResponse {
	s.mu.RLock()
	toolsCopy := make([]Tool, len(s.tools))
	copy(toolsCopy, s.tools)
	s.mu.RUnlock()

	return JSONRPCResponse{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result:  ToolsListResult{Tools: toolsCopy},
	}
}

func (s *Server) handleToolsCall(ctx context.Context, req *JSONRPCRequest) JSONRPCResponse {
	var params ToolCallParams
	if err := json.Unmarshal(req.Params, &params); err != nil {
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error:   &JSONRPCErr{Code: -32602, Message: "invalid params"},
		}
	}

	s.mu.RLock()
	handler, ok := s.handlers[params.Name]
	s.mu.RUnlock()

	if !ok {
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error:   &JSONRPCErr{Code: -32602, Message: fmt.Sprintf("unknown tool: %s", params.Name)},
		}
	}

	// Inject bearer token from HTTP Authorization header into tool args
	// so apiKeyFromArgs picks it up instead of falling back to server env vars.
	if token := AuthTokenFromContext(ctx); token != "" {
		if params.Arguments == nil {
			params.Arguments = make(map[string]interface{})
		}
		params.Arguments["api_key"] = token
	}

	result, err := handler(ctx, params.Arguments)
	if err != nil {
		slog.Error("tool execution failed", "tool", params.Name, "error", err)
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: ToolCallResult{
				Content: []Content{{Type: "text", Text: fmt.Sprintf("Error: %s", err.Error())}},
				IsError: true,
			},
		}
	}

	text, err := json.Marshal(result)
	if err != nil {
		text = []byte(fmt.Sprintf("%v", result))
	}

	return JSONRPCResponse{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result: ToolCallResult{
			Content: []Content{{Type: "text", Text: string(text)}},
		},
	}
}

func (s *Server) handlePromptsList(req *JSONRPCRequest) JSONRPCResponse {
	s.mu.RLock()
	promptsCopy := make([]Prompt, len(s.prompts))
	copy(promptsCopy, s.prompts)
	s.mu.RUnlock()

	return JSONRPCResponse{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result:  PromptsListResult{Prompts: promptsCopy},
	}
}

func (s *Server) handlePromptsGet(req *JSONRPCRequest) JSONRPCResponse {
	var params PromptGetParams
	if err := json.Unmarshal(req.Params, &params); err != nil {
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error:   &JSONRPCErr{Code: -32602, Message: "invalid params"},
		}
	}

	s.mu.RLock()
	handler, ok := s.promptHandlers[params.Name]
	s.mu.RUnlock()

	if !ok {
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error:   &JSONRPCErr{Code: -32602, Message: fmt.Sprintf("unknown prompt: %s", params.Name)},
		}
	}

	result := handler(params.Arguments)
	return JSONRPCResponse{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result:  result,
	}
}

// ──────────────────────────── Helpers ────────────────────────────

func wantsSSE(r *http.Request) bool {
	accept := r.Header.Get("Accept")
	return strings.Contains(accept, sseContentType)
}

func writeJSON(w http.ResponseWriter, v interface{}) {
	w.Header().Set(contentType, jsonContentType)
	if err := json.NewEncoder(w).Encode(v); err != nil {
		slog.Error("failed to write response", "error", err)
	}
}

func extractBearerToken(r *http.Request) string {
	auth := r.Header.Get("Authorization")
	if len(auth) > 7 && strings.EqualFold(auth[:7], "Bearer ") {
		return auth[7:]
	}
	return ""
}

func writeSSE(w http.ResponseWriter, v interface{}) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		writeJSON(w, v)
		return
	}

	w.Header().Set(contentType, sseContentType)
	w.Header().Set("Cache-Control", "no-cache")

	data, err := json.Marshal(v)
	if err != nil {
		slog.Error("failed to marshal SSE data", "error", err)
		return
	}

	fmt.Fprintf(w, "data: %s\n\n", data)
	flusher.Flush()
}
