import os
 
# 国内常用的PIP源
SOURCES = {
    "阿里云": "https://mirrors.aliyun.com/pypi/simple/",
    "清华大学": "https://pypi.tuna.tsinghua.edu.cn/simple/",
    "豆瓣": "https://pypi.doubanio.com/simple/",
    "中国科学技术大学": "https://pypi.mirrors.ustc.edu.cn/simple/",
    "官方源": "https://pypi.org/simple/"
}
 
 
def get_pip_config_path():
    """获取pip配置文件路径"""
    user_home = os.path.expanduser("~")
 
    if os.name == "nt":  # Windows
        pip_config_dir = os.path.join(user_home, "pip")
        pip_config_file = os.path.join(pip_config_dir, "pip.ini")
    else:  # Linux / macOS
        pip_config_dir = os.path.join(user_home, ".pip")
        pip_config_file = os.path.join(pip_config_dir, "pip.conf")
 
    return pip_config_dir, pip_config_file
 
 
def set_pip_source(source_url):
    """设置pip源"""
    config_dir, config_file = get_pip_config_path()
 
    if not os.path.exists(config_dir):
        os.makedirs(config_dir)
 
    host = source_url.split("//")[1].split("/")[0]
 
    with open(config_file, "w", encoding="utf-8") as f:
        f.write("[global]\n")
        f.write(f"index-url = {source_url}\n\n")
        f.write("[install]\n")
        f.write("trusted-host =\n")
        f.write(f"    {host}\n")
 
    print(f"✅ 成功设置PIP源为: {source_url}")
    print(f"配置文件已保存至: {config_file}")
 
 
def show_current_source():
    """显示当前pip源"""
    _, config_file = get_pip_config_path()
 
    if not os.path.exists(config_file):
        print("⚠️ 未找到pip配置文件，当前使用默认源")
        return
 
    try:
        with open(config_file, "r", encoding="utf-8") as f:
            content = f.read()
            for line in content.splitlines():
                if line.strip().startswith("index-url"):
                    source = line.split("=")[1].strip()
                    print(f"当前PIP源: {source}")
                    return
        print("⚠️ 配置文件中未找到源信息，当前使用默认源")
    except Exception as e:
        print(f"读取配置文件出错: {e}")
 
 
def main():
    print("===== PIP源一键修改工具 =====")
    print("当前源地址:")
    show_current_source()
 
    print("\n请选择要设置的源:")
    for i, name in enumerate(SOURCES.keys(), 1):
        print(f"{i}. {name}")
 
    try:
        choice = int(input("\n请输入序号 (1-5): "))
        if 1 <= choice <= len(SOURCES):
            source_name = list(SOURCES.keys())[choice - 1]  # &#9989; 修复这里
            source_url = SOURCES[source_name]
            set_pip_source(source_url)
        else:
            print("输入有误，请输入1到5之间的数字")
    except ValueError:
        print("输入有误，请输入数字")
    except Exception as e:
        print(f"操作出错: {e}")
 
 
if __name__ == "__main__":
    main()