"""Download Python embeddable package to a specific directory."""

from __future__ import annotations

import argparse
import logging
import platform
import shutil
import time
import zipfile
from pathlib import Path
from typing import Final
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

# Architecture mapping
ARCH_MAP: Final = {
    "AMD64": "amd64",
    "x86_64": "amd64",
    "x86": "amd64",
    "i686": "amd64",
    "i386": "amd64",
    "ARM64": "arm64",
    "aarch64": "arm64",
}

USER_AGENT: Final = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"

logging.basicConfig(level=logging.INFO, format="%(message)s")

cwd = Path.cwd()
logger = logging.getLogger(__name__)


def get_system_arch(manual_arch: str | None = None) -> str:
    """Get system architecture for Python embeddable package."""
    if manual_arch:
        logger.debug(f"Using manual architecture: {manual_arch}")
        return manual_arch

    machine = platform.machine().upper()

    logger.debug(f"System machine: {machine}")

    if machine in ARCH_MAP:
        arch = ARCH_MAP[machine]
        logger.debug(f"Detected architecture: {arch}")
        return arch

    logger.warning(f"Unknown architecture: {machine}, defaulting to amd64")
    return "amd64"


def get_download_url(base_url_template: str, version: str, arch: str) -> str:
    """Get download URL with correct architecture."""
    return base_url_template.format(version=version, arch=arch)


def get_official_url_template(arch: str) -> str:
    """Get official URL template based on architecture."""
    return f"https://www.python.org/ftp/python/{{version}}/python-{{version}}-embed-{arch}.zip"


def get_latest_patch_version(major_minor: str, arch: str, timeout: int = 5) -> str:
    """Get the latest patch version for a given major.minor version.

    Args:
        major_minor: Major.minor version (e.g., '3.13') or full version (e.g., '3.13.1')
        arch: Architecture (amd64 or arm64)
        timeout: Request timeout in seconds

    Returns:
        Full version string (e.g., '3.13.1') or original if already complete
    """
    # Check if version is already complete (has 3 or more parts)
    version_parts = major_minor.split(".")
    if len(version_parts) >= 3:
        logger.debug(f"Version already complete: {major_minor}, skipping auto-completion")
        return major_minor

    # Check if version format is valid
    if "." not in major_minor or len(version_parts) != 2:
        logger.debug(f"Invalid version format: {major_minor}, using as: `3.13.5`")
        return major_minor

    logger.info(f"Checking latest version for {major_minor}...")

    base_url = get_official_url_template(arch)
    major, minor = major_minor.split(".")

    # Try from highest patch number down to 0
    for patch in range(20, -1, -1):
        version = f"{major}.{minor}.{patch}"
        test_url = base_url.format(version=version, arch=arch)

        req = Request(test_url, headers={"User-Agent": USER_AGENT}, method="HEAD")
        try:
            with urlopen(req, timeout=timeout) as response:
                if response.status == 200:
                    logger.info(f"Latest version found: {version}")
                    return version
        except HTTPError as e:
            if e.code == 404:
                logger.debug(f"Version {version} not found")
            else:
                logger.debug(f"Version {version}: HTTP {e.code}")
        except (URLError, TimeoutError):
            logger.debug(f"Version {version}: connection failed")

    logger.warning(f"Could not find any patch version for {major_minor}, using as-is")
    return major_minor


def get_mirror_url_templates(arch: str) -> list[str]:
    """Get mirror URL templates based on architecture."""
    return [
        f"https://mirrors.huaweicloud.com/python/{{version}}/python-{{version}}-embed-{arch}.zip",
        f"https://mirrors.aliyun.com/python-release/{{version}}/python-{{version}}-embed-{arch}.zip",
        f"https://mirrors.tuna.tsinghua.edu.cn/python/{{version}}/python-{{version}}-embed-{arch}.zip",
        f"https://mirrors.pku.edu.cn/python/{{version}}/python-{{version}}-embed-{arch}.zip",
    ]


