// Copyright 2020-2025 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"

	"golang.org/x/sync/semaphore"

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

// 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

	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 }
}

// 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 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
}

// 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()
	for n := len(tasks); n > 0; n = len(tasks) {
		next := tasks[n-1]
		tasks = tasks[:n-1]

		next.parents.Range(func(k, _ any) bool {
			tasks = append(tasks, k.(*task)) //nolint:errcheck
			return true
		})

		// Remove the task from the map. Syncronized by the dirty lock.
		e.tasks.Delete(next.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
}
