// Copyright 2020-2026 Buf Technologies, Inc.
//
// 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 incremental

import (
	"context"
	"fmt"
	"runtime"
	"slices"
	"sync"
	"sync/atomic"
	"time"
	"weak"

	"golang.org/x/sync/semaphore"

	"github.com/bufbuild/protocompile/experimental/report"
	"github.com/bufbuild/protocompile/internal"
)

// Executor is a caching executor for incremental queries.
//
// See [New], [Run], and [Invalidate].
type Executor struct {
	reportOptions report.Options
	dirty         sync.RWMutex

	// TODO: Evaluate alternatives. sync.Map is pretty bad at having predictable
	// performance, and we may want to add eviction to keep memoization costs
	// in a long-running process (like, say, a language server) down.
	// See https://github.com/dgraph-io/ristretto as a potential alternative.
	tasks sync.Map // [any, *task]

	sema *semaphore.Weighted

	// The [time.Duration] to wait before running the GC when debug mode is on. See docs for
	// [WithDebugEvict].
	evictGCDeadline time.Duration

	counter atomic.Uint64 // Used for generating sequence IDs for Result.Unchanged.
}

// ExecutorOption is an option func for [New].
type ExecutorOption func(*Executor)

// New constructs a new executor with the given maximum parallelism.
func New(options ...ExecutorOption) *Executor {
	exec := &Executor{
		sema: semaphore.NewWeighted(int64(runtime.GOMAXPROCS(0))),
	}

	for _, opt := range options {
		opt(exec)
	}

	return exec
}

// WithParallelism sets the maximum number of queries that can execute in
// parallel. Defaults to GOMAXPROCS if not set explicitly.
func WithParallelism(n int64) ExecutorOption {
	return func(e *Executor) { e.sema = semaphore.NewWeighted(n) }
}

// WithReportOptions sets the report options for reports generated by this
// executor.
func WithReportOptions(options report.Options) ExecutorOption {
	return func(e *Executor) { e.reportOptions = options }
}

// WithDebugEvict takes a [time.Duration] and configures debug mode for evictions in the
// executor.
//
// If set and the compiler is built with the debug tag, when [Executor.EvictWithCleanup]
// is called, all evicted keys will be tracked. Then after eviction, a goroutine will be
// kicked off to sleep for the configured duration, force a GC run, and then print out all
// pointers that should be evicted but have not been GC'd.
func WithDebugEvict(wait time.Duration) ExecutorOption {
	return func(e *Executor) { e.evictGCDeadline = wait }
}

// Keys returns a snapshot of the keys of which queries are present (and
// memoized) in an Executor.
//
// The returned slice is sorted.
func (e *Executor) Keys() (keys []string) {
	e.tasks.Range(func(k, t any) bool {
		task := t.(*task) //nolint:errcheck // All values in this map are tasks.
		result := task.result.Load()
		if result == nil || !closed(result.done) {
			return true
		}
		keys = append(keys, fmt.Sprintf("%#v", k))
		return true
	})

	slices.Sort(keys)
	return
}

var runExecutorKey, timingsKey byte

// Run executes a set of queries on this executor in parallel.
//
// This function only returns an error if ctx expires during execution,
// in which case it returns nil and [context.Cause].
//
// Errors that occur during each query are contained within the returned results.
// Unlike [Resolve], these contain the *transitive* errors for each query!
//
// Implementations of [Query].Execute MUST NOT UNDER ANY CIRCUMSTANCES call
// Run. This will result in potential resource starvation or deadlocking, and
// defeats other correctness verification (such as cycle detection). Instead,
// use [Resolve], which takes a [Task] instead of an [Executor].
//
// Note: this function really wants to be a method of [Executor], but it isn't
// because it's generic.
func Run[T any](ctx context.Context, e *Executor, queries ...Query[T]) ([]Result[T], *report.Report, error) {
	e.dirty.RLock()
	defer e.dirty.RUnlock()

	// Verify we haven't reëntrantly called Run.
	if callers, ok := ctx.Value(&runExecutorKey).(*[]*Executor); ok {
		if slices.Contains(*callers, e) {
			panic("protocompile/incremental: reentrant call to Run")
		}
		*callers = append(*callers, e)
	} else {
		ctx = context.WithValue(ctx, &runExecutorKey, &[]*Executor{e})
	}
	ctx, cancel := context.WithCancelCause(ctx)
	defer cancel(nil)

	generation := e.counter.Add(1)
	root := &Task{
		ctx:             ctx,
		cancel:          cancel,
		exec:            e,
		result:          &result{done: make(chan struct{})},
		runID:           generation,
		onRootGoroutine: true,
	}

	// Need to acquire a hold on the global semaphore to represent the root
	// task we're about to execute.
	if !root.acquire() {
		return nil, nil, context.Cause(ctx)
	}
	defer root.release()

	results, expired := Resolve(root, queries...)
	if expired != nil {
		if _, aborted := expired.(*errAbort); aborted {
			panic(expired)
		}
		return nil, nil, expired
	}

	// Record all diagnostics generates by the queries.
	report := &report.Report{Options: e.reportOptions}
	dedup := make(map[*task]struct{})
	tasks := make([]*task, 0, len(queries))
	for _, query := range queries {
		task, ok := e.getTask(query.Key())
		if !ok {
			continue // Uncompleted query, can happen due to an abort.
		}
		tasks = append(tasks, task)
	}
	for n := len(tasks); n > 0; n = len(tasks) {
		node := tasks[n-1]
		tasks = tasks[:n-1]
		if _, ok := dedup[node]; ok {
			continue
		}
		node.deps.Range(func(depAny any, _ any) bool {
			dep := depAny.(*task) //nolint:errcheck
			tasks = append(tasks, dep)
			return true
		})
		dedup[node] = struct{}{}
		report.Diagnostics = append(report.Diagnostics, node.report.Diagnostics...)
	}
	report.Canonicalize()

	return results, report, nil
}

