package agent

import (
	"context"
	"errors"
	"fmt"
	"sync"
)

var (
	ErrTaskAdmissionClosed = errors.New("Agent task admission is paused")
	ErrTaskStopped         = errors.New("Agent task is stopped")
)

type TaskIdentity struct {
	DispatcherID string
	TenantID     int64
	TaskID       string
}

type activeTaskCall struct {
	cancel context.CancelFunc
	done   chan struct{}
}

type taskCalls struct {
	state  string
	active map[string]*activeTaskCall
}

// TaskCalls keeps an in-process execution barrier for task-level controls.
// Dispatcher owns durable task state; an Agent restart never authorizes a new
// call without a fresh Dispatcher instruction. The caller must authenticate
// the active Dispatcher session before invoking Register or Apply.
type TaskCalls struct {
	mu    sync.Mutex
	tasks map[TaskIdentity]*taskCalls
}

func (r *TaskCalls) Register(task TaskIdentity, callID string, cancel context.CancelFunc) (func(), error) {
	if r == nil || !validTaskIdentity(task) || callID == "" || cancel == nil {
		return nil, errors.New("Agent call requires a task identity, call ID and cancellation")
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	entry := r.entry(task)
	switch entry.state {
	case "stopped":
		return nil, ErrTaskStopped
	case "paused":
		return nil, ErrTaskAdmissionClosed
	}
	if _, exists := entry.active[callID]; exists {
		return nil, errors.New("Agent call is already active")
	}
	call := &activeTaskCall{cancel: cancel, done: make(chan struct{})}
	entry.active[callID] = call
	var once sync.Once
	return func() {
		once.Do(func() {
			r.mu.Lock()
			delete(entry.active, callID)
			close(call.done)
			r.mu.Unlock()
		})
	}, nil
}

// Apply closes admission before requesting hangup or waiting for drain. Its
// success means all calls active when the control arrived have finished; a
// timed-out wait leaves the pause/stop barrier in place for re-delivery.
func (r *TaskCalls) Apply(ctx context.Context, task TaskIdentity, action, policy string) error {
	if r == nil || ctx == nil || !validTaskIdentity(task) {
		return errors.New("Agent task control requires a context and task identity")
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	if (action != "pause" && action != "stop" && action != "resume") ||
		(action == "resume" && policy != "") ||
		(action != "resume" && policy != "hangup" && policy != "drain") {
		return errors.New("Agent task control action or active-call policy is invalid")
	}
	r.mu.Lock()
	entry := r.entry(task)
	if entry.state == "stopped" && action != "stop" {
		r.mu.Unlock()
		return ErrTaskStopped
	}
	if action == "resume" {
		entry.state = ""
		r.mu.Unlock()
		return nil
	}
	entry.state = "paused"
	if action == "stop" {
		entry.state = "stopped"
	}
	active := make([]*activeTaskCall, 0, len(entry.active))
	for _, call := range entry.active {
		active = append(active, call)
	}
	r.mu.Unlock()
	if policy == "hangup" {
		for _, call := range active {
			call.cancel()
		}
	}
	for _, call := range active {
		select {
		case <-call.done:
		case <-ctx.Done():
			return fmt.Errorf("Agent task control waiting for active calls: %w", ctx.Err())
		}
	}
	return nil
}

func (r *TaskCalls) entry(task TaskIdentity) *taskCalls {
	if r.tasks == nil {
		r.tasks = make(map[TaskIdentity]*taskCalls)
	}
	entry := r.tasks[task]
	if entry == nil {
		entry = &taskCalls{active: make(map[string]*activeTaskCall)}
		r.tasks[task] = entry
	}
	return entry
}

func validTaskIdentity(task TaskIdentity) bool {
	return task.DispatcherID != "" && task.TenantID > 0 && task.TaskID != ""
}