#!/usr/bin/env python3
import requests
# 2025-11-22

class Aria2Client:
	"""
	简洁易用的 Aria2 RPC 客户端封装
	支持添加、查询、控制下载任务
	"""

	def __init__(self, host="localhost", port=6800, secret=None):
		"""
		初始化 Aria2 客户端
		
		:param host: aria2 服务地址 (默认 localhost)
		:param port: RPC 端口 (默认 6800)
		:param secret: RPC 密钥 (在配置文件中设置的 rpc-secret)
		"""
		self.rpc_url = f"http://{host}:{port}/jsonrpc"
		self.secret = secret
		self.session = requests.Session()

	def _call(self, method, *params):
		"""内部调用 RPC 方法"""
		payload = {
			"jsonrpc": "2.0",
			"id": "aria2-python-client",
			"method": f"aria2.{method}",
			"params": [f"token:{self.secret}"] + list(params) if self.secret else list(params)
		}

		response = self.session.post(self.rpc_url, json=payload)
		response.raise_for_status()

		result = response.json()
		if "error" in result:
			raise Exception(f"RPC Error: {result['error']['message']}")
		return result.get("result")

	# ========== 核心操作 ==========

	def add_url(self, url, save_path=None, filename=None, options=None):
		"""
		添加下载任务
		
		:param url: 下载链接 (HTTP/FTP/BT/磁力链)
		:param save_path: 保存目录 (可选)
		:param filename: 自定义文件名 (可选)
		:param options: 额外参数字典 (可选)
		:return: 任务ID (gid)
		
		示例:
			add_url("http://example.com/file.zip")
			add_url("magnet:?xt=...", save_path="/download/movie")
		"""
		opts = options or {}
		if save_path:
			opts["dir"] = save_path
		if filename:
			opts["out"] = filename

		return self._call("addUri", [url], opts)

	# ========== 任务查询 ==========

	def get_all_tasks(self):
		"""获取所有任务（活动+等待+停止）"""
		active = self._call("tellActive") or []
		waiting = self._call("tellWaiting", 0, 999999) or []
		stopped = self._call("tellStopped", 0, 999999) or []
		return active + waiting + stopped

	def get_downloading_tasks(self):
		"""获取正在下载的任务"""
		return self._call("tellActive") or []

	def get_waiting_tasks(self):
		"""获取等待中的任务"""
		return self._call("tellWaiting", 0, 999999) or []

	def get_finished_tasks(self):
		"""获取已完成的任务"""
		return self._call("tellStopped", 0, 999999) or []

	def get_task_status(self, gid):
		"""
		获取指定任务的详细信息
		
		:param gid: 任务ID
		:return: 任务详情字典
		"""
		return self._call("tellStatus", gid)

	# ========== 任务控制 ==========

	def pause(self, gid):
		"""
		暂停单个任务
		
		:param gid: 任务ID
		:return: "OK"
		"""
		return self._call("pause", gid)

	def resume(self, gid):
		"""
		恢复单个任务
		
		:param gid: 任务ID
		:return: "OK"
		"""
		return self._call("unpause", gid)

	def remove(self, gid):
		"""
		删除任务（包括已完成的）
		
		:param gid: 任务ID
		:return: "OK"
		"""
		return self._call("remove", gid)

	def pause_all(self):
		"""暂停所有任务（活动中+等待中）"""
		active = self.get_downloading_tasks()
		waiting = self.get_waiting_tasks()
		for task in active + waiting:
			if isinstance(task, dict) and task.get("gid"):
				try:
					print(task["gid"])
					self._call("pause", task["gid"])
				except Exception:
					pass

	def remove_all(self):
		"""删除所有任务（活动中+等待中）"""
		active = self.get_downloading_tasks()
		waiting = self.get_waiting_tasks()
		for task in active + waiting:
			if isinstance(task, dict) and task.get("gid"):
				try:
					print(task["gid"])
					self._call("remove", task["gid"])
				except Exception:
					pass

	def clear_finished(self):
		"""清除已完成的下载记录"""
		stopped = self.get_finished_tasks()
		for task in stopped:
			if isinstance(task, dict) and task.get("gid"):
				try:
					print(task["gid"])
					self._call("removeDownloadResult", task["gid"])
				except Exception:
					pass

	# ========== 实用工具 ==========

	def get_version(self):
		"""获取 aria2 版本信息"""
		return self._call("getVersion")

	def get_global_stats(self):
		"""获取全局统计（速度、任务数等）"""
		return self._call("getGlobalStat")

	def get_download_speed(self):
		"""获取当前下载速度 (bytes/s)"""
		stats = self.get_global_stats()
		return int(stats.get("downloadSpeed", 0))

	def get_upload_speed(self):
		"""获取当前上传速度 (bytes/s)"""
		stats = self.get_global_stats()
		return int(stats.get("uploadSpeed", 0))

	def set_speed_limit(self, download_limit=None, upload_limit=None):
		"""
		设置速度限制
		
		:param download_limit: 下载限速 (bytes/s, None=不限速)
		:param upload_limit: 上传限速 (bytes/s, None=不限速)
		"""
		if download_limit is not None:
			self._call("changeGlobalOption", {"max-overall-download-limit": str(download_limit)})
		if upload_limit is not None:
			self._call("changeGlobalOption", {"max-overall-upload-limit": str(upload_limit)})

	def is_feature_installed(self, feature):
		"""
		检查 aria2 是否支持特定功能
		
		:param feature: 功能名称 (如 "BitTorrent", "AsyncDNS")
		:return: True/False
		"""
		features = self.get_version().get("enabledFeatures", [])
		return feature in features


# ========== 使用示例 ==========

if __name__ == "__main__":
	aria2 = Aria2Client(host="localhost", port=6800, secret="YOUR_SECRET_KEY")

	# # 1. 添加下载任务
	# gid = aria2.add_url(
	#     "http://example.com/file.zip",
	#     save_path="/downloads",
	#     filename="myfile.zip"
	# )
	# print(f"已添加任务: {gid}")

	# # 2. 获取所有任务
	# tasks = aria2.get_all_tasks()
	# for task in tasks:
	#     print(f"任务ID: {task['gid']}, 文件名: {task.get('files', [{}])[0].get('path', 'N/A')}, 进度: {task.get('completedLength', 0)}/{task.get('totalLength', 0)}")

	# # 3. 获取当前速度
	# print(f"下载速度: {aria2.get_download_speed() / 1024:.2f} KB/s")

	# # 4. 暂停任务
	# aria2.pause(gid)
	
	# # 5. 暂停所有任务
	# aria2.pause_all()
	
	# # 6. 删除所有任务
	# aria2.remove_all()
	
	# 7. 清除已完成记录
	aria2.clear_finished()