package workers

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/diillson/chatcli/llm/client"
	"github.com/diillson/chatcli/llm/manager"
	"go.uber.org/zap"
)

// DispatcherConfig holds configuration for the async dispatcher.
type DispatcherConfig struct {
	MaxWorkers    int           // max concurrent worker goroutines
	ParallelMode  bool          // whether parallel dispatch is enabled
	Provider      string        // LLM provider for worker instances
	Model         string        // LLM model for worker instances
	WorkerTimeout time.Duration // timeout per individual worker
}

// DefaultMaxWorkers is the default number of concurrent worker goroutines.
const DefaultMaxWorkers = 4

// DefaultWorkerTimeout is the default timeout for a single worker.
const DefaultWorkerTimeout = 10 * time.Minute

// Dispatcher orchestrates parallel agent execution.
type Dispatcher struct {
	registry      *Registry
	lockMgr       *FileLockManager
	llmMgr        manager.LLMManager
	config        DispatcherConfig
	policyChecker PolicyChecker
	pipeline      ExecutionPipeline
	logger        *zap.Logger
}

// NewDispatcher creates a Dispatcher with the given dependencies.
func NewDispatcher(
	registry *Registry,
	llmMgr manager.LLMManager,
	config DispatcherConfig,
	logger *zap.Logger,
) *Dispatcher {
	if config.MaxWorkers <= 0 {
		config.MaxWorkers = DefaultMaxWorkers
	}
	if config.WorkerTimeout <= 0 {
		config.WorkerTimeout = DefaultWorkerTimeout
	}
	return &Dispatcher{
		registry: registry,
		lockMgr:  NewFileLockManager(),
		llmMgr:   llmMgr,
		config:   config,
		logger:   logger,
	}
}

// MaxWorkers returns the maximum number of concurrent worker goroutines.
func (d *Dispatcher) MaxWorkers() int {
	return d.config.MaxWorkers
}

// UpdateProviderModel updates the provider and model used by worker instances.
// This ensures that runtime provider/model changes are respected by parallel agents.
func (d *Dispatcher) UpdateProviderModel(provider, model string) {
	d.config.Provider = provider
	d.config.Model = model
}

// SetPolicyChecker sets the policy checker for enforcing security policies
// on tool calls executed by parallel workers.
func (d *Dispatcher) SetPolicyChecker(pc PolicyChecker) {
	d.policyChecker = pc
}

// SetPipeline wires an execution pipeline (quality.Pipeline et al). When
// non-nil, executeAgent invokes pipeline.Run instead of agent.Execute. nil
// restores the direct-call default.
func (d *Dispatcher) SetPipeline(p ExecutionPipeline) {
	d.pipeline = p
}

// Pipeline returns the currently wired pipeline (may be nil).
func (d *Dispatcher) Pipeline() ExecutionPipeline {
	return d.pipeline
}

// Dispatch executes a batch of agent calls, respecting parallelism settings.
// Independent calls run concurrently (up to MaxWorkers), dependent calls run sequentially.
// Returns results in the same order as the input calls.
func (d *Dispatcher) Dispatch(ctx context.Context, calls []AgentCall) []AgentResult {
	return d.DispatchWithProgress(ctx, calls, nil)
}

// DispatchWithProgress executes a batch of agent calls and sends real-time
// AgentEvent notifications to the provided channel. The channel is closed
// when all agents finish. If progressCh is nil, behaves like Dispatch.
func (d *Dispatcher) DispatchWithProgress(ctx context.Context, calls []AgentCall, progressCh chan<- AgentEvent) []AgentResult {
	if len(calls) == 0 {
		if progressCh != nil {
			close(progressCh)
		}
		return nil
	}
	if !d.config.ParallelMode || len(calls) == 1 {
		results := d.dispatchSequentialWithProgress(ctx, calls, progressCh)
		if progressCh != nil {
			close(progressCh)
		}
		return results
	}
	results := d.dispatchParallelWithProgress(ctx, calls, progressCh)
	if progressCh != nil {
		close(progressCh)
	}
	return results
}