// WithTimings enhances a Context to pass into [Run] which it will record timing
// information for each query.
//
// Each referenced query's Key() will be included in timings, mapped to how long
// it took to calculate. If that query was a cache hit, the duration will be
// zero.
func WithTimings(ctx context.Context, timings map[any]time.Duration) context.Context {
	return context.WithValue(ctx, &timingsKey, timings)
}

// Evict marks query keys as invalid, requiring those queries, and their
// dependencies, to be recomputed. keys that are not cached are ignored.
//
// This function cannot execute in parallel with calls to [Run], and will take
// an exclusive lock (note that [Run] calls themselves can be run in parallel).
func (e *Executor) Evict(keys ...any) {
	e.EvictWithCleanup(keys, nil)
}

// EvictWithCleanup is like [Executor.Evict], but it executes the given cleanup
// function atomically with the eviction action.
//
// This function can be used to clean up after a query, or modify the result of
// the evicted query by writing to a variable, without risking concurrent calls
// to [Run] seeing inconsistent or stale state across multiple queries.
func (e *Executor) EvictWithCleanup(keys []any, cleanup func()) {
	var tasks []*task
	for _, key := range keys {
		if t, ok := e.getTask(key); ok {
			tasks = append(tasks, t)
		}
	}
	if len(tasks) == 0 && cleanup == nil {
		return
	}

	e.dirty.Lock()
	defer e.dirty.Unlock()

	var evicted []weak.Pointer[task]
	logEvictionDebug := internal.Debug && e.evictGCDeadline > 0
	for n := len(tasks); n > 0; n = len(tasks) {
		next := tasks[n-1]
		tasks = tasks[:n-1]

		for k := range next.callers.Range {
			tasks = append(tasks, k.(*task)) //nolint:errcheck
		}

		// Remove the task from the map. Syncronized by the dirty lock.
		t, _ := e.tasks.LoadAndDelete(next.query.Key())
		if logEvictionDebug {
			evicted = append(evicted, weak.Make(t.(*task))) //nolint:errcheck
		}

		// Remove the task from the callers of its deps.
		for k := range next.deps.Range {
			k.(*task).callers.Delete(next) //nolint:errcheck
		}
	}

	if evicted != nil {
		go func() {
			time.Sleep(e.evictGCDeadline)
			runtime.GC()
			evicted = slices.DeleteFunc(evicted, func(e weak.Pointer[task]) bool {
				return e.Value() == nil
			})
			for _, e := range evicted {
				internal.DebugLog(
					[]any{"exec %p", e},
					"EvictWithCleanup",
					"failed to GC evicted task %p: %v",
					e.Value(),
					e.Value().query.Key(),
				)
			}
		}()
	}

	if cleanup != nil {
		cleanup()
	}
}

// getTask returns a task pointer for the given key and whether it was found.
// The returned task is nil if found is false.
func (e *Executor) getTask(key any) (_ *task, found bool) {
	if t, ok := e.tasks.Load(key); ok {
		return t.(*task), true //nolint:errcheck
	}
	return nil, false
}

// getOrCreateTask returns (and creates if necessary) a task pointer for the given query.
// The returned task is never nil.
func (e *Executor) getOrCreateTask(query *AnyQuery) *task {
	// Avoid allocating a new task object in the common case.
	key := query.Key()
	if t, ok := e.tasks.Load(key); ok {
		return t.(*task) //nolint:errcheck
	}
	t, _ := e.tasks.LoadOrStore(key, &task{
		query:  query,
		report: report.Report{Options: e.reportOptions},
	})
	return t.(*task) //nolint:errcheck
}