#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
换源脚本：支持以下发行版与架构
- Debian 11 (bullseye), Debian 12 (bookworm) — amd64/arm64
- Ubuntu 22.04 (jammy), Ubuntu 24.04 (noble) — amd64/arm64

APT 镜像选项：official（官方）、aliyun（阿里云）、tsinghua（清华）

pip 永久换源选项：tsinghua、aliyun、tencent、douban、default(恢复官方)

用法示例：
	APT 交互：     sudo python3 changesource.py
	APT 指定：     sudo python3 changesource.py --mirror aliyun
	APT 仅查看：   python3 changesource.py --mirror tsinghua --dry-run
	APT 写入更新： sudo python3 changesource.py --mirror official --update

	pip 交互：     python3 changesource.py --pip-only
	pip 指定：     python3 changesource.py --pip-mirror tsinghua
	pip 恢复默认： python3 changesource.py --pip-mirror default

注意：写入 /etc/apt/sources.list 与执行 apt-get update 需要 root 权限；Termux 下不需要 root。
"""

from __future__ import annotations

import argparse
import datetime
import os
import platform
import re
import shutil
import subprocess
import sys
from typing import Dict, Tuple, List


SUPPORTED = {
	"debian": {"11": "bullseye", "12": "bookworm"},
	"ubuntu": {"22.04": "jammy", "24.04": "noble"},
}


# Termux 检测
def is_termux() -> bool:
	prefix = os.environ.get("PREFIX", "")
	if "com.termux" in prefix:
		return True
	if os.environ.get("TERMUX_VERSION"):
		return True
	# 兜底：常见默认 PREFIX
	if prefix == "/data/data/com.termux/files/usr":
		return True
	return False

def get_termux_prefix() -> str:
	return os.environ.get("PREFIX", "/data/data/com.termux/files/usr")


#阅读系统信息
def read_os_release() -> Dict[str, str]:
	data: Dict[str, str] = {}
	try:
		with open("/etc/os-release", "r", encoding="utf-8") as f:
			for line in f:
				line = line.strip()
				if not line or line.startswith("#") or "=" not in line:
					continue
				k, v = line.split("=", 1)
				v = v.strip().strip('"').strip("'")
				data[k] = v
	except Exception:
		pass
	return data

#规范化架构名称
def normalize_arch(uname_arch: str) -> str:
	a = uname_arch.lower()
	if a in ("x86_64", "amd64"):
		return "amd64"
	if a in ("aarch64", "arm64"):
		return "arm64"
	# 其它架构暂不写入 sources 约束（APT 不必声明架构），但用于提示
	return a


#获取发行版版本信息
def get_distro_info() -> Tuple[str, str, str]:
	"""返回 (id, version_id, codename)
	id: debian/ubuntu
	version_id: 如 '11', '12', '22.04', '24.04'
	codename: bullseye/bookworm/jammy/noble
	"""
	info = read_os_release()
	distro_id = info.get("ID", "").lower()
	version_id = info.get("VERSION_ID", "")
	codename = info.get("VERSION_CODENAME", "").lower()

	# 一些派生可能提供 UBUNTU_CODENAME
	if not codename:
		codename = info.get("UBUNTU_CODENAME", "").lower()

	# 规范化版本格式
	if distro_id == "debian":
		# Debian 通常为 '11' 或 '12'
		version_id = version_id.split(".")[0]
		if not codename and version_id in SUPPORTED["debian"]:
			codename = SUPPORTED["debian"][version_id]
	elif distro_id == "ubuntu":
		# Ubuntu 保留小版本以匹配 22.04 / 24.04
		m = re.match(r"(\d{2})\.(\d{2})", version_id or "")
		if m:
			version_id = f"{m.group(1)}.{m.group(2)}"
		if not codename and version_id in SUPPORTED["ubuntu"]:
			codename = SUPPORTED["ubuntu"][version_id]

	return distro_id, version_id, codename

#验证发行版与版本支持
def validate_supported(distro_id: str, version_id: str, codename: str) -> Tuple[bool, str]:
	if distro_id not in SUPPORTED:
		return False, f"不支持的发行版: {distro_id}"
	if distro_id == "debian":
		if version_id not in SUPPORTED[distro_id]:
			return False, f"不支持的 Debian 版本: {version_id}（仅支持 11/12）"
		expect = SUPPORTED[distro_id][version_id]
		if codename != expect:
			return False, f"版本代号不匹配: 期望 {expect}, 实际 {codename or '未知'}"
	elif distro_id == "ubuntu":
		if version_id not in SUPPORTED[distro_id]:
			return False, f"不支持的 Ubuntu 版本: {version_id}（仅支持 22.04/24.04）"
		expect = SUPPORTED[distro_id][version_id]
		if codename != expect:
			return False, f"版本代号不匹配: 期望 {expect}, 实际 {codename or '未知'}"
	return True, ""


# ---- Termux 支持（仅清华源）----
def _termux_apply_mirror_to_file(path: str, pattern: str, new_line: str, dry_run: bool) -> bool:
	"""在给定文件中，将匹配 pattern 的行注释掉并在下一行追加 new_line；如果未匹配，且文件存在且不包含 new_line，则在末尾追加。
	返回是否发生变更。
	"""
	if not os.path.exists(path):
		return False

	try:
		with open(path, "r", encoding="utf-8") as f:
			lines: List[str] = f.read().splitlines()
	except Exception:
		# 读取失败视为不变更
		return False

	import re as _re
	changed = False
	out_lines: List[str] = []
	matched_once = False
	for line in lines:
		m = _re.match(pattern, line)
		if m:
			matched_once = True
			if not dry_run:
				out_lines.append("#" + line)
				out_lines.append(new_line)
			changed = True
		else:
			out_lines.append(line)

	if not matched_once:
		# 未匹配时，如果没有新行，则追加
		if new_line not in lines:
			if not dry_run:
				if out_lines and out_lines[-1].strip():
					out_lines.append("")
				out_lines.append("# added by changesource.py")
				out_lines.append(new_line)
			changed = True

	if changed and not dry_run:
		# 备份并写回
		backup_file(path)
		with open(path, "w", encoding="utf-8") as f:
			f.write("\n".join(out_lines) + "\n")

	return changed

# Termux 切换到清华源
def termux_switch_to_tsinghua(dry_run: bool, update: bool, assume_yes: bool) -> int:
	prefix = get_termux_prefix()
	main_path = os.path.join(prefix, "etc/apt/sources.list")
	x11_path = os.path.join(prefix, "etc/apt/sources.list.d/x11.list")
	root_path = os.path.join(prefix, "etc/apt/sources.list.d/root.list")

	# 与清华源教程一致的替换模式
	main_pat = r"^(deb.*stable main)$"
	main_new = "deb https://mirrors.tuna.tsinghua.edu.cn/termux/apt/termux-main stable main"

	x11_pat = r"^(deb.*x11 main)$"
	x11_new = "deb https://mirrors.tuna.tsinghua.edu.cn/termux/apt/termux-x11 x11 main"

	root_pat = r"^(deb.*root main)$"
	root_new = "deb https://mirrors.tuna.tsinghua.edu.cn/termux/apt/termux-root root main"

	print("[Termux] 目标镜像：清华 TUNA")
	print(f"[Termux] PREFIX={prefix}")

	# 确认
	if not dry_run and not assume_yes:
		ans = input("确认为 Termux 写入清华镜像源？[y/N]: ").strip().lower()
		if ans not in ("y", "yes"):
			print("已取消。")
			return 0


	# 执行替换
	changed_any = False
	for p, pat, newl in [
		(main_path, main_pat, main_new),
		(x11_path, x11_pat, x11_new),
		(root_path, root_pat, root_new),
	]:
		if os.path.exists(p):
			print(f"[Termux] 处理 {p}")
			changed = _termux_apply_mirror_to_file(p, pat, newl, dry_run)
			changed_any = changed_any or changed
		else:
			print(f"[Termux] 跳过（不存在）：{p}")

	if dry_run:
		print("[Termux] dry-run: 未实际写入。")
		return 0

	if update:
		# 运行 apt update / upgrade
		try:
			rc1 = subprocess.run(["apt", "update"], check=False).returncode
			if rc1 != 0:
				print(f"apt update 失败，返回码 {rc1}")
				return rc1
			cmd = ["apt", "upgrade"]
			if assume_yes:
				cmd.insert(2, "-y")
			rc2 = subprocess.run(cmd, check=False).returncode
			if rc2 != 0:
				print(f"apt upgrade 失败，返回码 {rc2}")
				return rc2
		except FileNotFoundError:
			print("未找到 apt，请确认处于 Termux 环境。", file=sys.stderr)
			return 127

	print("[Termux] 已完成清华源处理。")
	return 0


#渲染 Debian 源列表
def render_debian_sources(codename: str, mirror: str) -> str:
	# 组件：Debian 12 含 non-free-firmware
	components = "main contrib non-free"
	if codename == "bookworm":
		components = "main contrib non-free non-free-firmware"

	if mirror == "official":
		base = "http://deb.debian.org/debian"
		sec = "http://security.debian.org/debian-security"
	elif mirror == "aliyun":
		base = "https://mirrors.aliyun.com/debian"
		sec = "https://mirrors.aliyun.com/debian-security"
	elif mirror == "tsinghua":
		base = "https://mirrors.tuna.tsinghua.edu.cn/debian"
		sec = "https://mirrors.tuna.tsinghua.edu.cn/debian-security"
	else:
		raise ValueError(f"未知镜像: {mirror}")

	lines = [
		f"deb {base} {codename} {components}",
		f"deb {sec} {codename}-security {components}",
		f"deb {base} {codename}-updates {components}",
	]
	return "\n".join(lines) + "\n"

#渲染 Ubuntu 源列表
def render_ubuntu_sources(codename: str, mirror: str) -> str:
	if mirror == "official":
		base = "http://archive.ubuntu.com/ubuntu"
		sec = "http://security.ubuntu.com/ubuntu"
	elif mirror == "aliyun":
		base = sec = "https://mirrors.aliyun.com/ubuntu"
	elif mirror == "tsinghua":
		base = sec = "https://mirrors.tuna.tsinghua.edu.cn/ubuntu"
	else:
		raise ValueError(f"未知镜像: {mirror}")

	components = "main restricted universe multiverse"
	lines = [
		f"deb {base} {codename} {components}",
		f"deb {base} {codename}-updates {components}",
		f"deb {base} {codename}-backports {components}",
		f"deb {sec} {codename}-security {components}",
	]
	return "\n".join(lines) + "\n"

#根据发行版渲染源列表
def render_sources(distro_id: str, codename: str, mirror: str) -> str:
	if distro_id == "debian":
		return render_debian_sources(codename, mirror)
	elif distro_id == "ubuntu":
		return render_ubuntu_sources(codename, mirror)
	else:
		raise ValueError(f"不支持的发行版: {distro_id}")

#确保以 root 权限运行
def ensure_root(for_what: str) -> None:
	if os.geteuid() != 0:
		print(f"[需要 root] {for_what} 请使用 sudo 运行此脚本。", file=sys.stderr)
		sys.exit(1)

#备份文件
def backup_file(path: str) -> str:
	ts = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
	bak = f"{path}.bak-{ts}"
	try:
		if os.path.exists(path):
			shutil.copy2(path, bak)
	except Exception as e:
		print(f"备份 {path} 失败: {e}", file=sys.stderr)
		sys.exit(1)
	return bak

#写入 sources.list
def write_sources(content: str, path: str = "/etc/apt/sources.list") -> None:
	# 先备份，再写入
	backup_file(path)
	try:
		with open(path, "w", encoding="utf-8") as f:
			f.write(content)
	except Exception as e:
		print(f"写入 {path} 失败: {e}", file=sys.stderr)
		sys.exit(1)

#执行 apt-get update
def apt_update() -> int:
	try:
		# 与用户共享输出
		proc = subprocess.run(["apt-get", "update"], check=False)
		return proc.returncode
	except FileNotFoundError:
		print("未找到 apt-get，请确认系统为 Debian/Ubuntu。", file=sys.stderr)
		return 127

#交互式选择镜像
def choose_mirror_interactive() -> str:
	print("\n================ Linux 软件源镜像切换 ================")
	print("请选择要切换的镜像源：")
	options = [
		("official", "默认官方源"),
		("aliyun", "阿里云"),
		("tsinghua", "清华源"),
	]
	print("  0. 跳过（不更改）")
	for idx, (_, label) in enumerate(options, start=1):
		print(f"  {idx}. {label}")
	raw = input("输入编号 (默认 1): ").strip()
	if not raw:
		return options[0][0]
	try:
		i = int(raw)
		if i == 0:
			return "skip"
		if 1 <= i <= len(options):
			return options[i - 1][0]
	except Exception:
		pass
	print("输入无效，默认选择 1. 默认官方源。")
	return options[0][0]

# 统一交互主面板
def show_main_menu() -> tuple[str, str | None]:
	"""
	显示统一交互面板并返回用户选择。
	返回 (mode, value)
	- mode: 'apt' 或 'pip' 或 'skip'
	- value: 对于 apt，为 'ubuntu'|'debian'|'termux'；对于 pip，为 'tsinghua'|'aliyun'|'tencent'|'default'
	"""
	print("============请选择需要换源的命令============")
	print("Linux发行版软件源：")
	print("1.Ubuntu（支持22 24 amd64 arm64 官方源 清华源 阿里源）")
	print("2.Debian（支持11 12 amd64 arm64 官方源 清华源 阿里源）")
	print("3.Termux（支持清华源）")
	print()
	print("Python的pip镜像源：")
	print("a.清华源")
	print("b.阿里源")
	print("c.腾讯源")
	print("d.官方源")
	print("===========================================")

	sel = input("请输入选项编号（1/2/3 或 a/b/c/d，其他跳过）：").strip().lower()
	if sel == "1":
		return ("apt", "ubuntu")
	if sel == "2":
		return ("apt", "debian")
	if sel == "3":
		return ("apt", "termux")
	if sel == "a":
		return ("pip", "tsinghua")
	if sel == "b":
		return ("pip", "aliyun")
	if sel == "c":
		return ("pip", "tencent")
	if sel == "d":
		return ("pip", "default")
	return ("skip", None)

#解析命令行参数
def parse_args() -> argparse.Namespace:
	p = argparse.ArgumentParser(description="APT 换源脚本")
	p.add_argument("--mirror", choices=["official", "aliyun", "tsinghua"], help="选择镜像源")
	p.add_argument("--dry-run", action="store_true", help="仅打印将要写入的 sources.list，不实际写入")
	p.add_argument("--update", action="store_true", help="写入后执行 apt-get update")
	p.add_argument("-y", "--yes", action="store_true", help="不提示，直接写入")
	p.add_argument("--path", default="/etc/apt/sources.list", help="sources.list 路径（默认 /etc/apt/sources.list）")
	# pip 相关
	p.add_argument("--pip-mirror", choices=["tsinghua", "aliyun", "tencent", "douban", "default"], help="设置 pip 全局镜像（default 为恢复官方）")
	p.add_argument("--pip-only", action="store_true", help="仅执行 pip 换源，不进行 APT 换源")
	return p.parse_args()


# pip 镜像映射
PIP_MIRRORS: Dict[str, str | None] = {
	"tsinghua": "https://pypi.tuna.tsinghua.edu.cn/simple",
	"aliyun": "https://mirrors.aliyun.com/pypi/simple/",
	"tencent": "http://mirrors.cloud.tencent.com/pypi/simple",
	"douban": "http://pypi.douban.com/simple/",
	"default": None,  # unset
}


def choose_pip_mirror_interactive() -> str:
	print("\n=============== Python 的 pip 镜像源切换 ===============")
	print("请选择 pip 镜像：")
	options = [
		("tsinghua", "清华 TUNA"),
		("aliyun", "阿里云"),
		("tencent", "腾讯云"),
		("douban", "豆瓣"),
		("default", "恢复官方默认"),
	]
	print("  0. 跳过（不更改）")
	for i, (_, label) in enumerate(options, 1):
		print(f"  {i}. {label}")
	raw = input("输入编号 (默认 1): ").strip()
	if not raw:
		return options[0][0]
	try:
		idx = int(raw)
		if idx == 0:
			return "skip"
		if 1 <= idx <= len(options):
			return options[idx - 1][0]
	except Exception:
		pass
	print("输入无效，默认选择 1. 清华 TUNA。")
	return options[0][0]


def run_pip_config(mirror_key: str, dry_run: bool, assume_yes: bool) -> int:
	url = PIP_MIRRORS.get(mirror_key)
	py = sys.executable or "python3"
	if url:
		cmd = [py, "-m", "pip", "config", "set", "global.index-url", url]
		desc = f"pip 使用镜像: {mirror_key} -> {url}"
	else:
		cmd = [py, "-m", "pip", "config", "unset", "global.index-url"]
		desc = "pip 恢复官方默认源"
	print(f"[pip] {desc}")
	print(f"[pip] 将执行: {' '.join(cmd)}")
	if dry_run:
		print("[pip] dry-run: 未实际执行。")
		return 0
	if not assume_yes:
		ans = input("确认执行 pip 配置变更？[y/N]: ").strip().lower()
		if ans not in ("y", "yes"):
			print("已取消。")
			return 0
	try:
		rc = subprocess.run(cmd, check=False).returncode
		if rc == 0:
			print("[pip] 已完成。")
		else:
			print(f"[pip] 失败，返回码 {rc}")
		return rc
	except Exception as e:
		print(f"[pip] 执行失败: {e}", file=sys.stderr)
		return 1

#主函数
def main() -> None:
	args = parse_args()

	distro_id, version_id, codename = get_distro_info()
	arch = normalize_arch(platform.machine())

	# 在换源前输出系统信息
	print(f"检测到系统：distro={distro_id or 'unknown'}, version={version_id or 'unknown'}, codename={codename or 'unknown'}, arch={arch}")

	# 两个板块：APT 与 pip
	final_rc = 0

	# 如果未提供任何镜像参数，则进入统一交互主面板一次
	invoked_by_menu = False
	if not any([args.mirror, args.pip_mirror, args.pip_only]):
		mode, value = show_main_menu()
		invoked_by_menu = True
		if mode == "apt":
			# 将菜单选择映射到 apt 的 mirror 与环境
			if value == "termux":
				# Termux 只支持清华
				args.mirror = "tsinghua"
			# 用户意图仅为 APT，故跳过后续 pip 交互
			args.pip_mirror = "skip"
		elif mode == "pip":
			# 直接设置 pip 目标镜像并跳过 APT 流程
			args.pip_mirror = value
			args.pip_only = True
		else:
			# 完全跳过：同时跳过 APT 与 pip 的后续交互
			print("已跳过。")
			args.mirror = "skip"
			args.pip_mirror = "skip"

	# APT 板块
	if not args.pip_only:
		mirror = args.mirror or choose_mirror_interactive()
		if mirror != "skip":
			if is_termux():
				if mirror != "tsinghua":
					print("Termux 环境当前仅支持切换到清华源（tsinghua）。请使用 --mirror tsinghua 或选择跳过。", file=sys.stderr)
					final_rc = final_rc or 2
				else:
					rc = termux_switch_to_tsinghua(args.dry_run, args.update, args.yes)
					final_rc = rc or final_rc
			else:
				ok, reason = validate_supported(distro_id, version_id, codename)
				if not ok:
					print(reason, file=sys.stderr)
					final_rc = final_rc or 2
				else:
					try:
						content = render_sources(distro_id, codename, mirror)
					except Exception as e:
						print(str(e), file=sys.stderr)
						sys.exit(3)

					header = (
						f"# Generated by changesource.py on {datetime.datetime.now().isoformat(timespec='seconds')}\n"
						f"# distro={distro_id} version={version_id} codename={codename} arch={arch}\n"
						f"# mirror={mirror}\n"
					)
					content = header + content

					if args.dry_run:
						print(content)
					else:
						if os.geteuid() != 0:
							print("写入 sources.list 需要 root 权限，请使用 sudo 运行或带 --dry-run 预览。", file=sys.stderr)
							final_rc = final_rc or 1
						else:
							# 确认
							if not args.yes:
								print("将写入以下内容到:", args.path)
								print("-" * 60)
								print(content)
								print("-" * 60)
								ans = input("确认写入？[y/N]: ").strip().lower()
								if ans not in ("y", "yes"):
									print("已取消。")
								else:
									write_sources(content, args.path)
									print(f"已写入 {args.path} 并备份原文件为 .bak-时间戳。")
									if args.update:
										rc = apt_update()
										if rc == 0:
											print("apt-get update 成功。")
										else:
											print(f"apt-get update 退出码 {rc}，请检查网络或源配置。")
							else:
								write_sources(content, args.path)
								print(f"已写入 {args.path} 并备份原文件为 .bak-时间戳。")
								if args.update:
									rc = apt_update()
									if rc == 0:
										print("apt-get update 成功。")
									else:
										print(f"apt-get update 退出码 {rc}，请检查网络或源配置。")

	# pip 板块
	pip_key: str | None = args.pip_mirror
	if pip_key is None:
		# 若未提供且也未指定仅 APT，则展示 pip 板块交互
		pip_key = choose_pip_mirror_interactive()
	if pip_key and pip_key != "skip":
		rc_pip = run_pip_config(pip_key, args.dry_run, args.yes)
		final_rc = rc_pip or final_rc

	sys.exit(final_rc)

	ok, reason = validate_supported(distro_id, version_id, codename)
	if not ok:
		print(reason, file=sys.stderr)
		sys.exit(2)

	try:
		content = render_sources(distro_id, codename, mirror)
	except Exception as e:
		print(str(e), file=sys.stderr)
		sys.exit(3)

	header = (
		f"# Generated by changesource.py on {datetime.datetime.now().isoformat(timespec='seconds')}\n"
		f"# distro={distro_id} version={version_id} codename={codename} arch={arch}\n"
		f"# mirror={mirror}\n"
	)
	content = header + content

	if args.dry_run:
		print(content)
		return

	if os.geteuid() != 0:
		print("写入 sources.list 需要 root 权限，请使用 sudo 运行或带 --dry-run 预览。", file=sys.stderr)
		sys.exit(1)

	# 确认
	if not args.yes:
		print("将写入以下内容到:", args.path)
		print("-" * 60)
		print(content)
		print("-" * 60)
		ans = input("确认写入？[y/N]: ").strip().lower()
		if ans not in ("y", "yes"):
			print("已取消。")
			return

	write_sources(content, args.path)
	print(f"已写入 {args.path} 并备份原文件为 .bak-时间戳。")

	if args.update:
		rc = apt_update()
		if rc == 0:
			print("apt-get update 成功。")
		else:
			print(f"apt-get update 退出码 {rc}，请检查网络或源配置。")


if __name__ == "__main__":
	main()