// dispatchSequentialWithProgress executes agent calls one by one, emitting progress events.
func (d *Dispatcher) dispatchSequentialWithProgress(ctx context.Context, calls []AgentCall, progressCh chan<- AgentEvent) []AgentResult {
	results := make([]AgentResult, len(calls))
	total := len(calls)

	for i, call := range calls {
		d.logger.Info("Dispatching agent (sequential)",
			zap.String("agent", string(call.Agent)),
			zap.String("task", truncateStr(call.Task, 80)),
			zap.String("callID", call.ID),
		)

		if progressCh != nil {
			progressCh <- AgentEvent{
				Type:   AgentEventStarted,
				CallID: call.ID,
				Agent:  call.Agent,
				Task:   call.Task,
				Index:  i,
				Total:  total,
			}
		}

		result := d.executeAgent(ctx, call)
		results[i] = result

		d.logger.Info("Agent completed (sequential)",
			zap.String("agent", string(call.Agent)),
			zap.String("callID", call.ID),
			zap.Duration("duration", result.Duration),
			zap.Bool("hasError", result.Error != nil),
		)

		if progressCh != nil {
			evtType := AgentEventCompleted
			if result.Error != nil {
				evtType = AgentEventFailed
			}
			progressCh <- AgentEvent{
				Type:     evtType,
				CallID:   result.CallID,
				Agent:    result.Agent,
				Task:     result.Task,
				Duration: result.Duration,
				Error:    result.Error,
				Index:    i,
				Total:    total,
			}
		}
	}

	return results
}

// dispatchParallelWithProgress executes agent calls concurrently, emitting progress events.
func (d *Dispatcher) dispatchParallelWithProgress(ctx context.Context, calls []AgentCall, progressCh chan<- AgentEvent) []AgentResult {
	results := make([]AgentResult, len(calls))
	sem := make(chan struct{}, d.config.MaxWorkers) // semaphore
	var wg sync.WaitGroup
	total := len(calls)

	for i, call := range calls {
		wg.Add(1)
		go func(idx int, ac AgentCall) {
			defer wg.Done()

			// Acquire semaphore slot
			select {
			case sem <- struct{}{}:
				defer func() { <-sem }()
			case <-ctx.Done():
				result := AgentResult{
					CallID: ac.ID,
					Agent:  ac.Agent,
					Task:   ac.Task,
					Error:  ctx.Err(),
				}
				results[idx] = result
				if progressCh != nil {
					progressCh <- AgentEvent{
						Type:   AgentEventFailed,
						CallID: ac.ID,
						Agent:  ac.Agent,
						Task:   ac.Task,
						Error:  ctx.Err(),
						Index:  idx,
						Total:  total,
					}
				}
				return
			}

			d.logger.Info("Dispatching agent (parallel)",
				zap.String("agent", string(ac.Agent)),
				zap.String("task", truncateStr(ac.Task, 80)),
				zap.String("callID", ac.ID),
				zap.Int("workerSlot", idx),
			)

			if progressCh != nil {
				progressCh <- AgentEvent{
					Type:   AgentEventStarted,
					CallID: ac.ID,
					Agent:  ac.Agent,
					Task:   ac.Task,
					Index:  idx,
					Total:  total,
				}
			}

			result := d.executeAgent(ctx, ac)
			results[idx] = result

			d.logger.Info("Agent completed (parallel)",
				zap.String("agent", string(ac.Agent)),
				zap.String("callID", ac.ID),
				zap.Duration("duration", result.Duration),
				zap.Bool("hasError", result.Error != nil),
			)

			if progressCh != nil {
				evtType := AgentEventCompleted
				if result.Error != nil {
					evtType = AgentEventFailed
				}
				progressCh <- AgentEvent{
					Type:     evtType,
					CallID:   result.CallID,
					Agent:    result.Agent,
					Task:     result.Task,
					Duration: result.Duration,
					Error:    result.Error,
					Index:    idx,
					Total:    total,
				}
			}
		}(i, call)
	}

	wg.Wait()
	return results
}

