from __future__ import annotations

import argparse
import logging
import os
import subprocess
from pathlib import Path
from typing import Final

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

# Conda mirror URLs
_CONDA_MIRROR_URLS: Final[dict[str, frozenset[str]]] = {
    "tsinghua": frozenset([
        "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/",
        "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/",
        "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r/",
        "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2/",
        "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/pro/",
        "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/",
        "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/bioconda/",
        "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/menpo/",
        "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/",
    ]),
    "ustc": frozenset([
        "https://mirrors.ustc.edu.cn/anaconda/pkgs/main/",
        "https://mirrors.ustc.edu.cn/anaconda/pkgs/free/",
        "https://mirrors.ustc.edu.cn/anaconda/pkgs/r/",
        "https://mirrors.ustc.edu.cn/anaconda/pkgs/msys2/",
        "https://mirrors.ustc.edu.cn/anaconda/pkgs/pro/",
        "https://mirrors.ustc.edu.cn/anaconda/pkgs/dev/",
        "https://mirrors.ustc.edu.cn/anaconda/cloud/conda-forge/",
        "https://mirrors.ustc.edu.cn/anaconda/cloud/bioconda/",
        "https://mirrors.ustc.edu.cn/anaconda/cloud/menpo/",
        "https://mirrors.ustc.edu.cn/anaconda/cloud/pytorch/",
    ]),
    "bsfu": frozenset([
        "https://mirrors.bsfu.edu.cn/anaconda/pkgs/main/",
        "https://mirrors.bsfu.edu.cn/anaconda/pkgs/free/",
        "https://mirrors.bsfu.edu.cn/anaconda/pkgs/r/",
        "https://mirrors.bsfu.edu.cn/anaconda/pkgs/msys2/",
        "https://mirrors.bsfu.edu.cn/anaconda/pkgs/pro/",
        "https://mirrors.bsfu.edu.cn/anaconda/pkgs/dev/",
        "https://mirrors.bsfu.edu.cn/anaconda/cloud/conda-forge/",
        "https://mirrors.bsfu.edu.cn/anaconda/cloud/bioconda/",
        "https://mirrors.bsfu.edu.cn/anaconda/cloud/menpo/",
        "https://mirrors.bsfu.edu.cn/anaconda/cloud/pytorch/",
    ]),
    "aliyun": frozenset([
        "https://mirrors.aliyun.com/anaconda/pkgs/main/",
        "https://mirrors.aliyun.com/anaconda/pkgs/free/",
        "https://mirrors.aliyun.com/anaconda/pkgs/r/",
        "https://mirrors.aliyun.com/anaconda/pkgs/msys2/",
        "https://mirrors.aliyun.com/anaconda/pkgs/pro/",
        "https://mirrors.aliyun.com/anaconda/pkgs/dev/",
        "https://mirrors.aliyun.com/anaconda/cloud/conda-forge/",
        "https://mirrors.aliyun.com/anaconda/cloud/bioconda/",
        "https://mirrors.aliyun.com/anaconda/cloud/menpo/",
        "https://mirrors.aliyun.com/anaconda/cloud/pytorch/",
    ]),
}


def set_conda_mirror(mirror: str = "tsinghua") -> None:
    """Set the Conda mirror for the given channel.

    Args:
        mirror: Mirror name (tsinghua, ustc, bsfu, or aliyun)
    """
    if mirror not in _CONDA_MIRROR_URLS:
        logger.error(f"Invalid mirror: {mirror}")
        return

    old_config = Path.home() / ".condarc"
    if old_config.exists():
        logger.info("Found existing .condarc file, backing it up")
        os.rename(old_config, Path.home() / ".condarc.bak")
    else:
        logger.debug("No existing .condarc file found")

    mirror_urls = _CONDA_MIRROR_URLS[mirror]
    for url in mirror_urls:
        logger.debug(f"Adding mirror: {url}")
        try:
            subprocess.run(["conda", "config", "--add", "channels", url], check=True)
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to add mirror {url}: {e}")
            return

    try:
        subprocess.run(
            ["conda", "config", "--set", "show_channel_urls", "yes"], check=True
        )
        logger.info("Conda mirror set successfully")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to set show_channel_urls: {e}")


def parse_args() -> argparse.Namespace:
    """Parse command line arguments.

    Returns:
        Parsed arguments
    """
    parser = argparse.ArgumentParser(description="Setup Conda environment for SFI")
    parser.add_argument(
        "mirror",
        type=str,
        choices=["tsinghua", "ustc", "bsfu", "aliyun"],
        nargs="?",
        default="tsinghua",
        help="Conda mirror to use",
    )
    parser.add_argument("-d", "--debug", help="Debug mode", action="store_true")

    return parser.parse_args()


def main() -> None:
    """Main entry point for condasetup CLI."""
    args = parse_args()

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

    set_conda_mirror(args.mirror)


if __name__ == "__main__":
    main()
