"""
测试隔离器
支持独立测试环境、状态隔离、自动清理
"""
import asyncio
import uuid
import json
import shutil
from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any, Callable, Awaitable
from pathlib import Path
from datetime import datetime
from contextlib import asynccontextmanager
from loguru import logger


@dataclass
class IsolationConfig:
    """隔离配置"""
    workspace_dir: str = "./test_isolation"
    cleanup_on_exit: bool = True
    preserve_on_failure: bool = True
    snapshot_state: bool = True
    isolate_cookies: bool = True
    isolate_storage: bool = True
    isolate_network: bool = False


@dataclass
class IsolationContext:
    """隔离上下文"""
    id: str
    workspace: str
    created_at: datetime = field(default_factory=datetime.now)
    test_name: Optional[str] = None
    parent_context: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    cleanup_tasks: List[Callable[[], Awaitable[None]]] = field(default_factory=list)

    # 隔离状态
    cookies_backup: Optional[List[Dict]] = None
    storage_backup: Optional[Dict[str, Any]] = None
    network_requests: List[Dict] = field(default_factory=list)


@dataclass
class IsolationResult:
    """隔离结果"""
    success: bool
    context_id: str
    test_name: Optional[str] = None
    error: Optional[str] = None
    duration_seconds: float = 0.0
    cleanup_performed: bool = False
    artifacts_preserved: bool = False


