"""AlignmentGuard - 对齐守卫：约束注入 + 级联故障检测"""

from __future__ import annotations

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

logger = logging.getLogger(__name__)


@dataclass
class AlignmentConfig:
    """对齐守卫配置"""

    constraints: list[str] = field(default_factory=list)
    cascade_max_interactions: int = 10
    cascade_max_depth: int = 3
    audit_enabled: bool = False
    audit_model: str = "default"


@dataclass
class AlignmentCheckResult:
    """对齐检查结果"""

    passed: bool
    violations: list[str] = field(default_factory=list)
    checked_by: str = ""  # "rule" or "llm"


@dataclass
class CascadeAlert:
    """级联故障告警"""

    session_id: str
    alert_type: str  # "interaction_limit" or "loop_depth"
    current_value: int
    threshold: int
    message: str


class ConstraintInjector:
    """将全局约束注入到任务 input_data 中"""

    def __init__(self, config: AlignmentConfig):
        self._config = config

    def inject(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """注入约束指令到 input_data

        在 input_data 中添加 'alignment_constraints' 键，值为约束列表。
        不修改原始 dict，返回新 dict。
        """
        result = {**input_data, "alignment_constraints": list(self._config.constraints)}
        return result


class AlignmentGuard:
    """对齐守卫 — 扩展 QualityGate，增加约束注入和级联检测"""

    def __init__(self, config: AlignmentConfig, llm_gateway=None):
        self._config = config
        self._injector = ConstraintInjector(config)
        self._llm_gateway = llm_gateway
        self._interaction_counts: dict[str, int] = {}
        self._loop_depths: dict[str, int] = {}

    def inject_constraints(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """委托给 ConstraintInjector"""
        return self._injector.inject(input_data)

    async def check_output(
        self,
        output: dict[str, Any],
        constraints: list[str] | None = None,
    ) -> AlignmentCheckResult:
        """检查输出是否符合约束

        - 系统级约束：基于规则的检查（关键词 + 正则匹配）
        - 组织级约束：LLM 语义检查（仅当 audit_enabled=True）
        """
        effective_constraints = constraints if constraints is not None else self._config.constraints
        if not effective_constraints:
            return AlignmentCheckResult(passed=True, checked_by="rule")

        # 1. 基于规则的检查：关键词/子串匹配
        violations = self._rule_check(output, effective_constraints)
        if violations:
            return AlignmentCheckResult(
                passed=False,
                violations=violations,
                checked_by="rule",
            )

        # 2. LLM 语义检查（仅当 audit_enabled=True 且有 llm_gateway）
        if self._config.audit_enabled and self._llm_gateway is not None:
            return await self._llm_check(output, effective_constraints)

        return AlignmentCheckResult(passed=True, checked_by="rule")

    def _rule_check(
        self, output: dict[str, Any], constraints: list[str]
    ) -> list[str]:
        """基于规则的约束检查：将 output 内容拼接后做关键词/子串匹配"""
        content = self._extract_text(output)
        violations: list[str] = []
        for constraint in constraints:
            # 简单子串匹配：约束关键词出现在输出中即视为违规
            if constraint.lower() in content.lower():
                violations.append(constraint)
        return violations

    @staticmethod
    def _extract_text(output: dict[str, Any]) -> str:
        """从 output dict 中提取所有文本内容"""
        parts: list[str] = []
        for value in output.values():
            if isinstance(value, str):
                parts.append(value)
            else:
                parts.append(str(value))
        return " ".join(parts)

    async def _llm_check(
        self, output: dict[str, Any], constraints: list[str]
    ) -> AlignmentCheckResult:
        """LLM 语义检查"""
        content = self._extract_text(output)
        constraint_text = "\n".join(f"- {c}" for c in constraints)
        messages = [
            {
                "role": "system",
                "content": (
                    "You are an alignment auditor. Check if the following output "
                    "violates any of the listed constraints. "
                    "Reply with 'PASS' if no violations, or list the violated constraints."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Constraints:\n{constraint_text}\n\nOutput:\n{content}"
                ),
            },
        ]
        try:
            response = await self._llm_gateway.chat(
                messages=messages, model=self._config.audit_model
            )
            reply = response.content.strip()
            if reply.upper().startswith("PASS"):
                return AlignmentCheckResult(passed=True, checked_by="llm")
            else:
                return AlignmentCheckResult(
                    passed=False,
                    violations=[reply],
                    checked_by="llm",
                )
        except Exception as e:
            logger.warning(f"LLM audit failed: {e}")
            return AlignmentCheckResult(
                passed=False,
                violations=[f"LLM audit unavailable: {e}"],
                checked_by="rule",
            )

    def record_interaction(self, session_id: str) -> CascadeAlert | None:
        """记录一次 agent 间交互，超过阈值返回 CascadeAlert"""
        self._interaction_counts[session_id] = (
            self._interaction_counts.get(session_id, 0) + 1
        )
        count = self._interaction_counts[session_id]
        if count > self._config.cascade_max_interactions:
            return CascadeAlert(
                session_id=session_id,
                alert_type="interaction_limit",
                current_value=count,
                threshold=self._config.cascade_max_interactions,
                message=(
                    f"Session {session_id} exceeded max interactions: "
                    f"{count} > {self._config.cascade_max_interactions}"
                ),
            )
        return None

    def record_loop_depth(self, session_id: str, depth: int) -> CascadeAlert | None:
        """记录循环深度，超过阈值返回 CascadeAlert"""
        self._loop_depths[session_id] = depth
        if depth > self._config.cascade_max_depth:
            return CascadeAlert(
                session_id=session_id,
                alert_type="loop_depth",
                current_value=depth,
                threshold=self._config.cascade_max_depth,
                message=(
                    f"Session {session_id} exceeded max loop depth: "
                    f"{depth} > {self._config.cascade_max_depth}"
                ),
            )
        return None

    def reset_session(self, session_id: str) -> None:
        """重置某个 session 的交互计数"""
        self._interaction_counts.pop(session_id, None)
        self._loop_depths.pop(session_id, None)

    def get_interaction_count(self, session_id: str) -> int:
        """获取某个 session 的当前交互计数"""
        return self._interaction_counts.get(session_id, 0)