"""
浏览器上下文管理
高级上下文管理功能
"""
import asyncio
from typing import Optional, Dict, Any, List
from playwright.async_api import BrowserContext, Page, Browser
from loguru import logger
from dataclasses import dataclass
from pathlib import Path


@dataclass
class ContextProfile:
    """上下文配置"""
    name: str
    storage_state: Optional[str] = None
    viewport: Optional[Dict[str, int]] = None
    user_agent: Optional[str] = None
    locale: Optional[str] = None
    timezone: Optional[str] = None
    geolocation: Optional[Dict[str, float]] = None
    permissions: Optional[List[str]] = None


class ContextManager:
    """
    浏览器上下文管理器
    管理多个浏览器上下文
    """
    
    def __init__(self, browser: Browser):
        self.browser = browser
        self._contexts: Dict[str, BrowserContext] = {}
        self._pages: Dict[str, Page] = {}
    
    async def create_context(
        self,
        name: str,
        profile: ContextProfile = None,
        **kwargs,
    ) -> BrowserContext:
        """
        创建浏览器上下文
        
        Args:
            name: 上下文名称
            profile: 上下文配置
            **kwargs: 额外参数
            
        Returns:
            BrowserContext: 浏览器上下文
        """
        if name in self._contexts:
            raise ValueError(f"上下文已存在: {name}")
        
        options = {}
        
        if profile:
            if profile.storage_state:
                options["storage_state"] = profile.storage_state
            if profile.viewport:
                options["viewport"] = profile.viewport
            if profile.user_agent:
                options["user_agent"] = profile.user_agent
            if profile.locale:
                options["locale"] = profile.locale
            if profile.timezone:
                options["timezone_id"] = profile.timezone
            if profile.geolocation:
                options["geolocation"] = profile.geolocation
            if profile.permissions:
                options["permissions"] = profile.permissions
        
        options.update(kwargs)
        
        context = await self.browser.new_context(**options)
        self._contexts[name] = context
        
        logger.info(f"创建上下文: {name}")
        return context
    
    async def get_context(self, name: str) -> Optional[BrowserContext]:
        """获取上下文"""
        return self._contexts.get(name)
    
    async def create_page(self, context_name: str, page_name: str) -> Page:
        """
        在上下文中创建页面
        
        Args:
            context_name: 上下文名称
            page_name: 页面名称
            
        Returns:
            Page: 页面对象
        """
        context = self._contexts.get(context_name)
        if not context:
            raise ValueError(f"上下文不存在: {context_name}")
        
        page = await context.new_page()
        self._pages[page_name] = page
        
        logger.info(f"创建页面: {page_name} (context: {context_name})")
        return page
    
    async def get_page(self, page_name: str) -> Optional[Page]:
        """获取页面"""
        return self._pages.get(page_name)
    
    async def close_context(self, name: str):
        """关闭上下文"""
        if name in self._contexts:
            await self._contexts[name].close()
            del self._contexts[name]
            logger.info(f"关闭上下文: {name}")
    
    async def close_all(self):
        """关闭所有上下文"""
        for name in list(self._contexts.keys()):
            await self.close_context(name)
    
    def list_contexts(self) -> List[str]:
        """列出所有上下文"""
        return list(self._contexts.keys())
    
    def list_pages(self) -> List[str]:
        """列出所有页面"""
        return list(self._pages.keys())
    
    async def save_state(self, context_name: str, filepath: str):
        """保存上下文状态"""
        context = self._contexts.get(context_name)
        if not context:
            raise ValueError(f"上下文不存在: {context_name}")
        
        await context.storage_state(path=filepath)
        logger.info(f"保存状态: {filepath}")
    
    async def load_state(self, context_name: str, filepath: str) -> BrowserContext:
        """加载上下文状态"""
        if context_name in self._contexts:
            raise ValueError(f"上下文已存在: {context_name}")
        
        context = await self.browser.new_context(storage_state=filepath)
        self._contexts[context_name] = context
        
        logger.info(f"加载状态: {filepath}")
        return context


