"""
Session - 会话管理

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

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

from provider import BaseModelClient
from provider.factory import ModelClientFactory
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


def _get_git_branch(cwd: str = None) -> str:
    """获取当前目录的 git 分支名，非 git 仓库返回空字符串"""
    try:
        result = subprocess.run(
            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
            cwd=cwd or os.getcwd(),
            capture_output=True,
            text=True,
            timeout=2,
        )
        if result.returncode == 0:
            branch = result.stdout.strip()
            # HEAD 表示 detached HEAD 状态，返回短 commit hash
            if branch == "HEAD":
                result2 = subprocess.run(
                    ["git", "rev-parse", "--short", "HEAD"],
                    cwd=cwd or os.getcwd(),
                    capture_output=True,
                    text=True,
                    timeout=2,
                )
                if result2.returncode == 0:
                    return result2.stdout.strip()
            return branch
    except Exception:
        pass
    return ""


# 权限请求回调类型
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._auto_mode: bool = False

        # 持久化存储
        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,
            )

            # 自动模式：跳过所有权限确认
            if self._auto_mode:
                log_event(
                    "permission_auto_approved",
                    session_id=self.session_id,
                    permission=tool_name,
                    pattern=pattern,
                )
                return PermissionResponse(approved=True)

            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
            new_max_tokens = None
            try:
                provider_name = getattr(new_client, "provider", None)
                models_by_provider = getattr(self._config.provider, "models_by_provider", {}) or {}
                infos = models_by_provider.get(provider_name) or []
                for mi in infos:
                    if mi.get("name") == new_client.model:
                        new_max_tokens = mi.get("context_window")
                        break
            except Exception:
                new_max_tokens = None

            if not isinstance(new_max_tokens, int) or new_max_tokens <= 0:
                # Fallback: use client-injected window (resolved by ProviderManager/Factory) or safe default.
                new_max_tokens = getattr(new_client, "context_window", 128000) or 128000
            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}

    def get_session_info(self) -> dict:
        """获取完整会话信息（状态栏所需的所有数据）

        聚合上下文 token 信息和会话状态，供 /context/{session_id} 接口一次性返回。
        """
        cwd = os.getcwd()
        context_info = self.get_context_info(self.session_id)
        return {
            # 上下文 token 信息（保持兼容）
            "token_count": context_info.get("token_count", 0),
            "max_tokens": context_info.get("max_tokens", 0),
            "percent": context_info.get("percent", 0.0),
            "threshold": context_info.get("threshold", 80),
            # 会话状态信息
            "session_id": self.session_id,
            "model_name": self.model_name,
            "agent_type": self.agent_type.value,
            "auto_mode": self._auto_mode,
            "working_path": cwd,
            "git_branch": _get_git_branch(cwd),
            "is_running": self._is_running,
        }

    # ==================== 自动模式 ====================

    def set_auto_mode(self, enabled: bool) -> bool:
        """设置自动模式（跳过所有权限确认）

        Args:
            enabled: 是否启用自动模式

        Returns:
            设置后的状态
        """
        self._auto_mode = enabled
        log_event(
            "auto_mode_changed",
            session_id=self.session_id,
            auto_mode=enabled,
        )
        return self._auto_mode

    def get_auto_mode(self) -> bool:
        """获取自动模式状态"""
        return self._auto_mode


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)
