import asyncio
import random
import time
from typing import Any, Callable, TypeVar

T = TypeVar("T")

class RetryableError(Exception):
    """瞬态故障，可重试"""
    
class NonRetryableError(Exception):
    """客户端故障，不可重试"""

async def retry_with_backoff(
    fn: Callable[..., T],
    max_retries: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    jitter_factor: float = 0.1,
    retryable_statuses: set[int] | None = None,
) -> T:
    """Exponential Backoff + Jitter，带错误分类"""
    retryable_statuses = retryable_statuses or {429, 502, 503, 504}
    last_exception = None
    
    for attempt in range(max_retries + 1):
        try:
            return await fn()
        except NonRetryableError:
            raise  # 不重试客户端错误
        except RetryableError as e:
            last_exception = e
        except Exception as e:
            status = _extract_status(e)
            if status in retryable_statuses or _is_transient(e):
                last_exception = RetryableError(str(e))
            else:
                raise  # 不认识这个错误，不重试
        
        if attempt < max_retries:
            delay = _compute_backoff(attempt, base_delay, max_delay, jitter_factor)
            print(f"  Retry {attempt + 1}/{max_retries} in {delay:.1f}s...")
            await asyncio.sleep(delay)
    
    raise RetryableError(
        f"All {max_retries} retries failed. Last error: {last_exception}"
    ) from last_exception


def _compute_backoff(attempt: int, base: float, max_delay: float, jitter: float) -> float:
    """指数退避 + 随机 jitter"""
    delay = min(base * (2 ** attempt), max_delay)
    noise = delay * jitter * (2 * random.random() - 1)
    return max(0.1, delay + noise)

class RetryBudget:
    """重试预算：限制总重试时间和次数"""
    
    def __init__(self, max_attempts: int = 3, max_duration_ms: int = 30_000):
        self.max_attempts = max_attempts
        self.max_duration_ms = max_duration_ms
        self.attempts = 0
        self.started_at: float | None = None
    
    def can_retry(self) -> bool:
        if self.started_at is None:
            self.started_at = time.time()
        elapsed = (time.time() - self.started_at) * 1000
        if self.attempts >= self.max_attempts:
            return False
        if elapsed >= self.max_duration_ms:
            return False
        return True
    
    def record_attempt(self):
        self.attempts += 1

from pydantic import BaseModel, Field, ValidationError
from typing_extensions import Literal

class WeatherToolInput(BaseModel):
    city: str = Field(..., min_length=1, description="城市名")
    country: str = Field(..., min_length=1, description="国家代码")
    units: Literal["celsius", "fahrenheit"] = Field(default="celsius")
    timezone: str | None = Field(None, pattern=r"^[+-]\d{2}:\d{2}$")

class ToolOutputValidator:
    """LLM Tool Output 校验 + 自动修复"""
    
    def __init__(self, model_class: type[BaseModel]):
        self.model_class = model_class
    
    def validate_and_fix(self, raw_output: dict) -> BaseModel:
        """校验输出，尝试自动修复，最后 fallback"""
        try:
            return self.model_class(**raw_output)
        except ValidationError as e:
            return self._attempt_fix(raw_output, e)
    
    def _attempt_fix(self, raw: dict, error: ValidationError) -> BaseModel:
        """尝试修复常见错误类型"""
        fixed = dict(raw)
        
        for err in error.errors():
            loc = err["loc"][0]  # 出错的字段名
            msg = err["msg"]
            _type = err["type"]
            
            # 类型转换修复
            if _type == "type_error" and loc in fixed:
                if msg.startswith("Input should be a valid integer"):
                    fixed[loc] = int(fixed[loc])
                elif msg.startswith("Input should be a valid float"):
                    fixed[loc] = float(fixed[loc])
                elif msg.startswith("Input should be a valid boolean"):
                    fixed[loc] = str(fixed[loc]).lower() in ("true", "1", "yes")
            
            # Enum 修正：尝试大小写不敏感匹配
            if _type == "enum":
                valid_values = _get_enum_values(self.model_class, loc)
                if loc in fixed and isinstance(fixed[loc], str):
                    for v in valid_values:
                        if fixed[loc].lower() == v.lower():
                            fixed[loc] = v
                            break
        
        # 二次校验
        try:
            return self.model_class(**fixed)
        except ValidationError:
            raise NonRetryableError(f"Tool output validation failed after fixes: {fixed}")