class TestIsolator:
    """
    测试隔离器
    提供测试环境和状态的隔离
    """

    def __init__(
        self,
        page,
        config: Optional[IsolationConfig] = None,
    ):
        self.page = page
        self.config = config or IsolationConfig()
        self._contexts: Dict[str, IsolationContext] = {}
        self._current_context: Optional[IsolationContext] = None
        self._counter = 0

    # ==================== 上下文管理 ====================

    async def create_context(
        self,
        test_name: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> IsolationContext:
        """
        创建隔离上下文

        Args:
            test_name: 测试名称
            metadata: 元数据

        Returns:
            IsolationContext: 隔离上下文
        """
        self._counter += 1
        context_id = f"iso_{self._counter}_{uuid.uuid4().hex[:6]}"

        # 创建工作空间
        workspace = Path(self.config.workspace_dir) / context_id
        workspace.mkdir(parents=True, exist_ok=True)

        context = IsolationContext(
            id=context_id,
            workspace=str(workspace),
            test_name=test_name,
            metadata=metadata or {},
        )

        # 备份状态
        if self.config.snapshot_state:
            await self._backup_state(context)

        self._contexts[context_id] = context
        self._current_context = context

        logger.debug(f"创建隔离上下文: {context_id}")
        return context

    async def enter_context(
        self,
        test_name: Optional[str] = None,
    ) -> IsolationContext:
        """
        进入隔离上下文

        Args:
            test_name: 测试名称

        Returns:
            IsolationContext: 隔离上下文
        """
        context = await self.create_context(test_name)

        # 应用隔离
        if self.config.isolate_cookies:
            await self._isolate_cookies(context)

        if self.config.isolate_storage:
            await self._isolate_storage(context)

        return context

    async def exit_context(
        self,
        context_id: str,
        success: bool = True,
    ) -> IsolationResult:
        """
        退出隔离上下文

        Args:
            context_id: 上下文 ID
            success: 是否成功

        Returns:
            IsolationResult: 隔离结果
        """
        context = self._contexts.get(context_id)
        if not context:
            return IsolationResult(
                success=False,
                context_id=context_id,
                error="上下文不存在",
            )

        result = IsolationResult(
            success=success,
            context_id=context_id,
            test_name=context.test_name,
        )

        # 执行清理
        if self.config.cleanup_on_exit:
            preserve = not success and self.config.preserve_on_failure
            await self._cleanup(context, preserve)
            result.cleanup_performed = True
            result.artifacts_preserved = preserve

        # 恢复状态
        if self.config.snapshot_state:
            await self._restore_state(context)

        if self._current_context == context:
            self._current_context = None

        return result

    @asynccontextmanager
    async def isolate(
        self,
        test_name: Optional[str] = None,
    ):
        """
        隔离上下文管理器

        Args:
            test_name: 测试名称

        Yields:
            IsolationContext: 隔离上下文
        """
        context = await self.enter_context(test_name)
        success = True
        error = None

        try:
            yield context
        except Exception as e:
            success = False
            error = str(e)
            raise
        finally:
            result = await self.exit_context(context.id, success)
            if error:
                logger.error(f"隔离测试失败: {test_name}, error={error}")

    # ==================== 状态备份/恢复 ====================

    async def _backup_state(self, context: IsolationContext):
        """备份当前状态"""
        try:
            # 备份 Cookies
            if self.config.isolate_cookies:
                context.cookies_backup = await self.page.context.cookies()

            # 备份本地存储
            if self.config.isolate_storage:
                storage = await self.page.evaluate("""
                    () => ({
                        localStorage: {...localStorage},
                        sessionStorage: {...sessionStorage},
                    })
                """)
                context.storage_backup = storage

        except Exception as e:
            logger.warning(f"状态备份失败: {e}")

    async def _restore_state(self, context: IsolationContext):
        """恢复状态"""
        try:
            # 恢复 Cookies
            if context.cookies_backup:
                await self.page.context.clear_cookies()
                await self.page.context.add_cookies(context.cookies_backup)

            # 恢复存储
            if context.storage_backup:
                await self.page.evaluate("""
                    (storage) => {
                        // 清除并恢复 localStorage
                        localStorage.clear();
                        Object.entries(storage.localStorage || {}).forEach(([k, v]) => {
                            localStorage.setItem(k, v);
                        });
                        // 清除并恢复 sessionStorage
                        sessionStorage.clear();
                        Object.entries(storage.sessionStorage || {}).forEach(([k, v]) => {
                            sessionStorage.setItem(k, v);
                        });
                    }
                """, context.storage_backup)

        except Exception as e:
            logger.warning(f"状态恢复失败: {e}")

    async def _isolate_cookies(self, context: IsolationContext):
        """隔离 Cookies"""
        await self.page.context.clear_cookies()

    async def _isolate_storage(self, context: IsolationContext):
        """隔离存储"""
        await self.page.evaluate("""
            () => {
                localStorage.clear();
                sessionStorage.clear();
            }
        """)

    # ==================== 清理机制 ====================

    async def _cleanup(
        self,
        context: IsolationContext,
        preserve_artifacts: bool = False,
    ):
        """执行清理"""
        # 执行注册的清理任务
        for task in context.cleanup_tasks:
            try:
                await task()
            except Exception as e:
                logger.warning(f"清理任务失败: {e}")

        # 清理工作空间
        if not preserve_artifacts:
            workspace = Path(context.workspace)
            if workspace.exists():
                try:
                    shutil.rmtree(workspace)
                except Exception as e:
                    logger.warning(f"清理工作空间失败: {e}")

    def register_cleanup(
        self,
        context_id: str,
        task: Callable[[], Awaitable[None]],
    ):
        """
        注册清理任务

        Args:
            context_id: 上下文 ID
            task: 清理任务
        """
        context = self._contexts.get(context_id)
        if context:
            context.cleanup_tasks.append(task)

    # ==================== 工作空间管理 ====================

    def get_workspace(self, context_id: str) -> Optional[str]:
        """获取工作空间路径"""
        context = self._contexts.get(context_id)
        return context.workspace if context else None

    async def save_artifact(
        self,
        context_id: str,
        name: str,
        content: bytes,
    ) -> str:
        """
        保存工件

        Args:
            context_id: 上下文 ID
            name: 工件名称
            content: 工件内容

        Returns:
            str: 工件路径
        """
        context = self._contexts.get(context_id)
        if not context:
            raise ValueError(f"上下文不存在: {context_id}")

        artifacts_dir = Path(context.workspace) / "artifacts"
        artifacts_dir.mkdir(exist_ok=True)

        artifact_path = artifacts_dir / name
        with open(artifact_path, 'wb') as f:
            f.write(content)

        return str(artifact_path)

    async def save_screenshot(
        self,
        context_id: str,
        name: str,
    ) -> str:
        """保存截图"""
        context = self._contexts.get(context_id)
        if not context:
            raise ValueError(f"上下文不存在: {context_id}")

        screenshots_dir = Path(context.workspace) / "screenshots"
        screenshots_dir.mkdir(exist_ok=True)

        path = screenshots_dir / f"{name}.png"
        await self.page.screenshot(path=str(path))

        return str(path)

    # ==================== 信息查询 ====================

    def get_context(self, context_id: str) -> Optional[IsolationContext]:
        """获取上下文"""
        return self._contexts.get(context_id)

    def list_contexts(self) -> List[str]:
        """列出所有上下文"""
        return list(self._contexts.keys())

    def get_current_context(self) -> Optional[IsolationContext]:
        """获取当前上下文"""
        return self._current_context

    async def cleanup_all(self):
        """清理所有上下文"""
        for context_id in list(self._contexts.keys()):
            await self.exit_context(context_id, success=True)

        # 清理工作空间根目录
        root = Path(self.config.workspace_dir)
        if root.exists() and self.config.cleanup_on_exit:
            try:
                shutil.rmtree(root)
            except Exception as e:
                logger.warning(f"清理工作空间根目录失败: {e}")


def create_test_isolator(
    page,
    config: Optional[IsolationConfig] = None,
) -> TestIsolator:
    """创建测试隔离器实例"""
    return TestIsolator(page, config)
