# utils/context_manager.py - 上下文管理器（集成对话持久化和Token统计）

import os
import json
import base64
import mimetypes
import io
import uuid
import platform
import shutil
import subprocess
from copy import deepcopy
from typing import Dict, List, Optional, Any
from pathlib import Path
from datetime import datetime
try:
    from config import (
        MAX_CONTEXT_SIZE,
        DATA_DIR,
        PROMPTS_DIR,
        TERMINAL_SANDBOX_MOUNT_PATH,
        TERMINAL_SANDBOX_CPUS,
        TERMINAL_SANDBOX_MEMORY,
        PROJECT_MAX_STORAGE_MB,
        TERMINAL_SANDBOX_MODE,
        LINUX_SAFETY,
    )
    from config.model_profiles import (
        get_model_prompt_replacements,
        get_registered_model_keys,
        model_supports_image,
        model_supports_video,
    )
except ImportError:
    import sys
    from pathlib import Path
    project_root = Path(__file__).resolve().parents[1]
    if str(project_root) not in sys.path:
        sys.path.insert(0, str(project_root))
    from config import (
        MAX_CONTEXT_SIZE,
        DATA_DIR,
        PROMPTS_DIR,
        TERMINAL_SANDBOX_MOUNT_PATH,
        TERMINAL_SANDBOX_CPUS,
        TERMINAL_SANDBOX_MEMORY,
        PROJECT_MAX_STORAGE_MB,
        TERMINAL_SANDBOX_MODE,
        LINUX_SAFETY,
    )
    from config.model_profiles import (
        get_model_prompt_replacements,
        get_registered_model_keys,
        model_supports_image,
        model_supports_video,
    )
from utils.conversation_manager import ConversationManager
from utils.host_workspace_debug import write_host_workspace_debug
from utils.media_store import MediaStore
from utils.token_usage import normalize_usage_payload

AUTO_SHALLOW_PLACEHOLDER = "过早的工具结果已经被自动压缩"
AUTO_SHALLOW_TOOL_WHITELIST = {
    "write_file",
    "read_file",
    "edit_file",
    "terminal_input",
    "terminal_snapshot",
    "web_search",
    "extract_webpage",
    "run_command",
    "view_image",
    "view_video",
}