class ValidatedToolCaller:
    """带 Schema 校验的 Tool 调用器"""
    
    def __init__(self, validator: ToolOutputValidator, llm_client):
        self.validator = validator
        self.llm = llm_client
        self.max_retries = 2
    
    async def call_with_validation(self, tool_name: str, llm_raw_response: dict) -> BaseModel:
        """调用 Tool + 校验 + 有条件重试"""
        for attempt in range(self.max_retries + 1):
            try:
                validated = self.validator.validate_and_fix(llm_raw_response)
                # 执行实际的 Tool 调用
                result = await execute_tool(tool_name, validated.model_dump())
                return result
            except NonRetryableError:
                raise  # 校验失败且无法修复
            except RetryableError:
                # 重试前优化 prompt
                llm_raw_response = await self.llm.resend_with_correction(
                    previous_error=str(e),
                    stricter_schema=True,
                )
        
        # 最终 fallback
        return self._fallback_response(tool_name)

class MCPConnectionPool:
    """MCP 连接池：自动重连 + 健康检查 + 优雅降级"""
    
    def __init__(self, server_configs: list[dict], heartbeat_interval: int = 30):
        self.connections: dict[str, MCPClient] = {}
        self.health: dict[str, bool] = {}
        self.configs = {cfg["name"]: cfg for cfg in server_configs}
        self._heartbeat_task = None
    
    async def connect(self, server_name: str) -> MCPClient:
        """建立连接，含重试"""
        cfg = self.configs[server_name]
        for attempt in range(3):
            try:
                client = MCPClient(
                    command=cfg["command"],
                    args=cfg.get("args", []),
                    env=cfg.get("env", {}),
                )
                await client.initialize()
                self.connections[server_name] = client
                self.health[server_name] = True
                return client
            except Exception as e:
                print(f"  MCP connect {server_name} attempt {attempt+1} failed: {e}")
                await asyncio.sleep(1 * (2 ** attempt))
        raise ConnectionError(f"MCP server {server_name} unreachable after 3 attempts")
    
    async def call_tool(self, server_name: str, tool_name: str, args: dict):
        """调用 Tool，带健康检查和故障处理"""
        # 检查健康状况
        if not self.health.get(server_name, True):
            raise RetryableError(f"MCP server {server_name} is marked unhealthy")
        
        try:
            client = self.connections.get(server_name)
            if client is None:
                client = await self.connect(server_name)
            
            result = await client.call_tool(tool_name, args)
            return result
            
        except (ConnectionError, WebSocketDisconnect) as e:
            # MCP 连接中断——记录并尝试重连
            self.health[server_name] = False
            raise RetryableError(f"MCP connection lost: {e}")
        
        except ToolExecutionError as e:
            # Tool 执行逻辑错误——不是连接问题，不标记 Server 不可用
            if e.is_transient:
                raise RetryableError(str(e))
            else:
                raise NonRetryableError(str(e))
    
    async def health_check_loop(self):
        """后台健康检查"""
        while True:
            for name in list(self.connections.keys()):
                client = self.connections[name]
                try:
                    # Ping-like 健康检查
                    await client.list_tools(timeout=5)
                    self.health[name] = True
                except:
                    self.health[name] = False
                    # 尝试后台重连
                    try:
                        self.connections[name] = await self.connect(name)
                    except:
                        pass
            await asyncio.sleep(30)

class MCPGracefulDegradation:
    """MCP 降级策略管理器"""
    
    # 每个 Tool 的降级配置
    FALLBACK_MAP = {
        "search_web": {"alternative": None, "critical": False},  # 非关键——跳过即可
        "send_email": {"alternative": "send_email_smtp", "critical": True},  # 有关键替代
        "read_database": {"alternative": None, "critical": True},  # 如果不可用，整个任务失败
    }
    
    @classmethod
    async def call_with_degradation(cls, tool_name: str, args: dict, pool: MCPConnectionPool):
        """带降级的 Tool 调用"""
        server_name = cls._find_server(tool_name)
        
        try:
            return await pool.call_tool(server_name, tool_name, args)
        except (RetryableError, ConnectionError) as e:
            # 尝试降级
            fallback = cls.FALLBACK_MAP.get(tool_name, {})
            alt_tool = fallback.get("alternative")
            
            if alt_tool:
                print(f"  Degrading {tool_name} → {alt_tool}")
                alt_server = cls._find_server(alt_tool)
                return await pool.call_tool(alt_server, alt_tool, args)
            
            if fallback.get("critical", False):
                raise NonRetryableError(
                    f"Critical tool {tool_name} unavailable, cannot continue task"
                )
            
            # 非关键工具——优雅跳过
            print(f"  Skipping non-critical tool {tool_name} (unavailable)")
            return {"status": "skipped", "reason": f"Tool {tool_name} unavailable"}

