"""
浏览器实例池
管理多个浏览器实例
"""
import asyncio
from typing import Optional, Dict, Any, Callable, Awaitable, TypeVar
from loguru import logger

from .types import BrowserPoolConfig, BrowserInstance

T = TypeVar('T')


class BrowserPool:
    """浏览器实例池，管理多个浏览器实例"""

    def __init__(self, config: Optional[BrowserPoolConfig] = None):
        self.config = config or BrowserPoolConfig()
        self._instances: Dict[str, BrowserInstance] = {}
        self._playwright = None
        self._lock = asyncio.Lock()
        self._instance_counter = 0
        self._is_initialized = False

    async def initialize(self):
        if self._is_initialized:
            return
        from playwright.async_api import async_playwright
        self._playwright = await async_playwright().start()
        self._is_initialized = True
        logger.info(f"浏览器池已初始化，最大实例数: {self.config.max_instances}")

    async def close(self):
        for instance_id in list(self._instances.keys()):
            await self._close_instance(instance_id)
        if self._playwright:
            await self._playwright.stop()
            self._playwright = None
        self._is_initialized = False

    async def __aenter__(self):
        await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    async def acquire(self) -> BrowserInstance:
        if not self._is_initialized:
            await self.initialize()
        async with self._lock:
            available = self._find_available_instance()
            if available:
                available.is_busy = True
                available.last_used = datetime.now()
                return available
            if len(self._instances) < self.config.max_instances:
                instance = await self._create_instance()
                instance.is_busy = True
                return instance
            raise RuntimeError("浏览器池已满，无可用实例")

    async def release(self, instance_id: str):
        async with self._lock:
            instance = self._instances.get(instance_id)
            if instance:
                instance.is_busy = False
                instance.last_used = datetime.now()

    async def _create_instance(self) -> BrowserInstance:
        from datetime import datetime
        self._instance_counter += 1
        instance_id = f"browser_{self._instance_counter}"
        browser_launcher = getattr(self._playwright, self.config.browser_type)
        browser = await browser_launcher.launch(headless=self.config.headless, timeout=self.config.timeout)
        instance = BrowserInstance(id=instance_id, browser=browser)
        self._instances[instance_id] = instance
        return instance

    async def _close_instance(self, instance_id: str):
        instance = self._instances.pop(instance_id, None)
        if instance:
            for context in instance.contexts:
                await context.close()
            await instance.browser.close()

    def _find_available_instance(self) -> Optional[BrowserInstance]:
        for instance in self._instances.values():
            if not instance.is_busy and instance.context_count < self.config.max_contexts_per_browser:
                return instance
        return None

    async def create_context(self, instance_id: Optional[str] = None, **options) -> Any:
        if instance_id:
            instance = self._instances.get(instance_id)
            if not instance:
                raise ValueError(f"实例不存在: {instance_id}")
        else:
            instance = await self.acquire()
        context = await instance.browser.new_context(**options)
        instance.contexts.append(context)
        return context

    async def execute(self, task: Callable[[Any], Awaitable[T]], **options) -> T:
        instance = await self.acquire()
        context = None
        page = None
        try:
            context = await self.create_context(instance.id, **options)
            page = await context.new_page()
            return await task(page)
        finally:
            if page:
                await page.close()
            if context:
                await context.close()
            await self.release(instance.id)

    def get_stats(self) -> Dict[str, Any]:
        return {
            "total_instances": len(self._instances),
            "busy_instances": sum(1 for i in self._instances.values() if i.is_busy),
            "total_contexts": sum(i.context_count for i in self._instances.values()),
            "is_initialized": self._is_initialized,
        }

    async def scale_to(self, target_count: int):
        """扩缩容到指定实例数"""
        async with self._lock:
            current = len(self._instances)
            if target_count > current:
                for _ in range(target_count - current):
                    await self._create_instance()
            elif target_count < current:
                idle_instances = [
                    i for i in self._instances.values()
                    if not i.is_busy
                ]
                for instance in idle_instances[current - target_count:]:
                    await self._close_instance(instance.id)
