// Package automation provides a CLI automation layer using claude-go
// for agent-based processing of tasks like audit result transformation.
package automation

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	claude "github.com/standardbeagle/claude-go"
)

// queryFunc is the signature of claude.Query. Process calls it through the
// Processor's queryFn field so tests can inject a fake without spawning a real
// subprocess.
type queryFunc func(ctx context.Context, prompt string, opts *claude.AgentOptions) ([]claude.MessageType, error)

// Processor handles agent-based automation tasks.
type Processor struct {
	mu          sync.RWMutex
	defaultOpts *claude.AgentOptions
	prompts     *PromptRegistry
	config      ProcessorConfig
	stats       ProcessorStats
	// queryFn is the seam to the underlying agent query. Defaults to
	// claude.Query in New; overridden in tests with a fake.
	queryFn queryFunc
	// totalDuration accumulates processing time across all tasks so Stats
	// can derive AverageDuration. Guarded by mu like the rest of stats.
	totalDuration time.Duration
	closed        atomic.Bool
}

// ProcessorConfig configures the automation processor.
type ProcessorConfig struct {
	// Model is the default model to use (haiku recommended for most tasks)
	Model string

	// MaxBudgetUSD is the maximum cost per task in USD
	MaxBudgetUSD float64

	// MaxTurns is the maximum conversation turns
	MaxTurns int

	// WorkingDir is the working directory context
	WorkingDir string

	// AllowedTools specifies tools the agent can use
	AllowedTools []string

	// DisallowedTools specifies tools the agent cannot use
	DisallowedTools []string

	// TimeoutSecs is the timeout for each task
	TimeoutSecs int
}

// ProcessorStats tracks processor statistics.
type ProcessorStats struct {
	TasksProcessed  int64
	TasksSucceeded  int64
	TasksFailed     int64
	TotalTokens     int64
	TotalCostUSD    float64
	AverageDuration time.Duration
}

// DefaultConfig returns config optimized for automation tasks.
func DefaultConfig() ProcessorConfig {
	return ProcessorConfig{
		Model:        "haiku", // Fast and cheap for processing
		MaxBudgetUSD: 0.01,    // $0.01 limit per task
		MaxTurns:     3,       // Usually 1-2 turns needed
		TimeoutSecs:  30,      // 30 second timeout
		DisallowedTools: []string{ // No file/bash access for processing
			"Bash", "Write", "Edit", "Read",
		},
	}
}

// New creates a new automation processor.
func New(cfg ProcessorConfig) (*Processor, error) {
	if cfg.Model == "" {
		cfg.Model = "haiku"
	}
	if cfg.MaxTurns == 0 {
		cfg.MaxTurns = 3
	}
	if cfg.TimeoutSecs == 0 {
		cfg.TimeoutSecs = 30
	}

	opts := &claude.AgentOptions{
		Model:           cfg.Model,
		MaxTurns:        cfg.MaxTurns,
		MaxBudgetUSD:    cfg.MaxBudgetUSD,
		AllowedTools:    cfg.AllowedTools,
		DisallowedTools: cfg.DisallowedTools,
		PermissionMode:  claude.PermissionModeBypassPermission,
	}

	if cfg.WorkingDir != "" {
		opts.WorkingDirectory = cfg.WorkingDir
	}

	return &Processor{
		defaultOpts: opts,
		prompts:     DefaultPromptRegistry(),
		config:      cfg,
		queryFn:     claude.Query,
	}, nil
}

