import asyncio
import hashlib
import json
import platform
from contextlib import asynccontextmanager
from typing import Optional, AsyncGenerator, Union, Callable, Awaitable, Dict

from loguru import logger
from playwright.async_api import (
    async_playwright,
    Browser,
    BrowserContext,
    Page,
    Playwright,
    Response,
)

from flow_xhs.runtime.app_utils import get_system_chrome_path
from flow_xhs.runtime.config import (
    BROWSER_LOCALE,
    BROWSER_TIMEZONE,
    resolve_browser_headless,
)

# Handle playwright-stealth version differences
try:
    from playwright_stealth import stealth_async
except ImportError:
    try:
        from playwright_stealth import Stealth


        async def stealth_async(page):
            await Stealth().apply_stealth_async(page)
    except ImportError:
        async def stealth_async(page):
            logger.warning("Stealth module not found, proceeding without stealth.")
            pass

class BrowserManager:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(BrowserManager, cls).__new__(cls)
            cls._instance.is_running = False
            cls._instance._local_playwrights: Dict[int, Playwright] = {}
            cls._instance._local_browsers: Dict[int, Browser] = {}
            cls._instance._loop_semaphores: Dict[int, asyncio.Semaphore] = {}
            cls._instance._loop_start_locks: Dict[int, asyncio.Lock] = {}
            cls._instance._profile_cache: Dict[str, dict] = {}

        return cls._instance

    @staticmethod
    def _pick_by_hash(digest: str, start: int, end: int, offset: int) -> int:
        span = end - start + 1
        return start + (int(digest[offset: offset + 8], 16) % span)

    def _resolve_profile(self, account_key: Optional[str]) -> dict:
        key = account_key or "default"
        cached = self._profile_cache.get(key)
        if cached:
            return cached

        digest = hashlib.sha256(key.encode("utf-8")).hexdigest()
        system = platform.system()

        if system == "Darwin":
            user_agent = (
                "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/131.0.0.0 Safari/537.36"
            )
            device_scale_factor = 2.0
        else:
            user_agent = (
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/131.0.0.0 Safari/537.36"
            )
            device_scale_factor = 1.0

        profile = {
            "user_agent": user_agent,
            "viewport": {
                "width": self._pick_by_hash(digest, 1366, 1920, 0),
                "height": self._pick_by_hash(digest, 768, 1080, 8),
            },
            "device_scale_factor": device_scale_factor,
        }
        self._profile_cache[key] = profile
        return profile

    @staticmethod
    def _normalize_storage_state(storage_state: Optional[Union[dict, str]]) -> Optional[dict]:
        if isinstance(storage_state, str) and storage_state:
            try:
                return json.loads(storage_state)
            except Exception:
                return None
        return storage_state if isinstance(storage_state, dict) else None

    @staticmethod
    def _get_loop_id() -> int:
        return id(asyncio.get_running_loop())

    def _get_loop_semaphore(self) -> asyncio.Semaphore:
        loop_id = id(asyncio.get_running_loop())
        if loop_id not in self._loop_semaphores:
            self._loop_semaphores[loop_id] = asyncio.Semaphore(10)
        return self._loop_semaphores[loop_id]

    def _get_loop_start_lock(self, loop_id: int) -> asyncio.Lock:
        if loop_id not in self._loop_start_locks:
            self._loop_start_locks[loop_id] = asyncio.Lock()
        return self._loop_start_locks[loop_id]

    async def _ensure_loop_browser(self, headless: bool | None = None) -> int:
        headless = resolve_browser_headless(headless)
        loop_id = self._get_loop_id()
        if loop_id in self._local_browsers:
            return loop_id

        async with self._get_loop_start_lock(loop_id):
            if loop_id in self._local_browsers:
                return loop_id

            logger.info(f"DEBUG: [BrowserManager] loop={loop_id} 正在以直连模式启动浏览器...")
            launch_args = [
                "--disable-blink-features=AutomationControlled",
                "--no-sandbox",
                "--disable-dev-shm-usage",
            ]

            launch_kwargs = {
                "headless": headless,
                "args": launch_args,
            }
            browser_bin = get_system_chrome_path()
            if browser_bin:
                launch_kwargs["executable_path"] = browser_bin

            playwright = await async_playwright().start()
            browser = await playwright.chromium.launch(**launch_kwargs)
            self._local_playwrights[loop_id] = playwright
            self._local_browsers[loop_id] = browser
            self._loop_semaphores.setdefault(loop_id, asyncio.Semaphore(10))
            self.is_running = True
            logger.info(f"DEBUG: [BrowserManager] loop={loop_id} 浏览器直连模式启动完成")
            return loop_id

    async def start(self, headless: bool | None = None):
        """为当前事件循环启动浏览器管理服务。"""
        await self._ensure_loop_browser(headless=headless)

    async def stop(self):
        """停止所有事件循环持有的浏览器资源。"""
        if not self.is_running:
            return

        for loop_id, browser in list(self._local_browsers.items()):
            try:
                await browser.close()
            except Exception as e:
                logger.warning(f"DEBUG: [BrowserManager] loop={loop_id} 关闭浏览器失败: {e}")
        self._local_browsers.clear()

        for loop_id, playwright in list(self._local_playwrights.items()):
            try:
                await playwright.stop()
            except Exception as e:
                logger.warning(f"DEBUG: [BrowserManager] loop={loop_id} 停止 playwright 失败: {e}")
        self._local_playwrights.clear()
        self._loop_semaphores.clear()
        self._loop_start_locks.clear()

        self.is_running = False
        logger.info("DEBUG: [BrowserManager] 浏览器已停止")

    @asynccontextmanager
    async def get_context(
            self,
            storage_state: Optional[Union[dict, str]] = None,
            headless: bool | None = None,
            account_key: Optional[str] = None,
            timezone_id: Optional[str] = None,
            locale: Optional[str] = None,
    ) -> AsyncGenerator[BrowserContext, None]:
        """获取浏览器上下文，确保账号画像稳定。"""
        headless = resolve_browser_headless(headless)
        loop_id = await self._ensure_loop_browser(headless=headless)
        browser = self._local_browsers.get(loop_id)
        if not browser:
            raise RuntimeError(f"Browser not initialized for loop={loop_id}")

        semaphore = self._get_loop_semaphore()
        profile = self._resolve_profile(account_key)
        normalized_state = self._normalize_storage_state(storage_state)

        async with semaphore:
            context = await browser.new_context(
                permissions=["geolocation"],
                user_agent=profile["user_agent"],
                viewport=profile["viewport"],
                device_scale_factor=profile["device_scale_factor"],
                storage_state=normalized_state,
                ignore_https_errors=True,
                locale=locale or BROWSER_LOCALE,
                timezone_id=timezone_id or BROWSER_TIMEZONE,
            )
            try:
                yield context
            finally:
                await context.close()

    @asynccontextmanager
    async def get_page(
            self,
            storage_state: Optional[Union[dict, str]] = None,
            headless: bool | None = None,
            account_key: Optional[str] = None,
            timezone_id: Optional[str] = None,
            locale: Optional[str] = None,
    ) -> AsyncGenerator[Page, None]:
        """获取页面上下文。"""
        async with self.get_context(
                storage_state=storage_state,
                headless=headless,
                account_key=account_key,
                timezone_id=timezone_id,
                locale=locale,
        ) as context:
            page = await context.new_page()
            await stealth_async(page)
            yield page

    @asynccontextmanager
    async def observe_responses(self, page: Page, handler: Callable[[Response], Awaitable[None]]):
        page.on("response", handler)
        try:
            yield
        finally:
            page.remove_listener("response", handler)


browser_manager = BrowserManager()
