package comm

import (
	"context"
	"fmt"
	"sync"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
)

// AskHandler 处理发给本 agent 的 Ask。它运行在总线的按订阅工作 goroutine 上；对单个
// agent 而言 ask 是串行的。处理函数应经 (p *Probe).Reply 或 Answer 作答；返回的非 nil
// 错误由探针记日志并丢弃。
type AskHandler func(ctx context.Context, a Ask) error

// EventHandler 处理定向给本 agent 或广播给本 agent 的 Event。由于一个 agent 订阅了两个
// 事件模式（广播与定向），总线可能从两个工作 goroutine 并发调用同一个处理函数，因此
// EventHandler 必须可安全并发使用。
type EventHandler func(ctx context.Context, e Event) error

// SubscribeAsk 登记 h 以接收发给本 agent 的每个 Ask，并返回一个解除登记的函数。ctx 被
// 取消时登记自动移除；返回的函数是幂等的。
func (p *Probe) SubscribeAsk(ctx context.Context, h AskHandler) (func(), error) {
	if h == nil {
		return nil, fmt.Errorf("comm: subscribe ask: handler must not be nil")
	}
	pattern := askReceivePattern(p.agentID)
	unsub, err := p.bus.Subscribe(ctx, pattern, func(bctx context.Context, e bus.Event) error {
		a, ok := e.Payload.(Ask)
		if !ok {
			return fmt.Errorf("comm: ask payload type %T, want Ask", e.Payload)
		}
		if err := h(bctx, a); err != nil {
			p.log.Warn("ask handler failed", "from", a.From, "ask_id", a.ID, "err", err)
		}
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("comm: subscribe ask: %w", err)
	}
	return func() { unsub() }, nil
}

// SubscribeEvent 登记 h 以同时接收广播事件与定向给本 agent 的事件。它返回一个移除两处
// 登记的函数；该函数幂等，多次调用也安全。ctx 被取消时两处登记都自动移除。
//
// 处理函数会从两个订阅工作 goroutine 抵达调用方，也可能与 SubscribeAsk 的处理函数并发
// 运行，因此它必须可安全并发使用。
func (p *Probe) SubscribeEvent(ctx context.Context, h EventHandler) (func(), error) {
	if h == nil {
		return nil, fmt.Errorf("comm: subscribe event: handler must not be nil")
	}

	wrap := func(bctx context.Context, e bus.Event) error {
		ev, ok := e.Payload.(Event)
		if !ok {
			return fmt.Errorf("comm: event payload type %T, want Event", e.Payload)
		}
		if err := h(bctx, ev); err != nil {
			p.log.Warn("event handler failed", "from", ev.From, "name", ev.Name, "err", err)
		}
		return nil
	}

	var (
		mu     sync.Mutex
		unsubs []func()
	)
	combined := func() {
		mu.Lock()
		defer mu.Unlock()
		for _, u := range unsubs {
			u()
		}
	}

	for _, pattern := range eventReceivePatterns(p.agentID) {
		u, err := p.bus.Subscribe(ctx, pattern, wrap)
		if err != nil {
			combined()
			return nil, fmt.Errorf("comm: subscribe event %q: %w", pattern, err)
		}
		mu.Lock()
		unsubs = append(unsubs, u)
		mu.Unlock()
	}
	return combined, nil
}