package handlers

import (
	"context"
	"sync"
	"time"
)

type key int

var pool *sync.Pool

const (
	TraceIDKey = key(1)
	OpKey      = key(2)
)

type hcontext struct {
	traceID  []byte
	op       string
	deadline time.Time
	cancel   context.CancelFunc
	ctx      context.Context
	mu       sync.Mutex
}

func NewContext(traceID []byte, op string, deadline time.Time) (context.Context, context.CancelFunc) {
	c := pool.Get().(*hcontext)
	if c == nil {
		c = &hcontext{
			mu: sync.Mutex{},
		}
	}
	c.traceID = traceID
	c.op = op
	c.deadline = deadline
	c.ctx = nil
	c.cancel = nil
	return c, func() {
		if c.cancel != nil {
			c.cancel()
		}
		pool.Put(c)
	}
}

func (c *hcontext) Deadline() (time.Time, bool) {
	return c.deadline, true
}

func (c *hcontext) Done() <-chan struct{} {
	return c.context().Done()
}

func (c *hcontext) context() context.Context {
	c.mu.Lock()
	defer c.mu.Unlock()
	ctx := context.Background()
	ctx = context.WithValue(ctx, TraceIDKey, c.traceID)
	ctx = context.WithValue(ctx, OpKey, c.op)
	ctx, c.cancel = context.WithDeadline(ctx, c.deadline)
	c.ctx = ctx
	return ctx
}

func (c *hcontext) Err() error {
	return c.context().Err()
}

func (c *hcontext) Value(key interface{}) interface{} {
	switch key {
	case TraceIDKey:
		return c.traceID
	case OpKey:
		return c.op
	}
	return nil
}