class TokenMixin:
    """ContextManager token mixin 能力 mixin。"""

    def _token_totals_path(self) -> Path:
        return self.data_dir / "token_totals.json"

    def _load_token_totals(self) -> Dict[str, Any]:
        path = self._token_totals_path()
        if not path.exists():
            return {
                "input_tokens": 0,
                "output_tokens": 0,
                "total_tokens": 0,
                "updated_at": None,
            }
        try:
            with open(path, 'r', encoding='utf-8') as fh:
                payload = json.load(fh) or {}
            return {
                "input_tokens": int(payload.get("input_tokens") or payload.get("total_input_tokens") or 0),
                "output_tokens": int(payload.get("output_tokens") or payload.get("total_output_tokens") or 0),
                "total_tokens": int(payload.get("total_tokens") or 0),
                "updated_at": payload.get("updated_at"),
            }
        except (OSError, json.JSONDecodeError, ValueError) as exc:
            print(f"[TokenStats] 读取累计Token失败: {exc}")
            return {
                "input_tokens": 0,
                "output_tokens": 0,
                "total_tokens": 0,
                "updated_at": None,
            }

    def _save_token_totals(self, data: Dict[str, Any]):
        path = self._token_totals_path()
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w', encoding='utf-8') as fh:
            json.dump(data, fh, ensure_ascii=False, indent=2)

    def _increment_workspace_token_totals(self, input_tokens: int, output_tokens: int, total_tokens: int):
        if input_tokens <= 0 and output_tokens <= 0 and total_tokens <= 0:
            return
        snapshot = self._load_token_totals()
        snapshot["input_tokens"] = snapshot.get("input_tokens", 0) + max(0, int(input_tokens))
        snapshot["output_tokens"] = snapshot.get("output_tokens", 0) + max(0, int(output_tokens))
        snapshot["total_tokens"] = snapshot.get("total_tokens", 0) + max(0, int(total_tokens))
        snapshot["updated_at"] = datetime.now().isoformat()
        self._save_token_totals(snapshot)

    def apply_usage_statistics(self, usage: Dict[str, Any]) -> bool:
        """
        根据模型返回的 usage 字段更新token统计
        """
        normalized_usage = normalize_usage_payload(usage) or {}
        prompt_tokens = int(normalized_usage.get("prompt_tokens") or 0)
        completion_tokens = int(normalized_usage.get("completion_tokens") or 0)
        total_tokens = int(normalized_usage.get("total_tokens") or (prompt_tokens + completion_tokens))
        # 当前上下文长度优先取专用字段；缺失时回退到 prompt_tokens
        current_context_tokens = int(normalized_usage.get("current_context_tokens") or prompt_tokens)

        try:
            self._increment_workspace_token_totals(prompt_tokens, completion_tokens, total_tokens)
        except Exception as exc:
            print(f"[TokenStats] 无法写入累计Token: {exc}")

        if not self.current_conversation_id:
            print("⚠️ 没有当前对话ID，跳过usage统计更新")
            return False
        
        try:
            success = self.conversation_manager.update_token_statistics(
                self.current_conversation_id,
                prompt_tokens,
                completion_tokens,
                total_tokens,
                current_context_tokens=current_context_tokens,
            )
            
            if success:
                self.safe_broadcast_token_update()
            
            return success
        except Exception as e:
            print(f"更新usage统计失败: {e}")
            return False

    def get_conversation_token_statistics(self, conversation_id: str = None) -> Optional[Dict]:
        """
        获取指定对话的token统计
        
        Args:
            conversation_id: 对话ID，默认为当前对话
        
        Returns:
            Dict: Token统计信息
        """
        target_id = conversation_id or self.current_conversation_id
        if not target_id:
            return None

        # 路由到正确的 conversation_manager（支持多智能体会话独立存储）
        target_manager = getattr(
            self, "_get_conversation_manager_for_id", lambda _: self.conversation_manager
        )(target_id)
        return target_manager.get_token_statistics(target_id)

    def get_current_context_tokens(self, conversation_id: str = None) -> int:
        """
        获取最近一次请求的上下文token数量
        """
        stats = self.get_conversation_token_statistics(conversation_id)
        if not stats:
            return 0
        return stats.get("current_context_tokens", 0)

    def safe_broadcast_token_update(self):
        """安全的token更新广播（只广播累计统计，不重新计算）"""
        try:
            print(f"[Debug] 尝试广播token更新")
            
            # 检查是否有回调函数
            if not hasattr(self, '_web_terminal_callback'):
                print(f"[Debug] 没有_web_terminal_callback属性")
                return
                
            if not self._web_terminal_callback:
                print(f"[Debug] _web_terminal_callback为None")
                return
            
            if not self.current_conversation_id:
                print(f"[Debug] 没有当前对话ID")
                return
                
            print(f"[Debug] 广播token统计，对话ID: {self.current_conversation_id}")
            
            # 只获取已有的累计token统计，不重新计算
            cumulative_stats = self.get_conversation_token_statistics()
            
            # 准备广播数据
            broadcast_data = {
                'conversation_id': self.current_conversation_id,
                'cumulative_input_tokens': cumulative_stats.get("total_input_tokens", 0) if cumulative_stats else 0,
                'cumulative_output_tokens': cumulative_stats.get("total_output_tokens", 0) if cumulative_stats else 0,
                'cumulative_total_tokens': cumulative_stats.get("total_tokens", 0) if cumulative_stats else 0,
                'current_context_tokens': cumulative_stats.get("current_context_tokens", 0) if cumulative_stats else 0,
                'updated_at': datetime.now().isoformat()
            }
            
            print(f"[Debug] Token统计: 累计输入={broadcast_data['cumulative_input_tokens']}, 累计输出={broadcast_data['cumulative_output_tokens']}")
            
            # 广播到前端
            self._web_terminal_callback('token_update', broadcast_data)
            
            print(f"[Debug] token更新已广播")
            
        except Exception as e:
            print(f"[Debug] 广播token更新失败: {e}")
            import traceback
            traceback.print_exc()
