"""
作者：Xiaoqiang
微信公众号：XiaoqiangClub
创建时间：2025/5/29 09:07
最后修改时间：2025/6/3 09:06
文件描述：Chrome绿色版下载工具
文件路径：/AutoChrome/AutoChrome/chrome_downloader.py
版权声明：© 2025 Xiaoqiang. All Rights Reserved.
"""

import os
import shutil
import zipfile
import tempfile
import platform
import requests
from tqdm import tqdm
from typing import Optional, Dict, Any, Literal


class ChromeDownloader:
    # Chromium最新快照版本号获取URL
    CHROMIUM_SNAPSHOT_URLS: Dict[str, str] = {
        "windows": "https://storage.googleapis.com/chromium-browser-snapshots/Win_x64/LAST_CHANGE",
        "darwin": "https://storage.googleapis.com/chromium-browser-snapshots/Mac/LAST_CHANGE",
        "linux": "https://storage.googleapis.com/chromium-browser-snapshots/Linux_x64/LAST_CHANGE",
    }
    # Chromium官方快照下载URL模板
    CHROMIUM_DOWNLOAD_URLS: Dict[str, str] = {
        "windows": "https://storage.googleapis.com/chromium-browser-snapshots/Win_x64/{rev}/chrome-win.zip",
        "darwin": "https://storage.googleapis.com/chromium-browser-snapshots/Mac/{rev}/chrome-mac.zip",
        "linux": "https://storage.googleapis.com/chromium-browser-snapshots/Linux_x64/{rev}/chrome-linux.zip",
    }

    def __init__(
            self,
            download_dir: Optional[str] = None,
            logger: Optional[Any] = None,
            custom_source: Optional[Dict[str, Dict[str, str]]] = None,
    ):
        """
        Chromium绿色版下载工具

        :param download_dir: 下载目录，默认为当前目录下的 chrome 文件夹
        :param logger: 日志对象
        :param custom_source: 自定义源，格式见文档
        """
        self.system: str = platform.system().lower()
        self.download_dir: str = download_dir or os.path.abspath("chrome")
        self.logger: Optional[Any] = logger
        self.custom_source: Optional[Dict[str, Dict[str, str]]] = custom_source

    def log(self, msg: str, level: str = "info") -> None:
        """
        日志输出

        :param msg: 日志内容
        :param level: 日志等级
        """
        if self.logger:
            getattr(self.logger, level, self.logger.info)(msg)
        else:
            print(msg)

    def get_latest_chromium_revision(
            self, system: Optional[str] = None
    ) -> Optional[str]:
        """
        获取最新的 Chromium 快照版本号

        :param system: 指定系统，默认当前系统
        :return: 版本号
        """
        sys_name: str = (system or self.system).lower()
        url: Optional[str] = None
        # 优先自定义源
        if self.custom_source and sys_name in {
            k.lower(): v for k, v in self.custom_source.items()
        }:
            url = self.custom_source[sys_name].get("latest")
        else:
            url = self.CHROMIUM_SNAPSHOT_URLS.get(sys_name)
        if not url:
            self.log(f"❌ 暂不支持当前系统：{sys_name}", "error")
            return None
        try:
            resp = requests.get(url, timeout=10)
            resp.raise_for_status()
            rev = resp.text.strip()
            self.log(f"✅ 最新Chromium快照版本号：{rev}")
            return rev
        except Exception as e:
            self.log(f"⚠️ 获取Chromium快照版本号失败：{e}", "warning")
            return None

    def download_chromium(
            self, revision: Optional[str] = None, system: Optional[str] = None
    ) -> Optional[str]:
        """
        下载 Chromium 快照

        :param revision: 指定版本号，默认为最新版本
        :param system: 指定系统，默认为当前系统
        :return: 解压目录
        """
        sys_name: str = (system or self.system).lower()
        rev: Optional[str] = revision or self.get_latest_chromium_revision(
            system=sys_name
        )
        if not rev:
            return None
        # 优先自定义源
        if self.custom_source and sys_name in self.custom_source:
            url_tpl = self.custom_source[sys_name].get("download")
        else:
            url_tpl = self.CHROMIUM_DOWNLOAD_URLS.get(sys_name)
        if not url_tpl:
            self.log(f"❌ 暂不支持当前系统：{sys_name}", "error")
            return None
        url = url_tpl.format(rev=rev)
        self.log(f"⬇️  开始下载 Chromium 快照版：{url}")
        extract_dir = os.path.join(self.download_dir, f"chromium_{rev}")
        return self._download_and_extract(url, extract_dir)

    def _find_chrome_exe(self, base_path: str) -> Optional[str]:
        """
        在给定路径下查找 chrome 可执行文件

        :param base_path: 基础路径
        :return: chrome可执行文件路径
        """
        if self.system == "windows":
            return os.path.join(base_path, "chrome-win", "chrome.exe")
        elif self.system == "linux":
            return os.path.join(base_path, "chrome-linux", "chrome")
        elif self.system == "darwin":
            return os.path.join(
                base_path, "chrome-mac", "Chromium.app", "Contents", "MacOS", "Chromium"
            )
        return None

    def _download_and_extract(
            self, url: str, extract_dir: str, desc: str = "⏳ 下载进度"
    ) -> Optional[str]:
        """
        通用下载和解压方法

        :param url: 下载地址
        :param extract_dir: 解压目录
        :param desc: 进度描述
        :return: 解压后的路径
        """
        try:
            with requests.get(url, stream=True, timeout=60) as r:
                r.raise_for_status()
                total = int(r.headers.get("content-length", 0))
                os.makedirs(self.download_dir, exist_ok=True)
                with tempfile.NamedTemporaryFile(
                        delete=False, suffix=".zip"
                ) as f, tqdm(
                    total=total, unit="B", unit_scale=True, desc=desc, ncols=80
                ) as pbar:
                    for chunk in r.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                            pbar.update(len(chunk))
                    zip_path = f.name
            self.log(f"✅ 下载完成，解压中...")
            if os.path.exists(extract_dir):
                shutil.rmtree(extract_dir)
            with zipfile.ZipFile(zip_path, "r") as zip_ref, tqdm(
                    total=len(zip_ref.infolist()), desc="⏳ 解压进度", ncols=80
            ) as pbar:
                for member in zip_ref.infolist():
                    zip_ref.extract(member, extract_dir)
                    pbar.update(1)
            os.remove(zip_path)
            self.log(f"✅ 解压完成，路径：{extract_dir}")
            return extract_dir
        except Exception as e:
            self.log(f"❌ {desc}失败：{e}", "error")
            return None

    def _download_file(
            self, url: str, save_path: str, desc: str = "⏳ 下载进度"
    ) -> Optional[str]:
        """
        通用文件下载方法

        :param url: 下载地址
        :param save_path: 保存路径
        :param desc: 进度描述
        :return: 文件路径
        """
        try:
            with requests.get(url, stream=True, timeout=60) as r:
                r.raise_for_status()
                total = int(r.headers.get("content-length", 0))
                os.makedirs(self.download_dir, exist_ok=True)
                with open(save_path, "wb") as f, tqdm(
                        total=total, unit="B", unit_scale=True, desc=desc, ncols=80
                ) as pbar:
                    for chunk in r.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                            pbar.update(len(chunk))
            self.log(f"✅ {desc}完成，路径：{save_path}")
            return save_path
        except Exception as e:
            self.log(f"❌ {desc}失败：{e}", "error")
            return None

    def download(
            self,
            revision: Optional[str] = None,
            system: Optional[Literal["windows", "linux", "darwin"]] = None,
            return_chrome_path: bool = False,
    ) -> Optional[str]:
        """
        下载 Chrome 浏览器

        :param revision: 指定 Chromium 快照版本号，默认为最新版本
        :param system: 指定系统，默认为自动识别
        :param return_chrome_path: 是否返回 chrome 可执行文件路径
        :return: 下载路径或 chrome 可执行文件路径
        """
        sys_name: str = (system or self.system).lower()

        # 优先自定义源
        if self.custom_source:
            self.log("🔍 使用自定义源进行下载")
            custom_rev = revision or self.get_latest_chromium_revision(sys_name)
            if custom_rev:
                path = self.download_chromium(custom_rev, sys_name)
                if path:
                    return self._handle_download_result(path, return_chrome_path)

        # 官方主源
        path = self.download_chromium(revision, sys_name)
        if path:
            return self._handle_download_result(path, return_chrome_path)

        self.log(f"❌ 下载失败！不支持的系统：{sys_name}", "error")
        return None

    def _handle_download_result(
            self, path: str, return_chrome_path: bool
    ) -> Optional[str]:
        """
        处理下载结果，返回路径或chrome可执行文件路径

        :param path: 下载或解压路径
        :param return_chrome_path: 是否返回chrome可执行文件路径
        :return: 路径或chrome可执行文件路径
        """
        if return_chrome_path:
            chrome_path = self._find_chrome_exe(path)
            if chrome_path and os.path.exists(chrome_path):
                self.log(f"✅ 找到 chrome 可执行文件：{chrome_path}")
                return chrome_path

            self.log(
                f"❌ 未找到 chrome 的可执行文件，请前往存放路径手动查找：{path}",
                "warning",
            )
            return None

        self.log(f"✅ Chrome下载完成，存放路径：{path}")
        return path