// executeAgent runs a single agent call with timeout and error handling.
func (d *Dispatcher) executeAgent(ctx context.Context, call AgentCall) AgentResult {
	startTime := time.Now()

	// Look up the agent
	agent, ok := d.registry.Get(call.Agent)
	if !ok {
		return AgentResult{
			CallID:   call.ID,
			Agent:    call.Agent,
			Task:     call.Task,
			Error:    fmt.Errorf("unknown agent type: %s", call.Agent),
			Duration: time.Since(startTime),
		}
	}

	// Resolve the effective client+effort for this worker:
	//   - agent.Model() non-empty → route through the shared model
	//     resolver (API cache → catalog → family heuristic → optimistic);
	//     cross-provider swaps are honored when the target provider is
	//     configured, otherwise the worker stays on the dispatcher's
	//     default model and the UserMessage is logged.
	//   - agent.Effort() non-empty → attached to ctx via WithEffortHint,
	//     so the provider's SendPrompt opts into extended thinking /
	//     reasoning_effort for every LLM call in this worker's ReAct loop.
	//
	// The dispatcher's config.Provider/Model are NOT mutated — swap is
	// scoped to this single worker invocation.
	effProvider := d.config.Provider
	effModel := d.config.Model

	var llmClient client.LLMClient
	if hint := strings.TrimSpace(agent.Model()); hint != "" {
		resolution := client.ResolveModelRouting(client.ResolveModelRoutingInput{
			Router:       d.llmMgr,
			UserProvider: d.config.Provider,
			UserModel:    d.config.Model,
			Hint:         hint,
			Logger:       d.logger,
		})
		if resolution.Changed {
			llmClient = resolution.Client
			effProvider = resolution.Provider
			effModel = resolution.Model
			d.logger.Info("worker model hint honored",
				zap.String("agent", string(call.Agent)),
				zap.String("note", resolution.Note),
				zap.String("from_provider", d.config.Provider),
				zap.String("to_provider", effProvider),
				zap.String("from_model", d.config.Model),
				zap.String("to_model", effModel))
		} else if resolution.UserMessage != "" {
			d.logger.Warn("worker model hint fell back",
				zap.String("agent", string(call.Agent)),
				zap.String("note", resolution.Note),
				zap.String("user_message", resolution.UserMessage))
		}
	}
	if llmClient == nil {
		var err error
		llmClient, err = d.llmMgr.GetClient(effProvider, effModel)
		if err != nil {
			return AgentResult{
				CallID:   call.ID,
				Agent:    call.Agent,
				Task:     call.Task,
				Error:    fmt.Errorf("failed to create LLM client for worker: %w", err),
				Duration: time.Since(startTime),
			}
		}
	}

	// Embed agent metadata in context so the policy checker can show
	// meaningful context (agent name, task description) in security prompts.
	agentCtx := context.WithValue(ctx, CtxKeyAgentName, string(call.Agent))
	agentCtx = context.WithValue(agentCtx, CtxKeyAgentTask, call.Task)

	// Attach effort hint to ctx (no-op when agent.Effort() is empty).
	if effort := client.NormalizeEffort(agent.Effort()); effort != client.EffortUnset {
		agentCtx = client.WithEffortHint(agentCtx, effort)
	}

	// Create worker context with timeout
	workerCtx, cancel := context.WithTimeout(agentCtx, d.config.WorkerTimeout)
	defer cancel()

	deps := &WorkerDeps{
		LLMClient:     llmClient,
		LockMgr:       d.lockMgr,
		PolicyChecker: d.policyChecker,
		Logger:        d.logger.With(zap.String("agent", string(call.Agent)), zap.String("callID", call.ID)),
	}

	// Execute the agent. When a quality pipeline is wired, route through
	// it so registered pre/post hooks (Self-Refine, CoVe, Reflexion, …)
	// can wrap the call. When nil, behavior is identical to a direct
	// agent.Execute — the no-pipeline performance contract.
	var (
		result  *AgentResult
		execErr error
	)
	if d.pipeline != nil {
		result, execErr = d.pipeline.Run(workerCtx, agent, call.Task, deps)
	} else {
		result, execErr = agent.Execute(workerCtx, call.Task, deps)
	}
	if result == nil {
		result = &AgentResult{}
	}

	result.CallID = call.ID
	result.Agent = call.Agent
	result.Task = call.Task
	result.Duration = time.Since(startTime)
	if execErr != nil && result.Error == nil {
		result.Error = execErr
	}

	return *result
}

// FormatResults formats a slice of AgentResult into a feedback string
// for injection into the orchestrator's LLM history.
func FormatResults(results []AgentResult) string {
	var b strings.Builder
	b.WriteString("--- Agent Results ---\n\n")

	for i, r := range results {
		fmt.Fprintf(&b, "[%s] (call %s, %s)\n", r.Agent, r.CallID, r.Duration.Round(time.Millisecond))
		fmt.Fprintf(&b, "Task: %s\n", r.Task)
		if r.Error != nil {
			fmt.Fprintf(&b, "Status: FAILED — %v\n", r.Error)
		} else {
			b.WriteString("Status: OK\n")
		}
		if r.Output != "" {
			b.WriteString("Output:\n")
			b.WriteString(r.Output)
			if !strings.HasSuffix(r.Output, "\n") {
				b.WriteString("\n")
			}
		}
		if i < len(results)-1 {
			b.WriteString("\n---\n\n")
		}
	}

	return b.String()
}

// truncateStr truncates a string to maxLen and appends "..." if needed.
func truncateStr(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}