def test_url_speed(url: str, timeout: int = 5) -> float:
    """Test URL speed by making a HEAD request and measuring response time.

    Falls back to GET request if HEAD is not supported.
    """
    start_time = time.time()

    req = Request(url, headers={"User-Agent": USER_AGENT}, method="HEAD")
    try:
        with urlopen(req, timeout=timeout) as response:
            if response.status == 200:
                elapsed = time.time() - start_time
                logger.debug(f"URL {url}: {elapsed:.3f}s")
                return elapsed
    except HTTPError as e:
        # Some servers don't support HEAD, fallback to GET with range request
        if e.code == 405 or e.code == 403:
            req = Request(
                url,
                headers={"User-Agent": USER_AGENT, "Range": "bytes=0-1023"},
                method="GET",
            )
            try:
                with urlopen(req, timeout=timeout) as response:
                    if response.status == 206 or response.status == 200:
                        elapsed = time.time() - start_time
                        logger.debug(f"URL {url}: {elapsed:.3f}s (GET fallback)")
                        return elapsed
            except (URLError, TimeoutError):
                logger.debug(f"URL {url}: failed (GET fallback)")
        else:
            logger.debug(f"URL {url}: failed (HTTP {e.code})")
    except (URLError, TimeoutError) as e:
        logger.debug(f"URL {url}: failed ({type(e).__name__})")

    return float("inf")


def find_available_urls(version: str, arch: str, timeout: int = 5) -> list[str]:
    """Find all available URLs for downloading Python embeddable package, sorted by speed."""
    official_url = get_official_url_template(arch).format(version=version)
    mirror_templates = get_mirror_url_templates(arch)
    mirror_urls = [template.format(version=version) for template in mirror_templates]
    urls_to_test = [*mirror_urls, official_url]

    logger.debug("Testing mirror speeds...")

    speed_results = []
    for url in urls_to_test:
        speed = test_url_speed(url, timeout)
        if speed != float("inf"):
            speed_results.append((speed, url))
            logger.debug(f"  ✓ {url} ({speed:.3f}s)")
        else:
            logger.debug(f"  ✗ {url} (failed)")

    if not speed_results:
        logger.warning("All mirrors failed, falling back to official URL")
        return [official_url]

    speed_results.sort(key=lambda x: x[0])
    sorted_urls = sorted(speed_results, key=lambda x: x[0])
    available_urls = [url for _, url in sorted_urls]

    logger.debug(f"Available URLs (sorted by speed): {len(available_urls)} found")
    for i, url in enumerate(available_urls, 1):
        logger.debug(f"  {i}. {url}")

    return available_urls


def download_with_progress(url: str, dest_path: Path) -> None:
    """Download file with progress bar and integrity check."""
    logger.info(f"Downloading from: {url}")

    req = Request(url, headers={"User-Agent": USER_AGENT})
    try:
        with urlopen(req) as response:
            total_size = int(response.headers.get("content-length", 0))
            downloaded = 0
            last_progress = 0

            dest_path.parent.mkdir(parents=True, exist_ok=True)

            with dest_path.open("wb") as f:
                while True:
                    chunk = response.read(65536)  # Larger chunk size for better performance
                    if not chunk:
                        break
                    f.write(chunk)
                    downloaded += len(chunk)

                    if total_size > 0:
                        progress = (downloaded / total_size) * 100
                        # Update progress when percentage changes (every 1%)
                        if int(progress) > last_progress or downloaded == total_size:
                            logger.info(
                                f"Progress: {downloaded / 1024 / 1024:.2f} MB / {total_size / 1024 / 1024:.2f} MB ({progress:.1f}%)"
                            )
                            last_progress = int(progress)

            # Verify downloaded file size
            if total_size > 0 and downloaded != total_size:
                raise URLError(f"Incomplete download: expected {total_size} bytes, got {downloaded} bytes")

            logger.info(f"Download completed: {downloaded / 1024 / 1024:.2f} MB")
    except URLError as e:
        logger.error(f"Failed to download {url}: {e}")
        raise


def extract_zip(zip_path: Path, dest_dir: Path) -> None:
    """Extract zip file to destination directory using optimized shutil method."""
    logger.info(f"Extracting to: {dest_dir}")

    dest_dir.mkdir(parents=True, exist_ok=True)

    # Use shutil.unpack_archive for better performance (uses optimized C implementation)
    shutil.unpack_archive(str(zip_path), str(dest_dir), "zip")

    # Count extracted files for logging
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        member_count = len(zip_ref.namelist())
    logger.info(f"Extracted {member_count} files")


