"""PlanChecker — 计划检查与复盘

每步执行后检查产出质量，全部完成后复盘总结并写入经验库。

核心能力：
1. QualityGate: 每步完成后验证产出（required_fields / min_word_count / 自定义校验）
2. LLMReflector: 使用 LLM 评估步骤质量（可选，回退到规则评估）
3. ReviewReport: 全部完成后生成复盘报告（成功路径、失败原因、耗时分布、优化建议）
4. ExperienceStore: 复盘结果写入经验库（可选依赖）

使用方式：
    checker = PlanChecker()
    result = await checker.check_step(step, exec_result)
    report = await checker.review_plan(plan, plan_result)
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Awaitable

from agentkit.core.plan_schema import ExecutionPlan, PlanStep, PlanStepStatus
from agentkit.core.plan_executor import PlanExecutionResult, StepExecutionResult
from agentkit.skills.base import QualityGateConfig

logger = logging.getLogger(__name__)


class CheckStatus(str, Enum):
    """检查结果状态"""

    PASS = "pass"
    FAIL = "fail"
    SKIP = "skip"


@dataclass
class CheckResult:
    """单步检查结果

    Attributes:
        step_id: 步骤 ID
        status: 检查状态（pass/fail/skip）
        reason: 检查原因说明
        quality_score: 质量评分（0.0 ~ 1.0）
        details: 详细检查项
    """

    step_id: str
    status: CheckStatus
    reason: str = ""
    quality_score: float = 0.0
    details: dict[str, Any] = field(default_factory=dict)


@dataclass
class ReviewReport:
    """复盘报告

    全部步骤完成后生成，包含成功路径、失败原因、耗时分布和优化建议。

    Attributes:
        plan_id: 计划 ID
        outcome: 整体结果（"success" / "partial" / "failure"）
        success_path: 成功步骤路径（按执行顺序）
        failure_reasons: 失败原因列表
        duration_distribution: 各步骤耗时分布
        optimization_tips: 优化建议
        quality_scores: 各步骤质量评分
        total_duration_ms: 总耗时
        success_rate: 成功率
    """

    plan_id: str
    outcome: str = "success"
    success_path: list[str] = field(default_factory=list)
    failure_reasons: list[str] = field(default_factory=list)
    duration_distribution: dict[str, float] = field(default_factory=dict)
    optimization_tips: list[str] = field(default_factory=list)
    quality_scores: dict[str, float] = field(default_factory=dict)
    total_duration_ms: float = 0.0
    success_rate: float = 1.0

    def to_dict(self) -> dict[str, Any]:
        """序列化为字典"""
        return {
            "plan_id": self.plan_id,
            "outcome": self.outcome,
            "success_path": self.success_path,
            "failure_reasons": self.failure_reasons,
            "duration_distribution": self.duration_distribution,
            "optimization_tips": self.optimization_tips,
            "quality_scores": self.quality_scores,
            "total_duration_ms": self.total_duration_ms,
            "success_rate": self.success_rate,
        }


# 自定义校验器类型：接收步骤结果，返回 (通过, 原因)
CustomValidator = Callable[[dict[str, Any] | None], tuple[bool, str]]


class QualityGate:
    """质量门控

    基于 QualityGateConfig 验证步骤产出：
    1. required_fields: 结果字典必须包含指定字段
    2. min_word_count: 结果文本字段最少字数
    3. custom_validator: 自定义校验函数
    """

    def __init__(
        self,
        config: QualityGateConfig | None = None,
        custom_validator: CustomValidator | None = None,
    ):
        self._config = config or QualityGateConfig()
        self._custom_validator = custom_validator

    def check(self, step: PlanStep, exec_result: StepExecutionResult) -> CheckResult:
        """检查步骤产出质量

        Args:
            step: 计划步骤
            exec_result: 步骤执行结果

        Returns:
            CheckResult: 检查结果
        """
        # 跳过非完成步骤
        if exec_result.status != PlanStepStatus.COMPLETED:
            return CheckResult(
                step_id=step.step_id,
                status=CheckStatus.SKIP,
                reason=f"Step status is {exec_result.status.value}, skipping quality check",
            )

        result = exec_result.result
        details: dict[str, Any] = {}
        failures: list[str] = []

        # 1. 检查 required_fields
        missing_fields = self._check_required_fields(result)
        if missing_fields:
            failures.append(f"Missing required fields: {', '.join(missing_fields)}")
            details["missing_fields"] = missing_fields

        # 2. 检查 min_word_count
        word_count_result = self._check_min_word_count(result)
        if word_count_result:
            failures.append(word_count_result)
            details["word_count_check"] = word_count_result

        # 3. 自定义校验
        custom_result = self._check_custom(result)
        if custom_result:
            failures.append(custom_result)
            details["custom_check"] = custom_result

        if failures:
            return CheckResult(
                step_id=step.step_id,
                status=CheckStatus.FAIL,
                reason="; ".join(failures),
                quality_score=self._compute_quality_score(len(failures)),
                details=details,
            )

        return CheckResult(
            step_id=step.step_id,
            status=CheckStatus.PASS,
            reason="All quality checks passed",
            quality_score=1.0,
            details=details,
        )

    def _check_required_fields(self, result: dict[str, Any] | None) -> list[str]:
        """检查必填字段"""
        if not self._config.required_fields:
            return []
        if result is None:
            return list(self._config.required_fields)
        return [f for f in self._config.required_fields if f not in result]

    def _check_min_word_count(self, result: dict[str, Any] | None) -> str:
        """检查最少字数"""
        if self._config.min_word_count <= 0:
            return ""
        if result is None:
            return f"Result is None, cannot check min_word_count ({self._config.min_word_count})"

        total_words = 0
        for value in result.values():
            if isinstance(value, str):
                total_words += len(value.split())

        if total_words < self._config.min_word_count:
            return (
                f"Word count ({total_words}) is below minimum "
                f"({self._config.min_word_count})"
            )
        return ""

    def _check_custom(self, result: dict[str, Any] | None) -> str:
        """执行自定义校验"""
        if self._custom_validator is None:
            return ""
        try:
            passed, reason = self._custom_validator(result)
            if not passed:
                return reason or "Custom validation failed"
        except Exception as e:
            return f"Custom validator error: {e}"
        return ""

    @staticmethod
    def _compute_quality_score(failure_count: int) -> float:
        """根据失败项数计算质量评分"""
        if failure_count == 0:
            return 1.0
        if failure_count == 1:
            return 0.5
        if failure_count == 2:
            return 0.25
        return 0.1


class RuleBasedStepReflector:
    """基于规则的步骤反思器

    评估步骤执行质量，生成质量评分和改进建议。
    当 LLM 不可用时的回退方案。
    """

    async def reflect_step(
        self,
        step: PlanStep,
        exec_result: StepExecutionResult,
    ) -> tuple[float, list[str]]:
        """对步骤执行结果进行反思

        Args:
            step: 计划步骤
            exec_result: 步骤执行结果

        Returns:
            (quality_score, suggestions): 质量评分和改进建议
        """
        suggestions: list[str] = []

        if exec_result.status != PlanStepStatus.COMPLETED:
            # 失败步骤
            score = 0.0
            if exec_result.error:
                if "timed out" in exec_result.error.lower():
                    suggestions.append(
                        f"Step '{step.name}' timed out: consider increasing timeout or decomposing the task"
                    )
                elif "no agent available" in exec_result.error.lower():
                    suggestions.append(
                        f"Step '{step.name}' had no available agent: check skill registry"
                    )
                else:
                    suggestions.append(
                        f"Step '{step.name}' failed: {exec_result.error}"
                    )
            return score, suggestions

        # 成功步骤
        score = 0.6  # 基础分

        # 有输出数据加分
        if exec_result.result and len(exec_result.result) > 0:
            score += 0.2

        # 无重试加分
        if exec_result.retry_count == 0:
            score += 0.1

        # 耗时合理加分
        if exec_result.duration_ms > 0 and exec_result.duration_ms < 30000:
            score += 0.1

        score = min(score, 1.0)

        # 生成建议
        if exec_result.retry_count > 0:
            suggestions.append(
                f"Step '{step.name}' required {exec_result.retry_count} retries: "
                f"consider improving step reliability"
            )

        if exec_result.duration_ms > 60000:
            suggestions.append(
                f"Step '{step.name}' took {exec_result.duration_ms / 1000:.1f}s: "
                f"consider optimizing for performance"
            )

        return score, suggestions


class PlanChecker:
    """计划检查器

    每步执行后检查产出质量，全部完成后复盘总结并写入经验库。

    检查环节：每步完成后，QualityGate 验证产出 + Reflector 评估是否达标
    复盘环节：全部完成后，生成复盘报告（成功路径、失败原因、耗时分布）
    经验写入：复盘结果写入 ExperienceStore（可选）
    闭环：检查不通过 → 触发重试或计划调整

    使用方式：
        # 独立使用
        checker = PlanChecker()
        result = await checker.check_step(step, exec_result)
        report = await checker.review_plan(plan, plan_result)

        # 与 PlanExecutor 集成
        checker = PlanChecker(experience_store=store)
        executor = PlanExecutor(
            agent_pool=pool,
            on_step_complete=checker.make_step_complete_callback(),
        )
    """

    def __init__(
        self,
        quality_gate: QualityGate | None = None,
        quality_gate_config: QualityGateConfig | None = None,
        custom_validator: CustomValidator | None = None,
        reflector: Any | None = None,
        experience_store: Any | None = None,
        max_check_retries: int = 1,
        quality_threshold: float = 0.5,
        step_quality_configs: dict[str, QualityGateConfig] | None = None,
    ):
        """初始化 PlanChecker

        Args:
            quality_gate: 质量门控实例（优先使用）
            quality_gate_config: 质量门控配置（quality_gate 为 None 时使用）
            custom_validator: 自定义校验函数
            reflector: 步骤反思器（None 时使用 RuleBasedStepReflector）
            experience_store: 经验存储（None 时不写入经验库）
            max_check_retries: 检查不通过时最大重试次数
            quality_threshold: 质量评分阈值，低于此值视为不通过
            step_quality_configs: 每步骤独立的质量门控配置
        """
        if quality_gate is not None:
            self._quality_gate = quality_gate
        else:
            self._quality_gate = QualityGate(
                config=quality_gate_config,
                custom_validator=custom_validator,
            )
        self._reflector = reflector or RuleBasedStepReflector()
        self._experience_store = experience_store
        self._max_check_retries = max_check_retries
        self._quality_threshold = quality_threshold
        self._step_quality_configs = step_quality_configs or {}

        # 内部状态：记录每步检查结果
        self._check_results: dict[str, CheckResult] = {}
        self._step_quality_gates: dict[str, QualityGate] = {}

        # 为有独立配置的步骤创建 QualityGate
        for step_id, config in self._step_quality_configs.items():
            self._step_quality_gates[step_id] = QualityGate(config=config)

    async def check_step(
        self,
        step: PlanStep,
        exec_result: StepExecutionResult,
    ) -> CheckResult:
        """检查单个步骤的产出质量

        在每步完成后调用，验证产出是否达标。

        Args:
            step: 计划步骤
            exec_result: 步骤执行结果

        Returns:
            CheckResult: 检查结果
        """
        # 选择步骤专属或默认 QualityGate
        gate = self._step_quality_gates.get(step.step_id, self._quality_gate)

        # 1. QualityGate 规则检查
        gate_result = gate.check(step, exec_result)

        # 2. Reflector 评估（仅对已完成步骤）
        if exec_result.status == PlanStepStatus.COMPLETED:
            try:
                reflect_score, suggestions = await self._reflector.reflect_step(
                    step, exec_result
                )
            except Exception as e:
                logger.warning(f"Reflector failed for step '{step.step_id}': {e}")
                reflect_score = gate_result.quality_score
                suggestions = []

            # 综合评分：取 QualityGate 和 Reflector 的加权平均
            combined_score = 0.4 * gate_result.quality_score + 0.6 * reflect_score

            # 如果 Reflector 评分低于阈值，标记为不通过
            if combined_score < self._quality_threshold and gate_result.status == CheckStatus.PASS:
                gate_result = CheckResult(
                    step_id=step.step_id,
                    status=CheckStatus.FAIL,
                    reason=f"Quality score ({combined_score:.2f}) below threshold ({self._quality_threshold})",
                    quality_score=combined_score,
                    details={
                        **gate_result.details,
                        "reflector_score": reflect_score,
                        "reflector_suggestions": suggestions,
                    },
                )
            elif gate_result.status != CheckStatus.PASS:
                # 已有不通过结果，更新评分
                gate_result = CheckResult(
                    step_id=step.step_id,
                    status=gate_result.status,
                    reason=gate_result.reason,
                    quality_score=combined_score,
                    details={
                        **gate_result.details,
                        "reflector_score": reflect_score,
                        "reflector_suggestions": suggestions,
                    },
                )
            else:
                # 通过，更新评分
                gate_result = CheckResult(
                    step_id=step.step_id,
                    status=gate_result.status,
                    reason=gate_result.reason,
                    quality_score=combined_score,
                    details={
                        **gate_result.details,
                        "reflector_score": reflect_score,
                        "reflector_suggestions": suggestions,
                    },
                )

        # 记录检查结果
        self._check_results[step.step_id] = gate_result

        logger.info(
            f"Check step '{step.step_id}': status={gate_result.status.value}, "
            f"score={gate_result.quality_score:.2f}, reason={gate_result.reason}"
        )

        return gate_result

    async def review_plan(
        self,
        plan: ExecutionPlan,
        plan_result: PlanExecutionResult,
        task_type: str = "",
        goal: str = "",
    ) -> ReviewReport:
        """复盘整个计划执行结果

        全部步骤完成后调用，生成复盘报告并写入经验库。

        Args:
            plan: 执行计划
            plan_result: 计划执行结果
            task_type: 任务类型（写入经验库用）
            goal: 任务目标（写入经验库用）

        Returns:
            ReviewReport: 复盘报告
        """
        # 1. 构建成功路径
        success_path = plan_result.completed_steps

        # 2. 收集失败原因
        failure_reasons = self._collect_failure_reasons(plan_result)

        # 3. 构建耗时分布
        duration_distribution = {
            sid: r.duration_ms
            for sid, r in plan_result.step_results.items()
        }

        # 4. 收集质量评分
        quality_scores = {
            sid: cr.quality_score
            for sid, cr in self._check_results.items()
        }

        # 5. 计算成功率
        total_steps = len(plan.steps)
        completed_count = len(plan_result.completed_steps)
        success_rate = completed_count / total_steps if total_steps > 0 else 0.0

        # 6. 判断整体结果
        outcome = self._determine_outcome(plan_result)

        # 7. 生成优化建议
        optimization_tips = self._generate_optimization_tips(
            plan_result, quality_scores
        )

        report = ReviewReport(
            plan_id=plan.plan_id,
            outcome=outcome,
            success_path=success_path,
            failure_reasons=failure_reasons,
            duration_distribution=duration_distribution,
            optimization_tips=optimization_tips,
            quality_scores=quality_scores,
            total_duration_ms=plan_result.total_duration_ms,
            success_rate=success_rate,
        )

        logger.info(
            f"Review plan '{plan.plan_id}': outcome={outcome}, "
            f"success_rate={success_rate:.2f}, "
            f"failures={len(failure_reasons)}, "
            f"tips={len(optimization_tips)}"
        )

        # 8. 写入经验库（可选）
        if self._experience_store is not None:
            await self._write_experience(report, plan, plan_result, task_type, goal)

        return report

    def should_retry(self, check_result: CheckResult, retry_count: int) -> bool:
        """判断是否应该重试

        检查不通过且重试次数未耗尽时返回 True。

        Args:
            check_result: 检查结果
            retry_count: 当前重试次数

        Returns:
            是否应该重试
        """
        if check_result.status != CheckStatus.FAIL:
            return False
        if check_result.status == CheckStatus.SKIP:
            return False
        return retry_count < self._max_check_retries

    def should_request_human(self, check_result: CheckResult, retry_count: int) -> bool:
        """判断是否应该请求人工介入

        检查不通过且重试次数已耗尽时返回 True。

        Args:
            check_result: 检查结果
            retry_count: 当前重试次数

        Returns:
            是否应该请求人工介入
        """
        if check_result.status != CheckStatus.FAIL:
            return False
        return retry_count >= self._max_check_retries

    def make_step_complete_callback(
        self,
    ) -> Callable[[PlanStep, StepExecutionResult], Awaitable[None]]:
        """创建步骤完成回调，用于与 PlanExecutor 集成

        用法：
            checker = PlanChecker()
            executor = PlanExecutor(
                agent_pool=pool,
                on_step_complete=checker.make_step_complete_callback(),
            )

        Returns:
            异步回调函数
        """

        async def on_step_complete(
            step: PlanStep, exec_result: StepExecutionResult
        ) -> None:
            await self.check_step(step, exec_result)

        return on_step_complete

    def _collect_failure_reasons(self, plan_result: PlanExecutionResult) -> list[str]:
        """收集失败原因"""
        reasons: list[str] = []

        for sid, r in plan_result.step_results.items():
            if r.status == PlanStepStatus.FAILED:
                reason = f"Step '{sid}' failed"
                if r.error:
                    reason += f": {r.error}"
                reasons.append(reason)
            elif r.status == PlanStepStatus.SKIPPED:
                reason = f"Step '{sid}' skipped"
                if r.error:
                    reason += f": {r.error}"
                reasons.append(reason)

        # 补充检查不通过的原因
        for sid, cr in self._check_results.items():
            if cr.status == CheckStatus.FAIL:
                reason = f"Step '{sid}' quality check failed: {cr.reason}"
                if reason not in reasons:
                    reasons.append(reason)

        return reasons

    def _determine_outcome(self, plan_result: PlanExecutionResult) -> str:
        """判断整体结果"""
        total = len(plan_result.step_results)
        if total == 0:
            return "success"

        completed = len(plan_result.completed_steps)
        failed = len(plan_result.failed_steps)
        skipped = len(plan_result.skipped_steps)

        if completed == total:
            return "success"
        if failed == total or (failed + skipped == total and completed == 0):
            return "failure"
        return "partial"

    def _generate_optimization_tips(
        self,
        plan_result: PlanExecutionResult,
        quality_scores: dict[str, float],
    ) -> list[str]:
        """生成优化建议"""
        tips: list[str] = []

        # 基于质量评分
        low_quality_steps = [
            sid for sid, score in quality_scores.items() if score < self._quality_threshold
        ]
        if low_quality_steps:
            tips.append(
                f"Steps with low quality scores: {', '.join(low_quality_steps)}. "
                f"Consider improving input data or step configuration."
            )

        # 基于重试
        high_retry_steps = [
            (sid, r.retry_count)
            for sid, r in plan_result.step_results.items()
            if r.retry_count > 0
        ]
        if high_retry_steps:
            steps_str = ", ".join(
                f"'{sid}' ({count} retries)" for sid, count in high_retry_steps
            )
            tips.append(
                f"Steps requiring retries: {steps_str}. "
                f"Consider improving step reliability."
            )

        # 基于耗时
        slow_steps = [
            (sid, r.duration_ms)
            for sid, r in plan_result.step_results.items()
            if r.duration_ms > 60000
        ]
        if slow_steps:
            steps_str = ", ".join(
                f"'{sid}' ({ms / 1000:.1f}s)" for sid, ms in slow_steps
            )
            tips.append(
                f"Slow steps detected: {steps_str}. "
                f"Consider optimizing for performance."
            )

        # 基于跳过步骤
        skipped = plan_result.skipped_steps
        if skipped:
            tips.append(
                f"Skipped steps: {', '.join(skipped)}. "
                f"Review dependency chain and failure handling strategy."
            )

        # 基于检查结果中的 Reflector 建议
        for sid, cr in self._check_results.items():
            reflector_suggestions = cr.details.get("reflector_suggestions", [])
            for suggestion in reflector_suggestions:
                if suggestion not in tips:
                    tips.append(suggestion)

        return tips

    async def _write_experience(
        self,
        report: ReviewReport,
        plan: ExecutionPlan,
        plan_result: PlanExecutionResult,
        task_type: str,
        goal: str,
    ) -> None:
        """将复盘结果写入经验库"""
        from agentkit.evolution.experience_schema import TaskExperience

        # 构建步骤摘要
        steps_summary_parts: list[str] = []
        for step in plan.steps:
            r = plan_result.step_results.get(step.step_id)
            if r:
                steps_summary_parts.append(
                    f"{step.name}: {r.status.value}"
                    + (f" ({r.duration_ms / 1000:.1f}s)" if r.duration_ms > 0 else "")
                )
        steps_summary = "; ".join(steps_summary_parts)

        experience = TaskExperience(
            task_type=task_type or "plan_execution",
            goal=goal or plan.goal,
            steps_summary=steps_summary,
            outcome=report.outcome,
            duration_seconds=report.total_duration_ms / 1000,
            success_rate=report.success_rate,
            failure_reasons=report.failure_reasons,
            optimization_tips=report.optimization_tips,
        )

        try:
            exp_id = await self._experience_store.record_experience(experience)
            logger.info(f"Experience recorded: {exp_id} outcome={report.outcome}")
        except Exception as e:
            logger.error(f"Failed to write experience to store: {e}")

    def reset(self) -> None:
        """重置内部状态（用于新一轮检查）"""
        self._check_results.clear()
