"""EvolutionMixin - 将进化引擎集成到 Agent 生命周期

在任务完成后自动触发反思 → 优化 → A/B 测试 → 应用/回滚的进化流程。
"""

import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any

from sqlalchemy.exc import DBAPIError

from agentkit.core.protocol import EvolutionEvent, TaskMessage, TaskResult
from agentkit.evolution.ab_tester import ABTestConfig, ABTestResult, ABTester
from agentkit.evolution.evolution_store import EvolutionStore
from agentkit.evolution.llm_reflector import LLMReflector
from agentkit.evolution.prompt_optimizer import (
    Module,
    PromptOptimizer,
)
from agentkit.evolution.reflector import Reflection, Reflector, RuleBasedReflector
from agentkit.evolution.strategy_tuner import StrategyConfig, StrategyTuner
from agentkit.memory.profile import MemoryStore

logger = logging.getLogger(__name__)


@dataclass
class SoulEvolutionConfig:
    """Soul 进化多维触发配置"""

    min_reflections: int = 3
    reflection_window_seconds: int = 3600
    time_decay_factor: float = 0.5
    task_type_weights: dict[str, float] = field(default_factory=dict)
    quality_gradient_threshold: float = -0.15


@dataclass
class EvolutionLogEntry:
    """进化日志条目"""
    task_id: str
    reflection: Reflection | None = None
    optimized_module: Module | None = None
    ab_test_result: ABTestResult | None = None
    applied: bool = False
    rolled_back: bool = False
    event_id: str | None = None
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


