package gobridge

import (
	"context"
	"sync"
)

// streamErrBox 是挂在 ctx.Value 树上的一块可变存储，key 是具体的 channel 值本身。
// 生命周期完全跟着调用方持有的 ctx/channel 走，不需要任何显式清理：
// 调用方不再引用它们时，整棵对象图（ctx → streamErrBox → errs → error）
// 会被 GC 自然回收——不依赖 ctx.Done()，context.Background() 一样能正常释放。
type streamErrBox struct {
	mu   sync.Mutex
	errs map[any]error
}

func (b *streamErrBox) get(ch any) error {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.errs[ch]
}

func (b *streamErrBox) set(ch any, err error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.errs[ch] = err
}

type streamErrBoxKey struct{}

// WithStreamErrors 返回一个包过的 ctx，以及一个用于查询流式调用错误的函数 streamErr。
// 之后用这个 ctx（或它的子 ctx，即使被别的 context.With* 再包一层也不受影响）发起的
// 流式输出/双向流调用，如果 Python handler 执行过程中抛异常，可以用 streamErr(ch)
// 查到具体错误，不需要再传一次 ctx。
//
// 不调用这个函数是默认行为：流式调用中途出错只会让 channel 静默提前关闭，
// 调用方拿不到任何错误信息。这是因为 Invoke[chan T] 在建立调用阶段就已经成功
// 返回了 (ch, nil)，后续 Python 生成器内部的异常没有天然的返回值可以携带。
//
//	ctx, streamErr := gobridge.WithStreamErrors(ctx)
//	ch, err := gobridge.Invoke[chan int](ctx, pool, "range_gen", 1, 10)
//	for v := range ch {
//	    fmt.Println(v)
//	}
//	if err := streamErr(ch); err != nil {
//	    // channel 是因为 Python 侧异常提前关闭的，而不是正常 yield 完
//	}
func WithStreamErrors(parent context.Context) (ctx context.Context, streamErr func(ch any) error) {
	box := &streamErrBox{errs: make(map[any]error)}
	ctx = context.WithValue(parent, streamErrBoxKey{}, box)
	return ctx, box.get
}

// recordStreamError 供 invokeStreamOut/invokeStreamBoth 在读到 TypeError 时调用。
// 如果调用方没有用 WithStreamErrors 包过 ctx，直接是个空操作。
func recordStreamError(ctx context.Context, ch any, err error) {
	box, ok := ctx.Value(streamErrBoxKey{}).(*streamErrBox)
	if !ok {
		return
	}
	box.set(ch, err)
}