"""
Session - 会话管理

管理 WebSocket 连接对应的会话，每个会话包含：
- Agent 实例（直接持有，不再通过 AgentSession 间接层）
- 工具注册表
- 权限管理
- 任务生命周期
"""

import asyncio
import uuid
from typing import Dict, Any, Optional, Callable, Awaitable, List, AsyncGenerator

from provider import ModelClientFactory, BaseModelClient
from core.logger import log_event
from core.config import AppConfig
from core.permission import PermissionRequest, PermissionResponse, PermissionManager
from core.session import get_session_manager as get_persistent_manager
from core.interaction import get_interaction_manager
from tools.registry import ToolRegistry
from tools.base import BaseTool
from agents import BaseAgent, create_agent, list_agent_types
from dataclasses import dataclass

from .constants import AgentType, EventType, PermissionConfig


# 权限请求回调类型
PermissionCallback = Callable[[str, str, Dict[str, Any]], Awaitable[PermissionResponse]]


@dataclass
class SessionState:
    """会话状态"""
    session_id: str
    agent_type: AgentType
    is_connected: bool = True


class Session:
    """
    会话

    管理单个 WebSocket 连接的完整生命周期，
    直接持有 Agent 实例并管理任务执行。
    """

    def __init__(
        self,
        session_id: str,
        model_name: str,
        config: AppConfig,
        mcp_tools_getter: Optional[Callable[[], List[BaseTool]]] = None,
    ):
        self.session_id = session_id
        self.model_name = model_name
        self._config = config

        # 存储 mcp_tools_getter 供子 Agent 使用
        self._mcp_tools_getter = mcp_tools_getter

        # 工具注册表（直接注入 mcp_tools_getter）
        self.tool_registry = ToolRegistry(
            session_id=session_id,
            interaction_manager=get_interaction_manager(),
            mcp_tools_getter=mcp_tools_getter,
        )

        # Agent（直接持有）
        self.agent: Optional[BaseAgent] = None
        self.agent_type: AgentType = AgentType.GENERAL

        # 任务状态（原 AgentSession 字段）
        self._current_task_id: Optional[str] = None
        self._is_running: bool = False
        self._cancel_requested: bool = False
        self._cancel_event: Optional[asyncio.Event] = None

        # 权限相关
        self._permission_future: Optional[asyncio.Future] = None
        self._permission_request_id: Optional[str] = None

        # 事件回调
        self._event_callback: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None

        # 连接状态
        self._is_connected = True

        # 持久化存储
        self._persistent_mgr = get_persistent_manager()
        self._persistent_session_created = False

    async def initialize(
        self,
        agent_type: AgentType = AgentType.GENERAL,
        event_callback: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None,
    ):
        """初始化会话"""
        self._event_callback = event_callback
        await self._create_agent(agent_type)

        log_event(
            "session_initialized",
            session_id=self.session_id,
            agent_type=agent_type.value,
        )

    async def _create_agent(self, agent_type: AgentType):
        """创建 Agent 实例"""
        self.agent_type = agent_type

        # 创建权限回调
        permission_callback = self._create_permission_callback()

        # 为主 Agent 动态注册 TaskTool
        if agent_type in [AgentType.GENERAL, AgentType.BUILD]:
            self._register_task_tool()

        # 创建模型客户端
        client = ModelClientFactory.create_client_from_model(self.model_name, self._config)

        # 创建 Agent
        self.agent = create_agent(
            agent_type=agent_type.value,
            client=client,
            tool_registry=self.tool_registry,
            permission_callback=permission_callback,
        )

        # 设置持久化
        self.agent.set_persistence(self.session_id)

    def _register_task_tool(self):
        """注册 TaskTool（仅主 Agent）"""
        from tools.task_tool import TaskTool

        task_tool = TaskTool(
            session_id=self.session_id,
            client_getter=lambda: self.agent.client,
            event_callback=self._event_callback,
            mcp_tools_getter=self._mcp_tools_getter,
            agent_types_getter=list_agent_types,
        )
        self.tool_registry.register(task_tool)

        log_event(
            "task_tool_registered",
            session_id=self.session_id,
        )

    def _create_permission_callback(self) -> PermissionCallback:
        """创建权限请求回调"""
        async def permission_callback(
            tool_name: str,
            pattern: str,
            args: Dict[str, Any],
        ) -> PermissionResponse:
            log_event(
                "permission_callback_start",
                session_id=self.session_id,
                permission=tool_name,
                pattern=pattern,
            )

            interaction_manager = get_interaction_manager()

            try:
                response = await interaction_manager.request_permission(
                    session_id=self.session_id,
                    permission=tool_name,
                    pattern=pattern,
                    metadata={"args": args},
                    timeout=None,
                )

                log_event(
                    "permission_response_received",
                    session_id=self.session_id,
                    request_id=response.request_id,
                    approved=response.approved,
                )

                return PermissionResponse(
                    approved=response.approved,
                    remember=response.remember,
                    remember_pattern=response.remember_pattern,
                )
            except RuntimeError as e:
                # 没有注册交互回调，默认放行（与旧逻辑保持一致）
                log_event(
                    "interaction_manager_not_available",
                    session_id=self.session_id,
                    error=str(e),
                )
                return PermissionResponse(approved=True)
            except asyncio.CancelledError:
                log_event(
                    "permission_request_cancelled",
                    session_id=self.session_id,
                    permission=tool_name,
                    pattern=pattern,
                )
                return PermissionResponse(approved=False)

        return permission_callback

    async def handle_permission_response(
        self,
        request_id: str,
        approved: bool,
        remember: bool = False,
        remember_pattern: Optional[str] = None,
    ):
        """处理权限响应"""
        if self._permission_future and self._permission_request_id == request_id:
            self._permission_future.set_result({
                "approved": approved,
                "remember": remember,
                "remember_pattern": remember_pattern,
            })

    # ==================== 任务执行 ====================

    def _request_cancel(self):
        """请求取消当前任务"""
        self._cancel_requested = True
        if self._cancel_event:
            self._cancel_event.set()

    async def _run_agent_stream(self, user_input: str) -> AsyncGenerator[Dict[str, Any], None]:
        """
        流式执行 Agent，包装生命周期事件

        注意：不在此处检查 _cancel_requested 并 break。
        中断信号通过 cancel_event 传播到 BaseAgent，由 BaseAgent 在安全检查点退出。
        """
        self._current_task_id = str(uuid.uuid4())
        self._is_running = True
        self._cancel_requested = False
        self._cancel_event = asyncio.Event()

        log_event(
            "agent_run_start",
            session_id=self.session_id,
            task_id=self._current_task_id,
            agent_type=self.agent.name,
        )

        yield {
            "type": EventType.TASK_STARTED,
            "data": {"task_id": self._current_task_id},
        }

        try:
            async for event in self.agent.run_stream(
                user_input, self.session_id,
                cancel_event=self._cancel_event,
            ):
                if event.get("type") == EventType.TASK_COMPLETED:
                    continue
                yield event

            if self._cancel_requested:
                yield {
                    "type": EventType.TASK_INTERRUPTED,
                    "data": {"task_id": self._current_task_id},
                }
                log_event(
                    "agent_run_interrupted",
                    session_id=self.session_id,
                    task_id=self._current_task_id,
                )
            else:
                yield {
                    "type": EventType.TASK_COMPLETED,
                    "data": {"task_id": self._current_task_id},
                }
                log_event(
                    "agent_run_end",
                    session_id=self.session_id,
                    task_id=self._current_task_id,
                )

        except Exception as e:
            log_event(
                "agent_run_error",
                session_id=self.session_id,
                error=str(e),
            )
            yield {
                "type": EventType.TASK_ERROR,
                "data": {
                    "task_id": self._current_task_id,
                    "message": f"Agent error: {str(e)}",
                },
            }

        finally:
            self._is_running = False
            self._current_task_id = None

    # ==================== 公开接口 ====================

    def _ensure_persistent_session(self):
        """确保持久化会话已创建"""
        if self._persistent_session_created:
            return
        existing = self._persistent_mgr.get_session(self.session_id)
        if not existing:
            self._persistent_mgr.create_session(
                session_id=self.session_id,
                title="New Session",
            )
        self._persistent_session_created = True

    async def execute(self, user_input: str):
        """执行用户输入"""
        if not self.agent:
            if self._event_callback:
                await self._event_callback({
                    "type": EventType.ERROR,
                    "data": {"message": "Agent not initialized"},
                })
            return

        self._ensure_persistent_session()

        # 流式执行并转发事件
        async for event in self._run_agent_stream(user_input):
            if self._event_callback:
                await self._event_callback(event)

    async def switch_agent(self, agent_type: AgentType) -> AgentType:
        """切换 Agent 类型"""
        old_type = self.agent_type

        if self._is_running:
            self._request_cancel()

        await self._create_agent(agent_type)

        log_event(
            "agent_switched",
            session_id=self.session_id,
            old_type=old_type.value,
            new_type=agent_type.value,
        )

        return old_type

    async def switch_model(self, model_name: str) -> Dict[str, str]:
        """切换模型（保留上下文）"""
        new_client = ModelClientFactory.create_client_from_model(model_name, self._config)
        old_model = self.model_name
        self.model_name = model_name

        if self._is_running:
            self._request_cancel()

        if self.agent:
            old_model = self.agent.client.model
            self.agent.client = new_client
            from provider.models import get_context_window
            new_max_tokens = get_context_window(new_client.model)
            for ctx in self.agent._contexts.values():
                ctx.config.max_tokens = new_max_tokens

        log_event(
            "model_switched",
            session_id=self.session_id,
            old_model=old_model,
            new_model=new_client.model,
        )

        return {"old_model": old_model, "new_model": new_client.model}

    def interrupt(self):
        """中断当前执行"""
        if self.agent and self._is_running:
            self._request_cancel()

    def disconnect(self):
        """断开连接"""
        self._is_connected = False
        if self._is_running:
            self._request_cancel()

    @property
    def is_connected(self) -> bool:
        return self._is_connected

    @property
    def is_running(self) -> bool:
        return self._is_running

    def clear_history(self):
        """清除会话历史"""
        if self.agent:
            self.agent.clear_history(self.session_id)

        log_event(
            "session_history_cleared",
            session_id=self.session_id,
        )

    async def compact_session(self, force: bool = False) -> Dict[str, Any]:
        """压缩会话上下文

        Args:
            force: 是否强制执行压缩（忽略自动压缩阈值检查）
        """
        if not self.agent:
            return {"compressed_count": 0, "success": False}

        if self._event_callback:
            await self._event_callback({
                "type": EventType.CONTEXT_COMPRESS_START,
                "data": {},
            })

        try:
            if hasattr(self.agent, 'compact'):
                await self.agent.compact(self.session_id)
            elif hasattr(self.agent, '_contexts'):
                context = self.agent._contexts.get(self.session_id)
                if context and hasattr(context, 'compress_if_needed'):
                    await context.compress_if_needed(force=force)

            result = {"compressed_count": 1, "success": True}

            if self._event_callback:
                await self._event_callback({
                    "type": EventType.CONTEXT_COMPRESS_END,
                    "data": {"success": True},
                })

            return result
        except Exception as e:
            log_event("session_compact_error", session_id=self.session_id, error=str(e))

            if self._event_callback:
                await self._event_callback({
                    "type": EventType.CONTEXT_COMPRESS_END,
                    "data": {"success": False, "error": str(e)},
                })

            return {"compressed_count": 0, "success": False, "error": str(e)}

    def get_context_info(self, thread_id: str) -> dict:
        """获取上下文信息"""
        if self.agent:
            context = self.agent._contexts.get(thread_id)
            if context:
                return context.get_context_info()
        return {"token_count": 0, "max_tokens": 0, "percent": 0.0, "threshold": 80}


class ConnectionSessionManager:
    """连接会话管理器"""

    def __init__(self):
        self._sessions: Dict[str, Session] = {}

    def create_session(
        self,
        session_id: str,
        model_name: str,
        config: AppConfig,
        mcp_tools_getter: Optional[Callable[[], List[BaseTool]]] = None,
    ) -> Session:
        """创建新会话"""
        session = Session(
            session_id=session_id,
            model_name=model_name,
            config=config,
            mcp_tools_getter=mcp_tools_getter,
        )
        self._sessions[session_id] = session

        log_event(
            "session_created",
            session_id=session_id,
            total_sessions=len(self._sessions),
        )

        return session

    def get_session(self, session_id: str) -> Optional[Session]:
        return self._sessions.get(session_id)

    def remove_session(self, session_id: str):
        if session_id in self._sessions:
            session = self._sessions.pop(session_id)
            session.disconnect()

            log_event(
                "session_removed",
                session_id=session_id,
                total_sessions=len(self._sessions),
            )

    @property
    def session_count(self) -> int:
        return len(self._sessions)