import time
from collections import deque
from enum import Enum

class CircuitState(Enum):
    CLOSED = "closed"      # 正常
    OPEN = "open"           # 跳闸
    HALF_OPEN = "half_open" # 试探中

class PluginCircuitBreaker:
    """Agent Plugin 熔断器——参考 Hystrix 设计"""
    
    def __init__(
        self,
        plugin_name: str,
        failure_threshold: int = 5,       # 5次连续失败触发熔断
        recovery_timeout: float = 30.0,    # 30秒冷却
        half_open_max_calls: int = 1,      # 半开状态试探请求数
        rolling_window: int = 60,          # 滑动窗口：60秒
    ):
        self.plugin_name = plugin_name
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.half_open_max_calls = half_open_max_calls
        self.rolling_window = rolling_window
        
        self.state = CircuitState.CLOSED
        self.failure_count = 0
        self.last_failure_time = 0.0
        self.half_open_calls = 0
        
        # 滑动窗口记录
        self._recent_results: deque[bool] = deque(maxlen=100)
    
    async def call(self, fn, *args, **kwargs):
        """执行函数，受熔断器保护"""
        if self.state == CircuitState.OPEN:
            if time.time() - self.last_failure_time > self.recovery_timeout:
                # 冷却时间到，进入 HALF-OPEN
                self.state = CircuitState.HALF_OPEN
                self.half_open_calls = 0
            else:
                raise RetryableError(
                    f"Circuit breaker OPEN for plugin '{self.plugin_name}'. "
                    f"Retry after {self.recovery_timeout - (time.time() - self.last_failure_time):.0f}s"
                )
        
        if self.state == CircuitState.HALF_OPEN:
            if self.half_open_calls >= self.half_open_max_calls:
                raise RetryableError(
                    f"Circuit breaker HALF-OPEN for '{self.plugin_name}', "
                    f"probe limit reached. Wait for next cycle."
                )
            self.half_open_calls += 1
        
        try:
            result = await fn(*args, **kwargs)
            self._on_success()
            return result
        except Exception as e:
            self._on_failure()
            raise
    
    def _on_success(self):
        self._recent_results.append(True)
        if self.state == CircuitState.HALF_OPEN:
            self.state = CircuitState.CLOSED  # 试探成功，恢复
            self.failure_count = 0
    
    def _on_failure(self):
        self._recent_results.append(False)
        self.failure_count += 1
        self.last_failure_time = time.time()
        
        if self.state == CircuitState.HALF_OPEN:
            self.state = CircuitState.OPEN  # 试探失败，回跳闸
        elif self.failure_count >= self.failure_threshold:
            self.state = CircuitState.OPEN

class PluginBulkhead:
    """舱壁模式：限制每个 Plugin 的并发数"""
    
    def __init__(self, max_concurrent: int = 3, queue_size: int = 10):
        self.semaphore = asyncio.Semaphore(max_concurrent)
        self.queue = asyncio.Queue(maxsize=queue_size)
    
    async def execute(self, fn, *args, **kwargs):
        async with self.semaphore:
            return await fn(*args, **kwargs)

class IdempotencyManager:
    """幂等性管理器——确保同一个 Tool 调用不被重复执行"""
    
    def __init__(self, storage: dict | None = None):
        self._store: dict[str, Any] = storage or {}
    
    def generate_key(self, tool_name: str, args: dict, session_id: str) -> str:
        """生成幂等性 Key：基于 Tool 名 + 参数 + 会话 ID 的确定哈希"""
        # 规范化参数：移除 LLM 可能引入的无关字段
        canonical_args = self._canonicalize(args)
        
        # 排序后的 JSON 确保相同参数产生相同 key
        payload = f"{session_id}:{tool_name}:{json.dumps(canonical_args, sort_keys=True)}"
        return hashlib.sha256(payload.encode()).hexdigest()[:32]
    
    def _canonicalize(self, args: dict) -> dict:
        """规范化参数——移除 LLM 可能会加的噪声字段"""
        NOISE_FIELDS = {"note", "reasoning", "thought", "explanation", "_meta"}
        return {k: v for k, v in args.items() if k not in NOISE_FIELDS}
    
    async def execute_idempotent(
        self,
        tool_name: str,
        args: dict,
        session_id: str,
        execute_fn,
    ):
        """幂等执行：如果已执行过，返回缓存结果"""
        key = self.generate_key(tool_name, args, session_id)
        
        if key in self._store:
            print(f"  Idempotent hit: {tool_name} already executed (key={key[:8]})")
            return self._store[key]
        
        result = await execute_fn(tool_name, args)
        self._store[key] = result
        return result

