/*
 * 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"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cloudwego/kitex/pkg/gofunc"
)

type Object interface {
	Close(exception error) error
}

type objectItem struct {
	object     Object
	lastActive time.Time
}

func NewObjectPool(idleTimeout time.Duration) *ObjectPool {
	s := new(ObjectPool)
	s.idleTimeout = idleTimeout
	s.objects = make(map[string]*Stack[objectItem])
	return s
}

type ObjectPool struct {
	L           sync.Mutex // STW
	objects     map[string]*Stack[objectItem]
	idleTimeout time.Duration
	closed      int32
	once        sync.Once // control the background cleaning goroutine, lazy init until ObjectPool.Push is invoked
}

func (s *ObjectPool) Push(key string, o Object) {
	s.once.Do(func() {
		gofunc.RecoverGoFuncWithInfo(context.Background(), s.cleaning, gofunc.NewBasicInfo("", ""))
	})
	s.L.Lock()
	stk := s.objects[key]
	if stk == nil {
		stk = NewStack[objectItem]()
		s.objects[key] = stk
	}
	s.L.Unlock()
	stk.Push(objectItem{object: o, lastActive: time.Now()})
}

func (s *ObjectPool) Pop(key string) Object {
	s.L.Lock()
	stk := s.objects[key]
	s.L.Unlock()
	if stk == nil {
		return nil
	}
	o, ok := stk.Pop()
	if !ok {
		return nil
	}
	return o.object
}

func (s *ObjectPool) Close() {
	atomic.CompareAndSwapInt32(&s.closed, 0, 1)
}

func (s *ObjectPool) cleaning() {
	cleanInternal := s.idleTimeout
	for atomic.LoadInt32(&s.closed) == 0 {
		time.Sleep(cleanInternal)

		now := time.Now()
		s.L.Lock()
		// clean objects
		for _, stk := range s.objects {
			deleted := 0
			stk.RangeDelete(func(o objectItem) (deleteNode, continueRange bool) {
				if o.object == nil {
					deleted++
					return true, true
				}

				// RangeDelete start from the stack bottom
				// we assume that the values on the top of last valid value are all valid
				if now.Sub(o.lastActive) < s.idleTimeout {
					return false, false
				}
				deleted++
				_ = o.object.Close(nil)
				return true, true
			})
		}
		s.L.Unlock()
	}
}
