package internal

import (
	"context"
	"sync"
	"sync/atomic"

	"go.uber.org/zap"
)

type contextKey string

var (
	contextKeyApp contextKey = "application-context"
)

func WithAppContext(ctx context.Context, appCtx *AppContext) context.Context {
	appCtx.parent = ctx
	return context.WithValue(ctx, contextKeyApp, appCtx)
}

func GetAppContext(ctx context.Context) *AppContext {
	return ctx.Value(contextKeyApp).(*AppContext)
}

type AppContext struct {
	parent context.Context

	Logger *zap.Logger
	Data   map[string]any
	Cancel context.CancelFunc

	task     []task
	taskCnt  atomic.Int32
	taskDone atomic.Int32
	wg       sync.WaitGroup
}

func (c AppContext) Init() AppContext {
	c.Data = make(map[string]any)
	return c
}

func (c *AppContext) GetContext() context.Context {
	ctx, cancel := context.WithCancel(c.parent)
	c.task = append(c.task, task{
		ctx:    ctx,
		cancel: cancel,
	})
	return ctx
}

// DoTask spawn task on a goroutine
func (p *AppContext) DoTask(f func(ctx context.Context)) {
	ctx := p.GetContext()
	p.wg.Add(1)
	p.taskCnt.Add(1)
	go func() {
		defer p.wg.Done()
		defer p.taskDone.Add(1)
		f(ctx)
	}()
}

func (p *AppContext) TaskCount() int {
	return int(p.taskCnt.Load())
}

func (p *AppContext) TaskDone() int {
	return int(p.taskDone.Load())
}

func (p *AppContext) TaskRemaining() int {
	return p.TaskCount() - p.TaskDone()
}

func (p *AppContext) WaitTask() {
	p.wg.Wait()
}

type task struct {
	ctx    context.Context
	cancel context.CancelFunc
}