/*
 * Copyright 2024 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package container

import (
	"context"
	"fmt"
	"io"
	"sync/atomic"
)

type pipeState = int32

const (
	pipeStateActive   pipeState = 0
	pipeStateClosed   pipeState = 1
	pipeStateCanceled pipeState = 2
)

var (
	ErrPipeEOF      = io.EOF
	ErrPipeCanceled = fmt.Errorf("pipe canceled")
	stateErrors     = map[pipeState]error{
		pipeStateClosed:   ErrPipeEOF,
		pipeStateCanceled: ErrPipeCanceled,
	}
)

type CtxDoneCallback func(ctx context.Context)

// Pipe implement a queue that never block on Write but block on Read if there is nothing to read
type Pipe[Item any] struct {
	queue    *Queue[Item]
	state    pipeState
	trigger  chan struct{}
	callback CtxDoneCallback
}

func NewPipe[Item any](opts ...PipeOption) *Pipe[Item] {
	options := PipeOptions{}
	for _, opt := range opts {
		opt(&options)
	}
	p := new(Pipe[Item])
	p.queue = NewQueue[Item]()
	p.trigger = make(chan struct{}, 1)
	p.state = pipeStateActive
	p.callback = options.callback
	return p
}

// Read will block if there is nothing to read
func (p *Pipe[Item]) Read(ctx context.Context, items []Item) (n int, err error) {
READ:
	// check readable items
	for i := 0; i < len(items); i++ {
		val, ok := p.queue.Get()
		if !ok {
			break
		}
		items[i] = val
		n++
	}
	if n > 0 {
		return n, nil
	}

	// check state
	state := atomic.LoadInt32(&p.state)
	if err = stateErrors[state]; err != nil {
		return 0, err
	}

	// no data to read, waiting writes
	for {
		if ctx.Done() != nil {
			select {
			case <-ctx.Done():
				if p.callback != nil {
					p.callback(ctx)
					// waiting for callback to close this pipe
					// Read is not concurrent-safe; the trigger signal can only be consumed once.
					<-p.trigger
					goto READ
				}
				return 0, ctx.Err()
			case <-p.trigger:
			}
		} else {
			<-p.trigger
		}
		goto READ
	}
}

func (p *Pipe[Item]) Write(ctx context.Context, items ...Item) (err error) {
	for _, item := range items {
		p.queue.Add(item)
	}
	// wake up
	select {
	case p.trigger <- struct{}{}:
	default:
	}
	return nil
}

func (p *Pipe[Item]) Close() {
	if !atomic.CompareAndSwapInt32(&p.state, pipeStateActive, pipeStateClosed) {
		return
	}
	select {
	case p.trigger <- struct{}{}:
	default:
	}
}

func (p *Pipe[Item]) Cancel() {
	if !atomic.CompareAndSwapInt32(&p.state, pipeStateActive, pipeStateCanceled) {
		return
	}
	select {
	case p.trigger <- struct{}{}:
	default:
	}
}

type PipeOptions struct {
	callback CtxDoneCallback
}

type PipeOption func(*PipeOptions)

func WithCtxDoneCallback(callback CtxDoneCallback) PipeOption {
	return func(options *PipeOptions) {
		options.callback = callback
	}
}