// Process runs a single task and returns the result.
func (p *Processor) Process(ctx context.Context, task Task) (*Result, error) {
	if p.closed.Load() {
		return nil, fmt.Errorf("processor is closed")
	}

	startTime := time.Now()

	// Get the system prompt for this task type
	systemPrompt := p.prompts.Get(task.Type)
	if systemPrompt == "" {
		return nil, fmt.Errorf("unknown task type: %s", task.Type)
	}

	// Build the user prompt from task input
	userPrompt, err := p.buildUserPrompt(task)
	if err != nil {
		return nil, fmt.Errorf("failed to build user prompt: %w", err)
	}

	// Create options with the system prompt
	opts := &claude.AgentOptions{
		Model:          p.config.Model,
		MaxTurns:       p.config.MaxTurns,
		MaxBudgetUSD:   p.config.MaxBudgetUSD,
		SystemPrompt:   systemPrompt,
		PermissionMode: claude.PermissionModeBypassPermission,
	}

	// Apply task-specific options
	if task.Options.Model != "" {
		opts.Model = task.Options.Model
	}

	// Run the query
	messages, err := p.queryFn(ctx, userPrompt, opts)
	if err != nil {
		failed := &Result{
			Type:     task.Type,
			Error:    err,
			Duration: time.Since(startTime),
		}
		p.recordResult(failed, false)
		return failed, nil
	}

	// Parse the result from messages
	result := &Result{
		Type:     task.Type,
		Duration: time.Since(startTime),
	}

	// Extract the final text content
	var textContent string
	for _, msg := range messages {
		switch m := msg.(type) {
		case claude.AssistantMessage:
			// Extract text from content blocks
			textContent = claude.GetText(m)
		case claude.ResultMessage:
			result.Cost = m.TotalCostUSD
			if m.Usage != nil {
				result.Tokens = m.Usage.InputTokens + m.Usage.OutputTokens
			}
		}
	}

	// Try to parse as JSON for structured output
	if textContent != "" {
		var output interface{}
		if err := json.Unmarshal([]byte(textContent), &output); err == nil {
			result.Output = output
		} else {
			// Return as plain text if not JSON
			result.Output = textContent
		}
	}

	p.recordResult(result, true)

	return result, nil
}

// recordResult folds one finished task into the processor stats under mu.
// All stats reads (Stats) and writes (here) go through mu so the counters,
// token/cost totals, and duration accumulator never race.
func (p *Processor) recordResult(result *Result, ok bool) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.stats.TasksProcessed++
	if ok {
		p.stats.TasksSucceeded++
	} else {
		p.stats.TasksFailed++
	}
	p.stats.TotalTokens += int64(result.Tokens)
	p.stats.TotalCostUSD += result.Cost
	p.totalDuration += result.Duration
}

// ProcessBatch runs multiple tasks concurrently.
func (p *Processor) ProcessBatch(ctx context.Context, tasks []Task) ([]*Result, error) {
	if p.closed.Load() {
		return nil, fmt.Errorf("processor is closed")
	}

	results := make([]*Result, len(tasks))
	var wg sync.WaitGroup
	var mu sync.Mutex
	var firstErr error

	for i, task := range tasks {
		wg.Add(1)
		go func(idx int, t Task) {
			defer wg.Done()

			result, err := p.Process(ctx, t)
			if err != nil {
				mu.Lock()
				if firstErr == nil {
					firstErr = err
				}
				mu.Unlock()
				return
			}

			mu.Lock()
			results[idx] = result
			mu.Unlock()
		}(i, task)
	}

	wg.Wait()

	if firstErr != nil {
		return results, firstErr
	}

	return results, nil
}

// Stats returns the processor statistics. AverageDuration is derived from
// the accumulated total duration over the number of processed tasks.
func (p *Processor) Stats() ProcessorStats {
	p.mu.RLock()
	defer p.mu.RUnlock()
	s := p.stats
	if s.TasksProcessed > 0 {
		s.AverageDuration = p.totalDuration / time.Duration(s.TasksProcessed)
	}
	return s
}

// Close shuts down the processor.
func (p *Processor) Close() error {
	p.closed.Store(true)
	return nil
}

// buildUserPrompt constructs the user prompt from task input.
func (p *Processor) buildUserPrompt(task Task) (string, error) {
	// Marshal the input to JSON
	inputJSON, err := json.MarshalIndent(task.Input, "", "  ")
	if err != nil {
		return "", fmt.Errorf("failed to marshal input: %w", err)
	}

	// Add context if present
	var contextStr string
	if len(task.Context) > 0 {
		contextJSON, err := json.MarshalIndent(task.Context, "", "  ")
		if err == nil {
			contextStr = fmt.Sprintf("\n\nContext:\n%s", string(contextJSON))
		}
	}

	return fmt.Sprintf("Process the following data:\n\n%s%s", string(inputJSON), contextStr), nil
}
