package router

import (
	"context"
	"sync"
	"time"

	"gokin/internal/agent"
	"gokin/internal/client"
	"gokin/internal/logging"
	"gokin/internal/tools"

	"google.golang.org/genai"
)

// ExampleStoreInterface defines the interface for example stores.
// This is implemented by memory.ExampleStore to avoid import cycles.
type ExampleStoreInterface interface {
	GetSimilarExamples(prompt string, limit int) []ExampleSummary
	GetExamplesForContext(taskType, prompt string, limit int) string
	LearnFromSuccess(taskType, prompt, agentType, output string, duration time.Duration, tokens int) error
}

// ExampleSummary contains a summary of a task example.
type ExampleSummary struct {
	ID          string
	TaskType    string
	InputPrompt string
	AgentType   string
	Duration    time.Duration
	Score       float64
}

// SmartRouterConfig holds configuration for the smart router.
type SmartRouterConfig struct {
	*RouterConfig

	// Adaptive selection settings
	AdaptiveEnabled bool
	MinDataPoints   int // Minimum executions before using learned data
	ExampleLimit    int // Maximum examples to include in context
}

// DefaultSmartRouterConfig returns the default smart router configuration.
func DefaultSmartRouterConfig() *SmartRouterConfig {
	return &SmartRouterConfig{
		RouterConfig: &RouterConfig{
			Enabled:            true,
			DecomposeThreshold: 4,
			ParallelThreshold:  7,
		},
		AdaptiveEnabled: true,
		MinDataPoints:   5,
		ExampleLimit:    3,
	}
}

// SmartRouter extends Router with adaptive agent selection and few-shot learning.
type SmartRouter struct {
	*Router

	// Parent context for background operations
	ctx context.Context

	// Strategy optimizer for learning from outcomes
	optimizer *agent.StrategyOptimizer

	// Example store for few-shot learning
	exampleStore ExampleStoreInterface

	// Configuration
	adaptiveEnabled bool
	minDataPoints   int
	exampleLimit    int

	// Protects optimizer and exampleStore from concurrent access
	mu sync.RWMutex
}

// NewSmartRouter creates a new smart router with adaptive capabilities.
func NewSmartRouter(ctx context.Context, cfg *SmartRouterConfig, executor *tools.Executor, agentRunner AgentRunner, c client.Client, workDir string) *SmartRouter {
	if cfg == nil {
		cfg = DefaultSmartRouterConfig()
	}

	baseRouter := NewRouter(cfg.RouterConfig, executor, agentRunner, c, nil, false, workDir)

	return &SmartRouter{
		Router:          baseRouter,
		ctx:             ctx,
		adaptiveEnabled: cfg.AdaptiveEnabled,
		minDataPoints:   cfg.MinDataPoints,
		exampleLimit:    cfg.ExampleLimit,
	}
}

// SetStrategyOptimizer sets the strategy optimizer for learning.
func (sr *SmartRouter) SetStrategyOptimizer(optimizer *agent.StrategyOptimizer) {
	sr.mu.Lock()
	defer sr.mu.Unlock()
	sr.optimizer = optimizer
}

// SetExampleStore sets the example store for few-shot learning.
func (sr *SmartRouter) SetExampleStore(store ExampleStoreInterface) {
	sr.mu.Lock()
	defer sr.mu.Unlock()
	sr.exampleStore = store
}

// Route determines the best execution strategy with adaptive learning.
func (sr *SmartRouter) Route(message string) *RoutingDecision {
	return sr.RouteWithContext(context.Background(), message)
}

// RouteWithContext is like Route but propagates the caller's context.
func (sr *SmartRouter) RouteWithContext(ctx context.Context, message string) *RoutingDecision {
	// Get base routing decision
	decision := sr.Router.RouteWithContext(ctx, message)

	// Snapshot fields under lock
	sr.mu.RLock()
	exampleStore := sr.exampleStore
	optimizer := sr.optimizer
	sr.mu.RUnlock()

	// Add learned examples if available
	if exampleStore != nil && sr.exampleLimit > 0 {
		examples := exampleStore.GetSimilarExamples(message, sr.exampleLimit)
		if len(examples) > 0 {
			decision.LearnedExamples = make([]LearnedExample, len(examples))
			for i, ex := range examples {
				decision.LearnedExamples[i] = LearnedExample{
					ID:        ex.ID,
					TaskType:  ex.TaskType,
					Prompt:    ex.InputPrompt,
					AgentType: ex.AgentType,
					Score:     ex.Score,
				}
			}

			runes := []rune(message)
			msgPreview := message
			if len(runes) > 50 {
				msgPreview = string(runes[:50])
			}
			logging.Debug("found similar examples for routing",
				"count", len(examples),
				"message", msgPreview)
		}
	}

	// Override agent selection if we have enough data
	if decision.Handler == HandlerSubAgent && sr.adaptiveEnabled && optimizer != nil {
		taskType := string(decision.Analysis.Type)
		metrics, hasData := optimizer.GetMetrics(taskType)

		if hasData && totalCount(metrics) >= sr.minDataPoints {
			recommended := optimizer.RecommendStrategy(taskType)
			if recommended != decision.SubAgentType {
				logging.Debug("adaptive override",
					"original", decision.SubAgentType,
					"recommended", recommended,
					"task_type", taskType,
					"data_points", totalCount(metrics))

				decision.SubAgentType = recommended
				decision.Reasoning = decision.Reasoning + " (adaptive: " + recommended + ")"
			}
		}
	}

	return decision
}

