package agent

import (
	"context"
	"encoding/json"
	"strings"

	"github.com/keepmind9/acp-sdk-go/rpc"
	"github.com/keepmind9/acp-sdk-go/schema"
)

// buildAgentRouter creates a rpc.Handler that dispatches incoming JSON-RPC
// messages to the appropriate Agent interface method.
func buildAgentRouter(agent Agent) rpc.Handler {
	return func(method string, params json.RawMessage, isRequest bool) (interface{}, error) {
		ctx := context.Background()

		// Extension methods start with "_"
		if strings.HasPrefix(method, "_") {
			return handleAgentExtension(agent, ctx, method, params, isRequest)
		}

		switch method {
		case schema.MethodInitialize:
			if !isRequest {
				return nil, nil
			}
			return handleRequest(params, func() (interface{}, error) {
				var req schema.InitializeRequest
				if err := json.Unmarshal(params, &req); err != nil {
					return nil, rpc.NewRPCError(-32602, "Invalid params: "+err.Error())
				}
				return agent.Initialize(ctx, &req)
			})

		case schema.MethodAuthenticate:
			if !isRequest {
				return nil, nil
			}
			return handleRequest(params, func() (interface{}, error) {
				var req schema.AuthenticateRequest
				if err := json.Unmarshal(params, &req); err != nil {
					return nil, rpc.NewRPCError(-32602, "Invalid params: "+err.Error())
				}
				return agent.Authenticate(ctx, &req)
			})

		case schema.MethodSessionNew:
			if !isRequest {
				return nil, nil
			}
			return handleRequest(params, func() (interface{}, error) {
				var req schema.NewSessionRequest
				if err := json.Unmarshal(params, &req); err != nil {
					return nil, rpc.NewRPCError(-32602, "Invalid params: "+err.Error())
				}
				return agent.NewSession(ctx, &req)
			})

		case schema.MethodSessionLoad:
			if !isRequest {
				return nil, nil
			}
			return handleRequest(params, func() (interface{}, error) {
				var req schema.LoadSessionRequest
				if err := json.Unmarshal(params, &req); err != nil {
					return nil, rpc.NewRPCError(-32602, "Invalid params: "+err.Error())
				}
				return agent.LoadSession(ctx, &req)
			})

		case schema.MethodSessionList:
			if !isRequest {
				return nil, nil
			}
			return handleRequest(params, func() (interface{}, error) {
				var req schema.ListSessionsRequest
				if err := json.Unmarshal(params, &req); err != nil {
					return nil, rpc.NewRPCError(-32602, "Invalid params: "+err.Error())
				}
				return agent.ListSessions(ctx, &req)
			})

		case schema.MethodSessionPrompt:
			if !isRequest {
				return nil, nil
			}
			return handleRequest(params, func() (interface{}, error) {
				var req schema.PromptRequest
				if err := json.Unmarshal(params, &req); err != nil {
					return nil, rpc.NewRPCError(-32602, "Invalid params: "+err.Error())
				}
				return agent.Prompt(ctx, &req)
			})

		case schema.MethodSessionCancel:
			if isRequest {
				return nil, nil
			}
			var notif schema.CancelNotification
			if err := json.Unmarshal(params, &notif); err != nil {
				return nil, nil // ignore malformed notifications
			}
			agent.Cancel(ctx, &notif)
			return nil, nil

		case schema.MethodSessionSet_mode:
			if !isRequest {
				return nil, nil
			}
			return handleRequest(params, func() (interface{}, error) {
				var req schema.SetSessionModeRequest
				if err := json.Unmarshal(params, &req); err != nil {
					return nil, rpc.NewRPCError(-32602, "Invalid params: "+err.Error())
				}
				return agent.SetSessionMode(ctx, &req)
			})

		case schema.MethodSessionSet_config_option:
			if !isRequest {
				return nil, nil
			}
			return handleRequest(params, func() (interface{}, error) {
				var req schema.SetSessionConfigOptionRequest
				if err := json.Unmarshal(params, &req); err != nil {
					return nil, rpc.NewRPCError(-32602, "Invalid params: "+err.Error())
				}
				return agent.SetSessionConfigOption(ctx, &req)
			})

		default:
			if isRequest {
				return nil, rpc.NewRPCError(-32601, "Method not found: "+method)
			}
			return nil, nil
		}
	}
}

// handleRequest is a helper that calls the handler function and returns the result.
func handleRequest(params json.RawMessage, fn func() (interface{}, error)) (interface{}, error) {
	return fn()
}

// handleAgentExtension dispatches extension methods (starting with "_").
func handleAgentExtension(agent Agent, ctx context.Context, method string, params json.RawMessage, isRequest bool) (any, error) {
	var p map[string]any
	if params != nil {
		json.Unmarshal(params, &p)
	}

	if isRequest {
		if h, ok := agent.(ExtMethodHandler); ok {
			return h.ExtMethod(ctx, method, p)
		}
		return nil, rpc.NewRPCError(-32601, "Method not found: "+method)
	}

	// Notification
	if h, ok := agent.(ExtNotificationHandler); ok {
		h.ExtNotification(ctx, method, p)
	}
	return nil, nil
}