#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@Author         : yanyongyu
@Date           : 2021-03-12 13:42:43
@LastEditors    : yanyongyu
@LastEditTime   : 2021-11-01 14:05:41
@Description    : None
@GitHub         : https://github.com/yanyongyu
"""
__author__ = 'yanyongyu'
import os
import platform
import subprocess
from contextlib import asynccontextmanager
from typing import AsyncIterator, Optional

from loguru import logger
from playwright.async_api import Browser, Error, Page, Playwright, async_playwright
from playwright.async_api import TimeoutError as PlaywrightTimeoutError

_browser: Optional[Browser] = None
_playwright: Optional[Playwright] = None
_use_browser: Optional[str] = None
_browser_pid: Optional[int] = None  # 保存浏览器进程 PID


async def init(use_browser, **kwargs) -> Browser:
    global _browser
    global _playwright
    global _use_browser
    _use_browser = use_browser
    _playwright = await async_playwright().start()
    try:
        _browser = await launch_browser(use_browser, **kwargs)
    except Error:
        await install_browser(use_browser)
        _browser = await launch_browser(use_browser, **kwargs)
    return _browser


async def launch_browser(use_browser, **kwargs) -> Browser:
    global _browser_pid
    assert _playwright is not None, 'Playwright 没有安装'
    if use_browser == 'firefox':
        logger.info('使用 firefox 启动')
        browser = await _playwright.firefox.launch(**kwargs)
    else:
        # 默认使用 chromium
        logger.info('使用 chromium 启动')
        browser = await _playwright.chromium.launch(**kwargs)

    # 保存浏览器进程 PID
    try:
        if browser.process:
            _browser_pid = browser.process.pid
            logger.info(f'浏览器进程 PID: {_browser_pid}')
    except Exception:
        pass

    return browser


async def get_browser(use_browser, **kwargs) -> Browser:
    """获取浏览器实例，如果不存在或已断开连接则创建新的"""
    if _browser and _browser.is_connected():
        # 检查浏览器是否真的健康
        try:
            # 尝试创建一个测试页面来验证浏览器是否正常工作
            test_page = await _browser.new_page()
            await test_page.close()
            return _browser
        except Exception:
            logger.warning('浏览器连接异常，准备重启浏览器')
            await force_restart_browser(use_browser, **kwargs)
            return _browser
    return await init(use_browser, **kwargs)


@asynccontextmanager
async def get_new_page(use_browser, **kwargs) -> AsyncIterator[Page]:
    """获取新页面，如果遇到超时异常会自动重启浏览器"""
    max_retries = 2
    retry_count = 0

    while retry_count < max_retries:
        try:
            browser = await get_browser(use_browser, **kwargs)
            page = await browser.new_page(**kwargs)
            # 设置默认超时时间为 60 秒，避免频繁超时
            page.set_default_timeout(60000)
            try:
                yield page
                break  # 成功执行，退出重试循环
            finally:
                try:
                    await page.close()
                except Exception:
                    pass  # 忽略关闭页面时的错误
        except PlaywrightTimeoutError as e:
            retry_count += 1
            logger.warning(f'Playwright 超时异常 (尝试 {retry_count}/{max_retries}): {e}')
            if retry_count < max_retries:
                logger.info('检测到浏览器可能假死，正在强制重启浏览器...')
                await force_restart_browser(use_browser, **kwargs)
            else:
                # 最后一次重试也失败，抛出异常
                logger.error('浏览器重启后仍然超时，放弃重试')
                raise
        except Exception as e:
            # 其他异常也尝试重启一次
            if retry_count == 0:
                logger.warning(f'浏览器操作异常，尝试重启: {e}')
                retry_count += 1
                await force_restart_browser(use_browser, **kwargs)
            else:
                raise


async def shutdown_browser():
    """关闭浏览器和 Playwright"""
    global _browser
    global _playwright
    global _use_browser
    if _browser:
        try:
            await _browser.close()
        except Exception:
            pass
        _browser = None
    if _playwright:
        try:
            await _playwright.stop()  # type: ignore
        except Exception:
            pass
        _playwright = None
    _use_browser = None


async def force_restart_browser(use_browser, **kwargs):
    """强制重启浏览器，只杀死 Playwright 启动的浏览器进程"""
    global _browser
    global _playwright
    global _use_browser
    global _browser_pid

    logger.info('开始强制重启浏览器...')

    # 1. 先尝试正常关闭
    if _browser:
        try:
            await _browser.close()
        except Exception as e:
            logger.warning(f'正常关闭浏览器失败: {e}')

    # 2. 关闭 Playwright
    if _playwright:
        try:
            await _playwright.stop()
        except Exception as e:
            logger.warning(f'关闭 Playwright 失败: {e}')

    # 3. 如果正常关闭失败，强制杀死 Playwright 启动的浏览器进程（只杀死我们自己的进程）
    if _browser_pid:
        try:
            kill_playwright_browser_process(_browser_pid, use_browser)
        except Exception as e:
            logger.warning(f'杀死浏览器进程时出错: {e}')

    # 4. 清理全局变量
    _browser = None
    _playwright = None
    _use_browser = None
    _browser_pid = None

    # 5. 重新初始化
    logger.info('重新初始化浏览器...')
    await init(use_browser, **kwargs)
    logger.info('浏览器重启完成')


def kill_playwright_browser_process(pid: int, use_browser: str):
    """只杀死指定的 Playwright 浏览器进程及其子进程（不影响其他浏览器）"""
    system = platform.system()

    try:
        if system == 'Windows':
            # Windows: 使用 taskkill 只杀死指定 PID 及其子进程
            # /T 参数会杀死该进程及其所有子进程
            # /F 强制杀死
            subprocess.run(
                ['taskkill', '/F', '/T', '/PID', str(pid)],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                check=False,
                timeout=5,
            )
            logger.info(f'已尝试杀死 Playwright 浏览器进程 PID: {pid} 及其子进程')
        elif system in ('Linux', 'Darwin'):
            # Linux/macOS: 使用 pkill 杀死进程树
            # -P 参数指定父进程 PID，会杀死该进程及其所有子进程
            subprocess.run(
                ['pkill', '-9', '-P', str(pid)],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                check=False,
                timeout=5,
            )
            # 也尝试直接杀死主进程
            subprocess.run(
                ['kill', '-9', str(pid)],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                check=False,
                timeout=5,
            )
            logger.info(f'已尝试杀死 Playwright 浏览器进程 PID: {pid} 及其子进程')
    except subprocess.TimeoutExpired:
        logger.warning(f'杀死进程超时: PID {pid}')
    except Exception as e:
        logger.warning(f'杀死浏览器进程失败: {e}')


async def install_browser(use_browser):
    import sys

    from playwright.__main__ import main

    logger.info('使用镜像源进行下载')
    os.environ['PLAYWRIGHT_DOWNLOAD_HOST'] = 'https://npmmirror.com/mirrors/playwright/'
    success = False

    if use_browser == 'firefox':
        logger.info('正在安装 firefox')
        sys.argv = ['', 'install', 'firefox']
    else:
        # 默认使用 chromium
        logger.info('正在安装 chromium')
        sys.argv = ['', 'install', 'chromium']
    try:
        logger.info('正在安装依赖')
        os.system('playwright install-deps')
        main()
    except SystemExit as e:
        if e.code == 0:
            success = True
    if not success:
        logger.error('浏览器更新失败, 请检查网络连通性')
