"""
Context 抽象基类

定义上下文管理的核心接口。
Agent 只需关心公开方法，压缩逻辑由 Context 内部自动处理。
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import List, Optional, Callable, AsyncGenerator, Awaitable, TYPE_CHECKING

if TYPE_CHECKING:
    from provider import Message, TokenUsage

# chat_fn 类型：调用 LLM 的异步生成器工厂
ChatFn = Callable[..., AsyncGenerator]

# count_tokens_fn 类型：估算 token 数的异步回调
CountTokensFn = Callable[..., Awaitable[int]]


@dataclass
class ContextConfig:
    """上下文配置"""

    # Token 限制（由 BaseModelClient.get_context_window() 动态获取，None 时使用默认值）
    max_tokens: Optional[int] = None

    # ==================== 触发阈值 ====================
    # 触发压缩的阈值
    threshold: float = 0.9  # 触发压缩

    # ==================== 弱压缩 (Prune) ====================
    prune_keep_recent: int = 10  # 保留最近 N 条消息
    prune_protect_chars: int = 10000  # 保护最近的工具输出字符数
    prune_min_chars: int = 5000  # 最小清理量
    prune_threshold: float = 0.6  # 弱压缩后目标阈值（60%）

    # ==================== 中压缩 (Summarize-Weak) ====================
    summarize_weak_turns: int = 3  # 保留最近 N 轮对话
    summarize_weak_min_msgs: int = 10  # 最小保留消息数
    summarize_weak_threshold: float = 0.6  # 中压缩后目标阈值（50%）

    # ==================== 强压缩 (Summarize-Strong) ====================
    summarize_strong_turns: int = 1  # 保留最近 N 轮对话
    summarize_strong_min_msgs: int = 4  # 最小保留消息数
    summarize_strong_threshold: float = 0.6  # 强压缩后目标阈值（40%）

    # ==================== 最强压缩 (Summarize-All) ====================
    # 最强压缩没有阈值，它是最终的兜底策略

    # Token 估算参数
    chars_per_token: int = 4


class BaseContext(ABC):
    """
    上下文基类 - 智能容器

    Agent 只需调用以下公开方法：
    - add_message(): 添加消息
    - get_messages(): 获取有效消息
    - update_token_usage(): 更新 token 使用量
    - set_system_prompt(): 设置系统提示词
    - clear(): 清空上下文

    压缩、溢出检测等逻辑完全内部自动处理。
    """

    def __init__(
        self,
        chat_fn: ChatFn,
        config: Optional[ContextConfig] = None,
    ):
        """
        初始化上下文

        Args:
            chat_fn: LLM 调用函数（异步生成器工厂，由 Agent 传入）
            config: 上下文配置
        """
        self.chat_fn = chat_fn
        self.config = config or ContextConfig()

    # ==================== 公开方法（Agent 调用） ====================

    @abstractmethod
    def add_message(self, message: "Message") -> None:
        """
        添加消息

        内部自动检查溢出并触发压缩。

        Args:
            message: 消息对象
        """
        pass

    @abstractmethod
    def get_messages(self) -> List["Message"]:
        """
        获取对话消息列表（不含系统提示词）

        Returns:
            消息列表
        """
        pass

    @abstractmethod
    def update_token_usage(self, usage: "TokenUsage") -> None:
        """
        更新 token 使用量

        使用 LLM 返回的实际 token 数，比字符估算更准确。

        Args:
            usage: LLM 返回的 token 使用量
        """
        pass

    @abstractmethod
    def set_system_prompt(self, prompts: List["Message"]) -> None:
        """
        设置系统提示词

        Args:
            prompts: 系统提示词 Message 列表（role 为 PROVIDER/AGENT/ENVIRONMENT/CUSTOM）
        """
        pass

    @property
    @abstractmethod
    def system_prompts(self) -> List["Message"]:
        """获取系统提示词列表"""
        pass

    @abstractmethod
    def clear(self) -> None:
        """清空上下文（保留系统提示词）"""
        pass

    # ==================== 可选覆盖 ====================

    def get_token_count(self) -> int:
        """
        获取当前 token 使用量

        Returns:
            token 数量
        """
        return 0

    def is_overflow(self) -> bool:
        """
        检查是否溢出（供外部查询，不触发压缩）

        Returns:
            是否溢出
        """
        tokens = self.get_token_count()
        return tokens > self.config.max_tokens * self.config.threshold

    async def compress_if_needed(self, force: bool = False) -> bool:
        """
        如果需要压缩则执行

        Args:
            force: 是否强制执行压缩（忽略阈值检查）

        由 agent 主循环在调用 LLM 之前显式调用。
        基类默认不做任何事，子类可覆盖实现压缩逻辑。

        Returns:
            是否执行了压缩
        """
        return False

    @property
    def needs_compress(self) -> bool:
        """
        是否需要压缩

        基类默认返回 False，子类可覆盖。
        """
        return False

    def get_context_info(self) -> dict:
        """
        获取上下文信息（供外部查询状态）

        Returns:
            包含 token 使用量、上下文窗口大小、使用百分比等信息的字典
        """
        token_count = self.get_token_count()
        max_tokens = self.config.max_tokens or 0
        percent = (token_count / max_tokens * 100) if max_tokens > 0 else 0
        
        return {
            "token_count": token_count,
            "max_tokens": max_tokens,
            "percent": round(percent, 1),
            "threshold": self.config.threshold * 100,
        }
