// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package context implements a mechanism to carry data across API boundaries.
// The context.T struct carries deadlines and cancellation as well as other
// arbitrary values.
//
// Application code receives contexts in two main ways:
//
// 1) A context.T is returned from v23.Init().  This will generally be
// used to set up servers in main, or for stand-alone client programs.
//    func main() {
//      ctx, shutdown := v23.Init()
//      defer shutdown()
//
//      doSomething(ctx)
//    }
//
// 2) A context.T is passed to every server method implementation as the first
// parameter.
//    func (m *myServer) Method(ctx *context.T, call rpc.ServerCall) error {
//      doSomething(ctx)
//    }
//
// Once you have a context you can derive further contexts to change settings.
// for example to adjust a deadline you might do:
//    func main() {
//      ctx, shutdown := v23.Init()
//      defer shutdown()
//      // We'll use cacheCtx to lookup data in memcache
//      // if it takes more than a second to get data from
//      // memcache we should just skip the cache and perform
//      // the slow operation.
//      cacheCtx, cancel := WithTimeout(ctx, time.Second)
//      if err := FetchDataFromMemcache(cacheCtx, key); err != nil {
//        // Here we use the original ctx, not the derived cacheCtx
//        // so we aren't constrained by the 1 second timeout.
//        RecomputeData(ctx, key)
//      }
//    }
//
// Contexts form a tree where derived contexts are children of the
// contexts from which they were derived.  Children inherit all the
// properties of their parent except for the property being replaced
// (the deadline in the example above).
//
// Contexts are extensible.  The Value/WithValue functions allow you to attach
// new information to the context and extend its capabilities.
// In the same way we derive new contexts via the 'With' family of functions
// you can create methods to attach new data:
//
//    package auth
//
//    import "v.io/v23/context"
//
//    type Auth struct{...}
//
//    type key int
//    const authKey = key(0)
//
//    function WithAuth(parent *context.T, data *Auth) *context.T {
//        return context.WithValue(parent, authKey, data)
//    }
//
//    function GetAuth(ctx *context.T) *Auth {
//        data, _ := ctx.Value(authKey).(*Auth)
//        return data
//    }
//
// Note that a value of any type can be used as a key, but you should
// use an unexported value of an unexported type to ensure that no
// collisions can occur.
package context

import (
	"errors"
	"sync"
	"time"

	"v.io/v23/logging"
)

type internalKey int

const (
	rootKey = internalKey(iota)
	cancelKey
	deadlineKey
)

// ContextLogger is a logger that uses a passed in T to configure
// the logging behavior.
type ContextLogger interface {
	// InfoDepth logs to the INFO log. depth is used to determine which call frame to log.
	InfoDepth(ctx *T, depth int, args ...interface{})

	// InfoStack logs the current goroutine's stack if the all parameter
	// is false, or the stacks of all goroutines if it's true.
	InfoStack(ctx *T, all bool)

	// VDepth returns true if the configured logging level is greater than or equal to its parameter. depth
	// is used to determine which call frame to test against.
	VDepth(ctx *T, depth int, level int) bool

	// VIDepth is like VDepth, except that it returns nil if there level is greater than the
	// configured log level.
	VIDepth(ctx *T, depth int, level int) ContextLogger
}

// CancelFunc is used to cancel a context.  The first call will
// cause the paired context and all decendants to close their Done()
// channels.  Further calls do nothing.
type CancelFunc func()

// Cancelled is returned by contexts which have been cancelled.
var Canceled = errors.New("context canceled")

// DeadlineExceeded is returned by contexts that have exceeded their
// deadlines and therefore been canceled automatically.
var DeadlineExceeded = errors.New("context deadline exceeded")

// T carries deadlines, cancellation and data across API boundaries.
// It is safe to use a T from multiple goroutines simultaneously.  The
// zero-type of context is uninitialized and will panic if used
// directly by application code. It also implements v23/logging.Logger and
// hence can be used directly for logging (e.g. ctx.Infof(...)).
type T struct {
	parent     *T
	logger     logging.Logger
	ctxLogger  ContextLogger
	key, value interface{}
}

// RootContext creates a new root context with no data attached.
// A RootContext is cancelable (see WithCancel).
// Typically you should not call this function, instead you should derive
// contexts from other contexts, such as the context returned from v23.Init
// or the result of the Context() method on a ServerCall.  This function
// is sometimes useful in tests, where it is undesirable to initialize a
// runtime to test a function that reads from a T.
func RootContext() (*T, CancelFunc) {
	return WithCancel(&T{logger: logging.Discard, key: rootKey})
}

// WithLogger returns a child of the current context that embeds the supplied
// logger.
func WithLogger(parent *T, logger logging.Logger) *T {
	child := *parent
	child.logger = logger
	return &child
}

// WithContextLogger returns a child of the current context that embeds the supplied
// context logger
func WithContextLogger(parent *T, logger ContextLogger) *T {
	child := *parent
	child.ctxLogger = logger
	return &child
}

// LoggerImplementation returns the implementation of the logger associated
// with this context. It should almost never need to be used by application
// code.
func (t *T) LoggerImplementation() interface{} {
	return t.logger
}

// Initialized returns true if this context has been properly initialized
// by a runtime.
func (t *T) Initialized() bool {
	return t != nil && t.key != nil
}

