package mcpserver

import (
	"context"
	"iter"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

type cleanupEventStore struct {
	base            mcp.EventStore
	onSessionClosed func(string)
}

func newCleanupEventStore(base mcp.EventStore, onSessionClosed func(string)) mcp.EventStore {
	return &cleanupEventStore{
		base:            base,
		onSessionClosed: onSessionClosed,
	}
}

func (s *cleanupEventStore) Open(ctx context.Context, sessionID, streamID string) error {
	return s.base.Open(ctx, sessionID, streamID)
}

func (s *cleanupEventStore) Append(ctx context.Context, sessionID, streamID string, data []byte) error {
	return s.base.Append(ctx, sessionID, streamID, data)
}

func (s *cleanupEventStore) After(ctx context.Context, sessionID, streamID string, index int) iter.Seq2[[]byte, error] {
	return s.base.After(ctx, sessionID, streamID, index)
}

func (s *cleanupEventStore) SessionClosed(ctx context.Context, sessionID string) error {
	if err := s.base.SessionClosed(ctx, sessionID); err != nil {
		return err
	}
	if s.onSessionClosed != nil {
		s.onSessionClosed(sessionID)
	}
	return nil
}