def install_embed_python(
    target_dir: Path,
    version: str,
    cache_dir: Path,
    offline: bool,
    keep_cache: bool,
    skip_speed_test: bool,
    arch: str,
    timeout: int = 5,
) -> bool:
    """Install Python embeddable package to target directory."""
    version_filename = f"python-{version}-embed-{arch}.zip"
    cache_file = cache_dir / version_filename

    if cache_file.exists():
        logger.debug(f"Using cached file: {cache_file}")
    elif offline:
        logger.error(f"Offline mode: no cached file found for version {version}")
        return False
    else:
        if skip_speed_test:
            url = get_official_url_template(arch).format(version=version)
            logger.info(f"Skipping speed test, using official URL: {url}")
            urls = [url]
        else:
            urls = find_available_urls(version, arch, timeout)

        download_success = False
        for i, url in enumerate(urls, 1):
            try:
                logger.info(f"Attempting download {i}/{len(urls)}: {url}")
                download_with_progress(url, cache_file)
                download_success = True
                break
            except URLError as e:
                logger.warning(f"Download failed from {url}: {e}")
                if i < len(urls):
                    logger.info("Retrying with next fastest URL...")
                else:
                    logger.error(
                        f"Failed to download Python {version} ({arch}) from all available URLs. "
                        f"Please check your internet connection and version."
                    )
                    return False

        if not download_success:
            return False

    try:
        extract_zip(cache_file, target_dir)
    except zipfile.BadZipFile:
        logger.error(f"Invalid zip file: {cache_file}")
        if not keep_cache:
            cache_file.unlink()
        return False
    except Exception as e:
        logger.error(f"Failed to extract: {e}")
        return False

    if not keep_cache and not offline:
        logger.debug(f"Cleaning up cache: {cache_file}")
        cache_file.unlink()

    python_exe = target_dir / "python.exe"
    if python_exe.exists():
        logger.info(f"Successfully installed Python {version} ({arch}) embeddable package to: {target_dir}")
        logger.info(f"Python executable: {python_exe}")
        return True

    logger.warning(f"Installation completed but python.exe not found at: {python_exe}")
    return True


def main() -> None:
    parser = argparse.ArgumentParser(
        prog="embedinstall",
        description="Download and install Python embeddable package to a specific directory.",
    )
    parser.add_argument(
        "--directory",
        "-D",
        type=str,
        default=str(cwd / "runtime"),
        help="Directory to install Python embeddable package (default: current directory)",
    )
    parser.add_argument("--debug", "-d", action="store_true", help="Enable debug mode")
    parser.add_argument(
        "--offline",
        "-o",
        action="store_true",
        help="Offline mode (use cached files only)",
    )
    parser.add_argument(
        "--version",
        "-v",
        type=str,
        default="3.8.10",
        help="Python version to install (default: 3.8.10)",
    )
    parser.add_argument(
        "--cache-dir",
        "-C",
        type=str,
        default=str(cwd / ".cache" / "embedinstall"),
        help="Cache directory for downloaded files",
    )
    parser.add_argument(
        "--keep-cache",
        "-k",
        action="store_true",
        default=True,
        help="Keep downloaded cache files",
    )
    parser.add_argument(
        "--skip-speed-test",
        "-s",
        action="store_true",
        help="Skip speed test and use official URL directly",
    )
    parser.add_argument(
        "--arch",
        "-a",
        type=str,
        help="Manually specify architecture (amd64, arm64). Default: auto-detect",
    )
    parser.add_argument(
        "--timeout",
        "-t",
        type=int,
        default=5,
        help="Timeout in seconds for URL speed test (default: 5)",
    )

    args = parser.parse_args()

    if args.debug:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)

    target_dir = Path(args.directory)
    cache_dir = Path(args.cache_dir)

    t0 = time.perf_counter()
    arch = get_system_arch(args.arch)

    # Auto-complete version if only major.minor is provided
    version = get_latest_patch_version(args.version, arch, args.timeout)
    if "." not in version:
        return

    logger.info(f"Installing Python {version} ({arch}) to: {target_dir}")
    logger.info(f"Cache directory: {cache_dir}")

    success = install_embed_python(
        target_dir=target_dir,
        version=version,
        cache_dir=cache_dir,
        offline=args.offline,
        keep_cache=args.keep_cache,
        skip_speed_test=args.skip_speed_test,
        arch=arch,
        timeout=args.timeout,
    )

    if not success:
        exit(1)

    logger.info(f"Installation completed in {time.perf_counter() - t0:.4f} seconds")