// Value is used to carry data across API boundaries.  This should be
// used only for data that is relevant across multiple API boundaries
// and not just to pass extra parameters to functions and methods.
// Any type that supports equality can be used as a key, but an
// unexported type should be used to prevent collisions.
func (t *T) Value(key interface{}) interface{} {
	for t != nil {
		if key == t.key {
			return t.value
		}
		t = t.parent
	}
	return nil
}

// Deadline returns the time at which this context will be automatically
// canceled.
func (t *T) Deadline() (deadline time.Time, ok bool) {
	if deadline, ok := t.Value(deadlineKey).(*deadlineState); ok {
		return deadline.deadline, true
	}
	return
}

// After the channel returned by Done() is closed, Err() will return
// either Canceled or DeadlineExceeded.
func (t *T) Err() error {
	if cancel, ok := t.Value(cancelKey).(*cancelState); ok {
		cancel.mu.Lock()
		defer cancel.mu.Unlock()
		return cancel.err
	}
	return nil
}

// Done returns a channel which will be closed when this context.T
// is canceled or exceeds its deadline.  Successive calls will
// return the same value.  Implementations may return nil if they can
// never be canceled.
func (t *T) Done() <-chan struct{} {
	if cancel, ok := t.Value(cancelKey).(*cancelState); ok {
		return cancel.done
	}
	return nil
}

// cancelState helps pass cancellation down the context tree.
type cancelState struct {
	done chan struct{}

	mu       sync.Mutex
	err      error                 // GUARDED_BY(mu)
	children map[*cancelState]bool // GUARDED_BY(mu)
}

func (c *cancelState) addChild(child *cancelState) {
	c.mu.Lock()

	if c.err != nil {
		err := c.err
		c.mu.Unlock()
		child.cancel(err)
		return
	}

	if c.children == nil {
		c.children = make(map[*cancelState]bool)
	}
	c.children[child] = true
	c.mu.Unlock()
}

func (c *cancelState) removeChild(child *cancelState) {
	c.mu.Lock()
	delete(c.children, child)
	c.mu.Unlock()
}

func (c *cancelState) cancel(err error) {
	var children map[*cancelState]bool

	c.mu.Lock()
	if c.err == nil {
		c.err = err
		children = c.children
		c.children = nil
		close(c.done)
	}
	c.mu.Unlock()

	for child, _ := range children {
		child.cancel(err)
	}
}

// A deadlineState helps cancel contexts when a deadline expires.
type deadlineState struct {
	deadline time.Time
	timer    *time.Timer
}

// WithValue returns a child of the current context that will return
// the given val when Value(key) is called.
func WithValue(parent *T, key interface{}, val interface{}) *T {
	if !parent.Initialized() {
		panic("Trying to derive a context from an uninitialized context.")
	}
	if key == nil {
		panic("Attempting to store a context value with an untyped nil key.")
	}
	return &T{logger: parent.logger, ctxLogger: parent.ctxLogger, parent: parent, key: key, value: val}
}

func withCancelState(parent *T) (*T, func(error)) {
	cs := &cancelState{done: make(chan struct{})}
	cancelParent, ok := parent.Value(cancelKey).(*cancelState)
	if ok {
		cancelParent.addChild(cs)
	}
	return WithValue(parent, cancelKey, cs), func(err error) {
		if ok {
			cancelParent.removeChild(cs)
		}
		cs.cancel(err)
	}
}

// WithCancel returns a child of the current context along with
// a function that can be used to cancel it.  After cancel() is
// called the channels returned by the Done() methods of the new context
// (and all context further derived from it) will be closed.
func WithCancel(parent *T) (*T, CancelFunc) {
	t, cancel := withCancelState(parent)
	return t, func() { cancel(Canceled) }
}

func withDeadlineState(parent *T, deadline time.Time, timeout time.Duration) (*T, CancelFunc) {
	t, cancel := withCancelState(parent)
	ds := &deadlineState{deadline, time.AfterFunc(timeout, func() { cancel(DeadlineExceeded) })}
	return WithValue(t, deadlineKey, ds), func() {
		ds.timer.Stop()
		cancel(Canceled)
	}
}

// WithRootContext returns a context derived from parent, but that is
// detached from the deadlines and cancellation hierarchy so that this
// context will only ever be canceled when the returned CancelFunc is
// called, or the RootContext from which this context is ultimately
// derived is canceled.
func WithRootCancel(parent *T) (*T, CancelFunc) {
	var root *cancelState
	for ancestor := parent; ancestor != nil; ancestor = ancestor.parent {
		if cs, ok := ancestor.value.(*cancelState); ok {
			root = cs
		}
	}
	cs := &cancelState{done: make(chan struct{})}
	if root != nil {
		root.addChild(cs)
	}
	return WithValue(parent, cancelKey, cs), func() {
		if root != nil {
			root.removeChild(cs)
		}
		cs.cancel(Canceled)
	}
}

// WithDeadline returns a child of the current context along with a
// function that can be used to cancel it at any time (as from
// WithCancel).  When the deadline is reached the context will be
// automatically cancelled.
// Contexts should be cancelled when they are no longer needed
// so that resources associated with their timers may be released.
func WithDeadline(parent *T, deadline time.Time) (*T, CancelFunc) {
	return withDeadlineState(parent, deadline, deadline.Sub(time.Now()))
}

// WithTimeout is similar to WithDeadline except a Duration is given
// that represents a relative point in time from now.
func WithTimeout(parent *T, timeout time.Duration) (*T, CancelFunc) {
	return withDeadlineState(parent, time.Now().Add(timeout), timeout)
}