class PaymentToolHandler:
    """支付类 Tool 的幂等策略：先确认，再执行"""
    
    async def process_refund(self, refund_id: str, amount: float, session_id: str):
        """带确认的幂等退款"""
        idem_key = self.idempotency.generate_key(
            "process_refund", {"refund_id": refund_id, "amount": amount}, session_id
        )
        
        # 第一步：记录意图
        self.db.save_intent(
            key=idem_key,
            action="refund",
            params={"refund_id": refund_id, "amount": amount},
            status="pending"
        )
        
        # 第二步：先确认状态（防止上一次超时后实际已执行）
        existing = self.payment_provider.query_refund(refund_id)
        if existing and existing["status"] == "completed":
            return {"status": "already_processed", "refund_id": refund_id}
        
        # 第三步：执行（带幂等 key）
        result = await self.payment_provider.create_refund(
            idempotency_key=idem_key,
            refund_id=refund_id,
            amount=amount,
        )
        
        # 第四步：更新状态
        self.db.update_intent_status(idem_key, "completed")
        return result

class ReliableToolCaller:
    """五层容错集成：最外层的统一 Tool 调用入口"""
    
    def __init__(self, mcp_pool, llm_client, session_id: str):
        self.validator_map: dict[str, ToolOutputValidator] = {}
        self.idempotency = IdempotencyManager()
        self.circuit_breakers: dict[str, PluginCircuitBreaker] = {}
        self.bulkheads: dict[str, PluginBulkhead] = {}
        self.mcp = mcp_pool
        self.llm = llm_client
        self.session_id = session_id
    
    def register_plugin(self, name: str, validator, breaker, bulkhead):
        self.validator_map[name] = validator
        self.circuit_breakers[name] = breaker
        self.bulkheads[name] = bulkhead
    
    async def call_tool(self, plugin_name: str, tool_name: str, raw_args: dict) -> Any:
        """完整容错管线"""
        # L1: Schema 校验
        validator = self.validator_map.get(plugin_name)
        if validator:
            try:
                validated = validator.validate_and_fix(raw_args)
            except NonRetryableError as e:
                # 优化 prompt 后重试一次
                corrected = await self.llm.resend_with_correction(
                    tool_name, raw_args, str(e)
                )
                validated = validator.validate_and_fix(corrected)
            args = validated.model_dump()
        else:
            args = raw_args
        
        # L2: 幂等性检查
        async def execute():
            return await self._call_with_breaker(plugin_name, tool_name, args)
        
        return await self.idempotency.execute_idempotent(
            tool_name, args, self.session_id, execute
        )
    
    async def _call_with_breaker(self, plugin: str, tool: str, args: dict):
        """L3+L4+L5: 熔断器 → 舱壁 → 重试"""
        breaker = self.circuit_breakers.get(plugin)
        bulkhead = self.bulkheads.get(plugin)
        
        async def call_internal():
            return await retry_with_backoff(
                lambda: self.mcp.call_tool(plugin, tool, args),
                max_retries=2,
                base_delay=1.0,
            )
        
        if breaker and bulkhead:
            # 熔断器保护 → 舱壁隔离
            return await breaker.call(
                lambda: bulkhead.execute(call_internal)
            )
        elif breaker:
            return await breaker.call(call_internal)
        else:
            return await call_internal()

# 每个熔断器跳闸时记录事件
class ObservableCircuitBreaker(PluginCircuitBreaker):
    def _on_failure(self):
        super()._on_failure()
        if self.state == CircuitState.OPEN:
            logger.warning(
                "Circuit breaker OPEN",
                extra={
                    "plugin": self.plugin_name,
                    "failure_count": self.failure_count,
                    "recovery_timeout": self.recovery_timeout,
                }
            )
            metrics.increment("circuit_breaker.tripped", tags={"plugin": self.plugin_name})