class EvolutionMixin:
    """进化混入类，将进化引擎集成到 Agent 生命周期。

    用法：
        class MyAgent(BaseAgent, EvolutionMixin):
            def __init__(self, ...):
                BaseAgent.__init__(self, ...)
                EvolutionMixin.__init__(self, reflector=..., ...)
    """

    _UNSET = object()  # 用于区分"未传入"和"显式传入 None"

    def __init__(
        self,
        reflector: Any = _UNSET,
        prompt_optimizer: PromptOptimizer | None = None,
        strategy_tuner: StrategyTuner | None = None,
        ab_tester: ABTester | None = None,
        evolution_store: EvolutionStore | None = None,
        reflector_type: str | None = None,
        llm_gateway: Any | None = None,
        auxiliary_model: str | None = None,
        strategy_tuning_enabled: bool = False,
        evolution_config: SoulEvolutionConfig | None = None,
    ):
        if reflector is not EvolutionMixin._UNSET:
            # 显式传入了 reflector 参数（包括 None）
            self._reflector = reflector
        elif reflector_type is not None:
            # 未传入 reflector，但指定了 reflector_type → 自动创建
            self._reflector = self._create_reflector(
                reflector_type, llm_gateway, auxiliary_model
            )
        else:
            # 都未指定：保持向后兼容，reflector 为 None
            self._reflector = None
        self._prompt_optimizer = prompt_optimizer
        self._strategy_tuner = strategy_tuner
        self._ab_tester = ab_tester
        self._evolution_store = evolution_store
        self._evolution_log: list[EvolutionLogEntry] = []
        self._current_module: Module | None = None
        self._strategy_tuning_enabled = strategy_tuning_enabled
        self._evolution_config = evolution_config
        self.pending_soul_updates: dict[str, list] = {}

    @staticmethod
    def _create_reflector(
        reflector_type: str,
        llm_gateway: Any | None = None,
        auxiliary_model: str | None = None,
    ) -> Reflector | None:
        """根据 reflector_type 创建对应的反思器

        Args:
            reflector_type: "llm" / "rule" / "auto"
            llm_gateway: LLMGateway 实例，llm/auto 模式需要
            auxiliary_model: LLM 反思使用的模型名称
        """
        if reflector_type == "llm":
            if llm_gateway is None:
                logger.warning(
                    "reflector_type='llm' but no llm_gateway provided, "
                    "falling back to RuleBasedReflector"
                )
                return RuleBasedReflector()
            model = auxiliary_model or "default"
            return LLMReflector(llm_gateway=llm_gateway, model=model)

        if reflector_type == "rule":
            return RuleBasedReflector()

        # "auto" 模式：优先 LLM，降级到规则
        if llm_gateway is not None:
            model = auxiliary_model or "default"
            return LLMReflector(llm_gateway=llm_gateway, model=model)

        return RuleBasedReflector()

    async def evolve_after_task(
        self,
        task: TaskMessage,
        result: TaskResult,
        memory_store: MemoryStore | None = None,
    ) -> EvolutionLogEntry:
        """任务完成后执行进化流程。

        流程：
        1. Reflector 反思 → 得到 Reflection
        2. Soul 进化检查（如果 memory_store 可用）
        3. 如果 Reflection 有改进建议 → PromptOptimizer 优化
        4. 如果优化产生了新 Prompt → ABTester 验证
        5. 如果 AB 测试通过 → EvolutionStore 应用变更
        6. 如果 AB 测试失败 → 回滚
        7. 如果策略调优启用 → StrategyTuner 调优
        """
        log_entry = EvolutionLogEntry(task_id=task.task_id)

        # Step 1: 反思
        if self._reflector is None:
            logger.debug("No reflector configured, skipping evolution")
            self._evolution_log.append(log_entry)
            return log_entry

        reflection = await self._reflector.reflect(task, result)
        log_entry.reflection = reflection

        logger.info(
            f"Evolution reflection for task {task.task_id}: "
            f"outcome={reflection.outcome}, quality={reflection.quality_score:.2f}, "
            f"suggestions={len(reflection.suggestions)}"
        )

        # Step 2: Soul 进化检查
        if memory_store is not None:
            await self.evolve_soul(task, result, memory_store, reflection=reflection)

        # Step 3: 如果有改进建议，触发 Prompt 优化
        if not reflection.suggestions:
            logger.debug("No improvement suggestions, skipping optimization")
            self._evolution_log.append(log_entry)
            return log_entry

        if self._prompt_optimizer is None or self._current_module is None:
            logger.debug("No prompt optimizer or current module configured, skipping optimization")
            self._evolution_log.append(log_entry)
            return log_entry

        # 将反思结果作为训练样本
        self._prompt_optimizer.add_example(
            input_data=task.input_data,
            output_data=result.output_data or {},
            quality_score=reflection.quality_score,
        )

        # Pass trace and reflection to LLMPromptOptimizer if available
        optimized = await self._optimize_with_context(self._current_module, reflection)

        # 检查是否真正产生了变化
        if optimized.name == self._current_module.name and not optimized.demos:
            logger.debug("Optimization produced no meaningful changes")
            self._evolution_log.append(log_entry)
            return log_entry

        log_entry.optimized_module = optimized

        # Step 3: A/B 测试验证
        if self._ab_tester is None:
            logger.debug("No AB tester configured, applying change directly")
            applied = await self._apply_change(task, result, optimized, reflection)
            log_entry.applied = applied
            # Strategy tuning (if enabled)
            if self._strategy_tuning_enabled and self._strategy_tuner is not None:
                await self._run_strategy_tuning(task, result, reflection)
            self._evolution_log.append(log_entry)
            return log_entry

        # Run A/B test
        ab_result = await self._run_ab_test(task, result, optimized, reflection)
        log_entry.ab_test_result = ab_result

        if ab_result is None or not ab_result.is_significant:
            # Insufficient samples or inconclusive
            if ab_result is None:
                logger.info("Insufficient data for A/B test, keeping current prompt")
            else:
                logger.info(
                    f"A/B test inconclusive (p={ab_result.p_value}), keeping current prompt"
                )
            # Don't apply the change, don't rollback either — just keep current
            self._evolution_log.append(log_entry)
            return log_entry

        if ab_result.winner == "experiment":
            # Treatment wins → apply optimized prompt
            logger.info("A/B test significant: treatment wins, applying optimized prompt")
            applied = await self._apply_change(task, result, optimized, reflection)
            log_entry.applied = applied
        else:
            # Control wins → rollback, keep original
            logger.info("A/B test significant: control wins, keeping original prompt")
            rolled_back = await self._rollback_change(log_entry)
            log_entry.rolled_back = rolled_back

        # Step 4: Strategy tuning (if enabled)
        if self._strategy_tuning_enabled and self._strategy_tuner is not None:
            await self._run_strategy_tuning(task, result, reflection)

        self._evolution_log.append(log_entry)
        return log_entry

    async def _optimize_with_context(
        self, module: Module, reflection: Reflection
    ) -> Module:
        """Run optimization, passing reflection context if optimizer supports it"""
        from agentkit.evolution.prompt_optimizer import LLMPromptOptimizer

        if isinstance(self._prompt_optimizer, LLMPromptOptimizer):
            return await self._prompt_optimizer.optimize(module, trace=None, reflection=reflection)

        return await self._prompt_optimizer.optimize(module)

    async def _run_ab_test(
        self,
        task: TaskMessage,
        result: TaskResult,
        optimized: Module,
        reflection: Reflection,
    ) -> ABTestResult | None:
        """Run A/B test: assign group → record result → evaluate"""
        test_id = f"evolve_{task.task_id}"

        # Create test if not exists
        if test_id not in self._ab_tester._tests:
            self._ab_tester.create_test(ABTestConfig(
                test_id=test_id,
                agent_name=result.agent_name,
                change_type="prompt",
            ))

        # Assign group deterministically based on task_id
        group = self._ab_tester.assign_group(test_id, task_id=task.task_id)

        # Record the current task result
        self._ab_tester.record_result(test_id, group, reflection.quality_score)

        # Persist results if store is available
        await self._ab_tester.persist_results(test_id)

        # Evaluate
        return await self._ab_tester.evaluate(test_id)

    async def _run_strategy_tuning(
        self,
        task: TaskMessage,
        result: TaskResult,
        reflection: Reflection,
    ) -> None:
        """Run strategy tuning with trace metrics"""
        if self._strategy_tuner is None:
            return

        # Build current strategy config from result metrics
        current_config = StrategyConfig(
            temperature=0.5,
            max_iterations=5,
        )

        # Record the current result
        self._strategy_tuner.record(current_config, reflection.quality_score)

        # Get suggestion
        suggested = await self._strategy_tuner.suggest(current_config)
        logger.info(
            f"Strategy tuning suggestion for task {task.task_id}: "
            f"temperature={suggested.temperature:.2f}, "
            f"max_iterations={suggested.max_iterations}"
        )

    def get_evolution_history(self) -> list[dict[str, Any]]:
        """获取进化历史记录"""
        history = []
        for entry in self._evolution_log:
            record: dict[str, Any] = {
                "task_id": entry.task_id,
                "applied": entry.applied,
                "rolled_back": entry.rolled_back,
                "event_id": entry.event_id,
                "created_at": entry.created_at.isoformat(),
            }
            if entry.reflection:
                record["reflection"] = {
                    "outcome": entry.reflection.outcome,
                    "quality_score": entry.reflection.quality_score,
                    "suggestions": entry.reflection.suggestions,
                }
            if entry.optimized_module:
                record["optimized_module"] = entry.optimized_module.name
            if entry.ab_test_result:
                record["ab_test"] = {
                    "winner": entry.ab_test_result.winner,
                    "is_significant": entry.ab_test_result.is_significant,
                }
            history.append(record)
        return history

    def set_current_module(self, module: Module | None = None) -> None:
        """设置当前 Prompt 模块

        Args:
            module: Module 实例。如果为 None，子类应自行创建。
        """
        self._current_module = module

    async def _apply_change(
        self,
        task: TaskMessage,
        result: TaskResult,
        optimized: Module,
        reflection: Reflection,
    ) -> bool:
        """应用优化变更"""
        if self._evolution_store is None:
            # 无存储时直接更新内存中的模块
            self._current_module = optimized
            return True

        event = EvolutionEvent(
            agent_name=result.agent_name,
            change_type="prompt",
            before={"module_name": self._current_module.name if self._current_module else ""},
            after={"module_name": optimized.name, "demos_count": len(optimized.demos)},
            metrics={"quality_score": reflection.quality_score},
        )

        try:
            event_id = await self._evolution_store.record(event)
            self._current_module = optimized
            # 回写 event_id 到对应的 log entry
            for entry in reversed(self._evolution_log):
                if entry.task_id == task.task_id and entry.event_id is None:
                    entry.event_id = event_id
                    break
            return True
        except (DBAPIError, RuntimeError, ValueError, KeyError) as e:
            logger.error(f"Failed to apply evolution change: {e}")
            return False

    async def _rollback_change(self, log_entry: EvolutionLogEntry) -> bool:
        """回滚进化变更"""
        if self._evolution_store is None or log_entry.event_id is None:
            return True

        try:
            return await self._evolution_store.rollback(log_entry.event_id)
        except (DBAPIError, RuntimeError, ValueError, KeyError) as e:
            logger.error(f"Failed to rollback evolution change: {e}")
            return False

    def record_reflection(
        self,
        pattern: str,
        reflection: Reflection,
        task_type: str = "",
        score: float | None = None,
    ) -> None:
        """记录反思到待处理列表，附带时间戳、分数和任务类型。"""
        if pattern not in self.pending_soul_updates:
            self.pending_soul_updates[pattern] = []
        self.pending_soul_updates[pattern].append(
            {
                "reflection": reflection,
                "timestamp": datetime.now(timezone.utc),
                "score": score if score is not None else reflection.quality_score,
                "task_type": task_type,
            }
        )

    async def evolve_soul(
        self,
        task: TaskMessage,
        result: TaskResult,
        memory_store: MemoryStore | None = None,
        reflection: Reflection | None = None,
        task_type: str = "",
        score: float | None = None,
    ) -> bool:
        """Check if soul should be updated based on accumulated reflections.

        Multi-dimensional triggers:
        - Time decay: older reflections contribute less
        - Quality gradient: declining scores trigger early
        - Task type weight: different task types have different trigger thresholds
        - Trigger threshold: effective_count * weight >= min_reflections
        """
        if memory_store is None:
            return False

        if reflection is None:
            if self._reflector is None:
                return False
            reflection = await self._reflector.reflect(task, result)

        # 只关注低质量且有建议的反思
        if reflection.quality_score >= 0.5:
            return False

        if not reflection.suggestions:
            return False

        config = self._evolution_config or SoulEvolutionConfig()

        # 按 pattern 分类累积反思（patterns为空时使用默认category）
        categories = reflection.patterns if reflection.patterns else ["default"]
        for pattern in categories:
            self.record_reflection(
                pattern, reflection, task_type=task_type, score=score
            )

        # 检查是否有类别满足触发条件
        for category, reflections in list(self.pending_soul_updates.items()):
            # --- Quality gradient: 3+ declining scores trigger early ---
            scores = [r["score"] for r in reflections if r["score"] is not None]
            quality_gradient_triggered = False
            if len(scores) >= 3:
                last_3 = scores[-3:]
                declines = [
                    last_3[i] - last_3[i - 1] for i in range(1, len(last_3))
                ]
                if all(d <= config.quality_gradient_threshold for d in declines):
                    quality_gradient_triggered = True

            # --- Time decay: compute effective_count ---
            now = datetime.now(timezone.utc)
            effective_count = 0.0
            for r in reflections:
                age_seconds = (now - r["timestamp"]).total_seconds()
                age_hours = age_seconds / 3600.0
                effective_count += config.time_decay_factor ** age_hours
            # Round to avoid floating-point precision issues
            # (e.g. 3 recent reflections should yield exactly 3.0)
            effective_count = round(effective_count, 6)

            # --- Task type weight ---
            weight = 1.0
            if task_type and task_type in config.task_type_weights:
                weight = config.task_type_weights[task_type]

            # --- Trigger threshold: effective_count * weight >= min_reflections ---
            weighted_count = effective_count * weight
            if weighted_count >= config.min_reflections or quality_gradient_triggered:
                # 触发 soul 更新
                from agentkit.tools.memory_tool import MemoryTool

                tool = MemoryTool(memory_store)
                section = category
                # 汇总所有累积反思的建议（去重，最多取 5 条）
                all_suggestions: list[str] = []
                seen: set[str] = set()
                for r in reflections:
                    for suggestion in r["reflection"].suggestions:
                        if suggestion not in seen:
                            seen.add(suggestion)
                            all_suggestions.append(suggestion)
                content = "; ".join(all_suggestions[:5])
                reason = f"连续{len(reflections)}次低质量反思 (category: {category})"

                update_result = await tool.execute(
                    action="update_soul",
                    file="soul",
                    section=section,
                    content=content,
                    reason=reason,
                )

                if update_result.get("success"):
                    logger.info(
                        f"Soul evolved: category={category}, "
                        f"version={update_result.get('version')}"
                    )
                    # 清除已处理的类别
                    del self.pending_soul_updates[category]
                    return True

        return False