class IsolatedSession:
    """
    隔离会话
    提供完全隔离的浏览器会话
    """
    
    def __init__(
        self,
        name: str,
        browser: Browser,
        isolated_storage: bool = True,
    ):
        self.name = name
        self.browser = browser
        self.isolated_storage = isolated_storage
        self._context: Optional[BrowserContext] = None
        self._page: Optional[Page] = None
        self._storage_path: Optional[str] = None
    
    async def start(self) -> Page:
        """启动会话"""
        options = {}
        
        if self.isolated_storage:
            self._storage_path = f"./sessions/{self.name}_state.json"
            Path(self._storage_path).parent.mkdir(parents=True, exist_ok=True)
            
            if Path(self._storage_path).exists():
                options["storage_state"] = self._storage_path
        
        self._context = await self.browser.new_context(**options)
        self._page = await self._context.new_page()
        
        logger.info(f"隔离会话已启动: {self.name}")
        return self._page
    
    async def save(self):
        """保存会话状态"""
        if self._context and self._storage_path:
            await self._context.storage_state(path=self._storage_path)
            logger.info(f"会话状态已保存: {self.name}")
    
    async def close(self):
        """关闭会话"""
        if self._context:
            await self.save()
            await self._context.close()
            self._context = None
            self._page = None
            logger.info(f"隔离会话已关闭: {self.name}")
    
    @property
    def page(self) -> Optional[Page]:
        """获取页面对象"""
        return self._page
    
    async def __aenter__(self) -> Page:
        return await self.start()
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()


class MultiUserSession:
    """
    多用户会话管理
    同时管理多个用户的浏览器会话
    """
    
    def __init__(self, browser: Browser):
        self.browser = browser
        self._sessions: Dict[str, IsolatedSession] = {}
    
    async def create_user(self, username: str) -> Page:
        """创建用户会话"""
        if username in self._sessions:
            raise ValueError(f"用户会话已存在: {username}")
        
        session = IsolatedSession(username, self.browser)
        self._sessions[username] = session
        
        return await session.start()
    
    async def get_user_page(self, username: str) -> Optional[Page]:
        """获取用户页面"""
        session = self._sessions.get(username)
        return session.page if session else None
    
    async def switch_user(self, username: str) -> Page:
        """切换用户"""
        page = await self.get_user_page(username)
        if not page:
            raise ValueError(f"用户会话不存在: {username}")
        return page
    
    async def close_user(self, username: str):
        """关闭用户会话"""
        session = self._sessions.get(username)
        if session:
            await session.close()
            del self._sessions[username]
    
    async def close_all(self):
        """关闭所有用户会话"""
        for username in list(self._sessions.keys()):
            await self.close_user(username)
    
    def list_users(self) -> List[str]:
        """列出所有用户"""
        return list(self._sessions.keys())
    
    async def broadcast(self, action):
        """向所有用户广播操作"""
        import asyncio
        
        tasks = []
        for session in self._sessions.values():
            if session.page:
                if asyncio.iscoroutinefunction(action):
                    tasks.append(action(session.page))
                else:
                    action(session.page)
        
        if tasks:
            await asyncio.gather(*tasks)


class SessionPool:
    """
    会话池
    复用浏览器会话
    """
    
    def __init__(
        self,
        browser: Browser,
        pool_size: int = 5,
    ):
        self.browser = browser
        self.pool_size = pool_size
        self._pool: List[BrowserContext] = []
        self._in_use: Dict[int, BrowserContext] = {}
        self._counter = 0
    
    async def initialize(self):
        """初始化会话池"""
        for _ in range(self.pool_size):
            context = await self.browser.new_context()
            self._pool.append(context)
        
        logger.info(f"会话池已初始化: {self.pool_size} 个会话")
    
    async def acquire(self) -> tuple:
        """
        获取会话
        
        Returns:
            tuple: (session_id, page)
        """
        if not self._pool:
            raise RuntimeError("会话池为空，请先初始化")
        
        context = self._pool.pop()
        page = await context.new_page()
        
        self._counter += 1
        session_id = self._counter
        self._in_use[session_id] = context
        
        logger.debug(f"获取会话: {session_id}")
        return session_id, page
    
    async def release(self, session_id: int):
        """释放会话"""
        if session_id not in self._in_use:
            return
        
        context = self._in_use.pop(session_id)
        
        # 关闭所有页面
        for page in context.pages:
            await page.close()
        
        self._pool.append(context)
        logger.debug(f"释放会话: {session_id}")
    
    async def close(self):
        """关闭会话池"""
        for context in self._pool:
            await context.close()
        
        for context in self._in_use.values():
            await context.close()
        
        self._pool.clear()
        self._in_use.clear()
        
        logger.info("会话池已关闭")
