package requests

import (
	"context"
	"fmt"

	broker "github.com/pharos-ai-factory/pharos-platform/services/beto/adaptors/messaging/amqp"
	names "github.com/pharos-ai-factory/pharos-platform/services/beto/contracts/requests/generated/names"
	a "github.com/pharos-ai-factory/pharos-platform/services/beto/foundation/application"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protoreflect"
)

type queryHandler struct {
	reply  protoreflect.FieldDescriptor
	handle func(context.Context, proto.Message) (proto.Message, error)
}

// RabbitMQRequestRegistry dispatches explicitly registered context queries.
// Register handlers during composition, before starting workers.
type RabbitMQRequestRegistry struct {
	codec    *codec
	handlers map[protoreflect.Name]queryHandler
	commands map[protoreflect.Name]commandHandler
	policy   *broker.RequestPolicy
}

// RequireAuthenticatedRequests opts this owner into configured broker principals.
// Configure once during composition, before starting workers.
func (r *RabbitMQRequestRegistry) RequireAuthenticatedRequests(policy *broker.RequestPolicy) error {
	if !policy.Valid() || r.policy != nil {
		return fmt.Errorf("invalid or duplicate request authentication policy")
	}
	r.policy = policy
	return nil
}

func New(owner string, request, reply proto.Message) (*RabbitMQRequestRegistry, error) {
	codec, err := newCodec(owner, request, reply)
	if err != nil {
		return nil, err
	}
	return &RabbitMQRequestRegistry{codec: codec, handlers: make(map[protoreflect.Name]queryHandler), commands: make(map[protoreflect.Name]commandHandler)}, nil
}

// Query binds generated identifiers, owner mappings and a plain query handler.
// Registrations are validated during composition, before workers start.
func Query[P, W proto.Message, Q, R any](r *RabbitMQRequestRegistry, name names.QueryName, replyName names.ReplyName, parse func(P) Q, execute func(context.Context, Q) (R, error), reply func(Q, R) (W, error)) error {
	if r.codec.queryType == nil || name == nil || replyName == nil || !name.Valid() || !replyName.Valid() || name.Owner() != r.codec.owner || replyName.Owner() != r.codec.owner {
		return fmt.Errorf("query identifiers disagree with owner")
	}
	return registerQuery(r, protoreflect.Name(name.String()), protoreflect.Name(replyName.String()), parse, execute, reply)
}

func registerQuery[P, W proto.Message, Q, R any](r *RabbitMQRequestRegistry, name, replyName protoreflect.Name, parse func(P) Q, execute func(context.Context, Q) (R, error), reply func(Q, R) (W, error)) error {
	if _, exists := r.handlers[name]; exists {
		return fmt.Errorf("duplicate query handler %s", name)
	}
	if parse == nil || execute == nil || reply == nil || replyName == "error" {
		return fmt.Errorf("invalid query registration %s", name)
	}
	input, err := payloadField(r.codec.queryType.Descriptor(), name)
	if err != nil {
		return err
	}
	if _, ok := r.codec.queryType.New().NewField(input).Message().Interface().(P); !ok {
		return fmt.Errorf("query %s input type disagrees with its contract", name)
	}
	output, err := payloadField(r.codec.reply.Descriptor(), replyName)
	if err != nil {
		return err
	}
	if _, ok := r.codec.reply.New().Interface().(W); !ok {
		return fmt.Errorf("query %s reply type disagrees with its contract", name)
	}
	var bus a.QueryBusPort[Q, R] = handlerQueryBus[Q, R]{execute: execute}
	r.handlers[name] = queryHandler{reply: output, handle: func(ctx context.Context, body proto.Message) (proto.Message, error) {
		input, ok := body.(P)
		if !ok {
			return nil, fmt.Errorf("query payload type disagrees with its registration")
		}
		query := parse(input)
		result, err := bus.Ask(ctx, query)
		if err != nil {
			return nil, err
		}
		return reply(query, result)
	}}
	return nil
}

type handlerQueryBus[Q, R any] struct {
	execute func(context.Context, Q) (R, error)
}

func (b handlerQueryBus[Q, R]) Ask(ctx context.Context, query Q) (R, error) {
	return b.execute(ctx, query)
}

// Handle rejects invalid or unregistered requests before invoking application code.
// Provider failures become typed replies; cancellation leaves delivery unsettled.
func (r *RabbitMQRequestRegistry) Handle(ctx context.Context, id string, data []byte) ([]byte, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if r.policy != nil && !r.policy.AcceptsContext(ctx, r.codec.owner, id) {
		return nil, fmt.Errorf("request authentication required")
	}
	if r.codec.commandField != nil {
		message := r.codec.request.New()
		if err := proto.Unmarshal(data, message.Interface()); err != nil {
			return nil, err
		}
		if message.WhichOneof(message.Descriptor().Oneofs().ByName("payload")) == r.codec.commandField {
			return r.handleCommand(ctx, id, data)
		}
	}
	name, body, err := r.codec.decode(id, data)
	if err != nil {
		return nil, err
	}
	handler, exists := r.handlers[name]
	if !exists {
		return nil, fmt.Errorf("unregistered query %s", name)
	}
	reply, err := handler.handle(ctx, body)
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}
	if err != nil || reply == nil || !reply.ProtoReflect().IsValid() || reply.ProtoReflect().Descriptor() != r.codec.reply.Descriptor() {
		return r.codec.unavailable(id)
	}
	message := reply.ProtoReflect()
	field := message.WhichOneof(message.Descriptor().Oneofs().ByName("payload"))
	if field != handler.reply && field != r.codec.errorField {
		return r.codec.unavailable(id)
	}
	payload := message.Get(field).Message()
	if !payload.IsValid() {
		return r.codec.unavailable(id)
	}
	if err := validatePayload(payload, true); err != nil {
		return r.codec.unavailable(id)
	}
	return r.codec.encode(id, field, payload.Interface())
}

// Validate requires every declared query to have a handler before workers start.
func (r *RabbitMQRequestRegistry) Validate() error {
	if r.codec.queryType != nil {
		fields := r.codec.queryType.Descriptor().Oneofs().ByName("payload").Fields()
		for i := 0; i < fields.Len(); i++ {
			if _, exists := r.handlers[fields.Get(i).Name()]; !exists {
				return fmt.Errorf("unregistered query %s", fields.Get(i).Name())
			}
		}
	}
	if r.codec.commandType != nil {
		fields := r.codec.commandType.Descriptor().Oneofs().ByName("payload").Fields()
		for i := 0; i < fields.Len(); i++ {
			if _, exists := r.commands[fields.Get(i).Name()]; !exists {
				return fmt.Errorf("unregistered command %s", fields.Get(i).Name())
			}
		}
	}
	return nil
}