// Execute routes the task with learning from outcomes.
func (sr *SmartRouter) Execute(ctx context.Context, history []*genai.Content, message string) ([]*genai.Content, string, error) {
	startTime := time.Now()
	decision := sr.RouteWithContext(ctx, message)

	// Execute using base router
	resultHistory, output, err := sr.Router.Execute(ctx, history, message)
	duration := time.Since(startTime)

	// Learn from the outcome
	success := err == nil && output != ""
	taskType := string(decision.Analysis.Type)
	agentType := decision.SubAgentType
	if agentType == "" {
		agentType = string(decision.Handler)
	}

	// Snapshot fields under lock
	sr.mu.RLock()
	execOptimizer := sr.optimizer
	execExampleStore := sr.exampleStore
	sr.mu.RUnlock()

	// Record with strategy optimizer
	if execOptimizer != nil {
		execOptimizer.RecordExecution(agentType, taskType, success, duration)
	}

	// Record successful executions for few-shot learning
	// Use a bounded goroutine with timeout to prevent hanging
	if success && execExampleStore != nil {
		go func() {
			// Timeout to prevent goroutine leak if store is slow
			ctx, cancel := context.WithTimeout(sr.ctx, 10*time.Second)
			defer cancel()

			done := make(chan struct{})
			go func() {
				defer close(done)
				defer func() {
					if r := recover(); r != nil {
						logging.Warn("panic in LearnFromSuccess goroutine", "panic", r)
					}
				}()
				if err := execExampleStore.LearnFromSuccess(taskType, message, agentType, output, duration, 0); err != nil {
					logging.Debug("failed to learn from success", "error", err)
				}
			}()

			select {
			case <-done:
				// Completed successfully
			case <-ctx.Done():
				logging.Warn("learning from success timed out")
			}
		}()
	}

	return resultHistory, output, err
}

// RouteWithLearning returns a routing decision with learned examples included.
func (sr *SmartRouter) RouteWithLearning(message string) *RoutingDecision {
	return sr.Route(message)
}

// GetExamplesContext returns formatted examples for prompt injection.
func (sr *SmartRouter) GetExamplesContext(taskType, prompt string) string {
	if sr.exampleStore == nil {
		return ""
	}
	return sr.exampleStore.GetExamplesForContext(taskType, prompt, sr.exampleLimit)
}

// GetAdaptiveStats returns statistics about adaptive routing.
func (sr *SmartRouter) GetAdaptiveStats() AdaptiveStats {
	stats := AdaptiveStats{
		Enabled:       sr.adaptiveEnabled,
		MinDataPoints: sr.minDataPoints,
	}

	if sr.optimizer != nil {
		allMetrics := sr.optimizer.GetAllMetrics()
		stats.TotalStrategies = len(allMetrics)
		stats.TopStrategies = make([]StrategyStats, 0)

		for name, metrics := range allMetrics {
			stats.TopStrategies = append(stats.TopStrategies, StrategyStats{
				Name:        name,
				SuccessRate: metrics.SuccessRate(),
				UseCount:    totalCount(metrics),
			})
		}
	}

	return stats
}

// AdaptiveStats contains statistics about adaptive routing.
type AdaptiveStats struct {
	Enabled         bool
	MinDataPoints   int
	TotalStrategies int
	TopStrategies   []StrategyStats
}

// StrategyStats contains statistics about a strategy.
type StrategyStats struct {
	Name        string
	SuccessRate float64
	UseCount    int
}

// totalCount returns the total number of executions for a strategy.
func totalCount(m *agent.StrategyMetrics) int {
	return m.SuccessCount + m.FailureCount
}

