Medios-Macina/cmdlet/download_torrent.py

"""Download torrent/magnet links via AllDebrid in a dedicated cmdlet.

Features:
- Accepts magnet links and .torrent files/url
- Uses AllDebrid API for background downloads
- Progress tracking and worker management
- Self-registering class-based cmdlet
"""

from __future__ import annotations
import sys
import uuid
import threading
from pathlib import Path
from typing import Any, Dict, Optional, Sequence

from SYS.logger import log
from . import _shared as sh


class Download_Torrent(sh.Cmdlet):
    """Class-based download-torrent cmdlet with self-registration."""

    def __init__(self) -> None:
        super().__init__(
            name="download-torrent",
            summary="Download torrent/magnet links via AllDebrid",
            usage="download-torrent <magnet|.torrent> [options]",
            alias=["torrent",
                   "magnet"],
            arg=[
                sh.CmdletArg(
                    name="magnet",
                    type="string",
                    required=False,
                    description="Magnet link or .torrent file/URL",
                    variadic=True,
                ),
                sh.CmdletArg(
                    name="output",
                    type="string",
                    description="Output directory for downloaded files",
                ),
                sh.CmdletArg(
                    name="wait",
                    type="float",
                    description="Wait time (seconds) for magnet processing timeout",
                ),
                sh.CmdletArg(
                    name="background",
                    type="flag",
                    alias="bg",
                    description="Start download in background",
                ),
            ],
            detail=["Download torrents/magnets via AllDebrid API."],
            exec=self.run,
        )
        self.register()

    def run(self, result: Any, args: Sequence[str], config: Dict[str, Any]) -> int:
        parsed = sh.parse_cmdlet_args(args, self)
        magnet_args = parsed.get("magnet", [])
        output_dir = Path(parsed.get("output") or Path.home() / "Downloads")
        wait_timeout = int(float(parsed.get("wait", 600)))
        background_mode = parsed.get("background", False)
        api_key = None
        try:
            from Provider.alldebrid import _get_debrid_api_key  # type: ignore

            api_key = _get_debrid_api_key(config)
        except Exception:
            api_key = None
        if not api_key:
            log(
                "AllDebrid API key not configured (check config.conf [provider=alldebrid] api_key=...)",
                file=sys.stderr,
            )
            return 1
        for magnet_url in magnet_args:
            if background_mode:
                self._start_background_worker(
                    magnet_url,
                    output_dir,
                    config,
                    api_key,
                    wait_timeout
                )
                log(f"⧗ Torrent download queued in background: {magnet_url}")
            else:
                # Foreground mode: submit quickly, then continue processing in background
                # so we return control to the REPL immediately.
                worker_id = str(uuid.uuid4())
                magnet_id = self._submit_magnet(worker_id, magnet_url, api_key)
                if magnet_id <= 0:
                    continue
                self._start_background_magnet_worker(
                    worker_id,
                    magnet_id,
                    output_dir,
                    api_key,
                    wait_timeout
                )
                log(f"⧗ Torrent processing started (ID: {magnet_id})")
        return 0

    @staticmethod
    def _submit_magnet(worker_id: str, magnet_url: str, api_key: str) -> int:
        """Submit a magnet and return its AllDebrid magnet ID.

        This is intentionally fast so the caller can return to the REPL.
        """
        try:
            from API.alldebrid import AllDebridClient

            client = AllDebridClient(api_key)
            log(f"[Worker {worker_id}] Submitting magnet to AllDebrid...")
            magnet_info = client.magnet_add(magnet_url)
            magnet_id = int(magnet_info.get("id", 0))
            if magnet_id <= 0:
                log(f"[Worker {worker_id}] Magnet add failed", file=sys.stderr)
                return 0
            log(f"[Worker {worker_id}] ✓ Magnet added (ID: {magnet_id})")
            return magnet_id
        except Exception as e:
            log(f"[Worker {worker_id}] Magnet submit failed: {e}", file=sys.stderr)
            return 0

    def _start_background_magnet_worker(
        self,
        worker_id: str,
        magnet_id: int,
        output_dir: Path,
        api_key: str,
        wait_timeout: int
    ) -> None:
        thread = threading.Thread(
            target=self._download_magnet_worker,
            args=(worker_id,
                  magnet_id,
                  output_dir,
                  api_key,
                  wait_timeout),
            daemon=True,
            name=f"TorrentWorker_{worker_id}",
        )
        thread.start()

    @staticmethod
    def _download_magnet_worker(
        worker_id: str,
        magnet_id: int,
        output_dir: Path,
        api_key: str,
        wait_timeout: int = 600,
    ) -> None:
        """Poll AllDebrid magnet status until ready, then download the files."""
        try:
            from API.alldebrid import AllDebridClient

            client = AllDebridClient(api_key)

            # Poll for ready status (simplified)
            import time

            elapsed = 0
            while elapsed < wait_timeout:
                status = client.magnet_status(magnet_id)
                if status.get("ready"):
                    break
                time.sleep(5)
                elapsed += 5
            if elapsed >= wait_timeout:
                log(f"[Worker {worker_id}] Timeout waiting for magnet", file=sys.stderr)
                return

            files_result = client.magnet_links([magnet_id])
            magnet_files = files_result.get(str(magnet_id),
                                            {})
            files_array = magnet_files.get("files", [])
            if not files_array:
                log(f"[Worker {worker_id}] No files found", file=sys.stderr)
                return
            for file_info in files_array:
                file_url = file_info.get("link")
                file_name = file_info.get("name")
                if file_url and file_name:
                    Download_Torrent._download_file(file_url, output_dir / file_name)
                    log(f"[Worker {worker_id}] ✓ Downloaded {file_name}")
        except Exception as e:
            log(f"[Worker {worker_id}] Torrent download failed: {e}", file=sys.stderr)

    @staticmethod
    def _download_torrent_worker(
        worker_id: str,
        magnet_url: str,
        output_dir: Path,
        config: Dict[str,
                     Any],
        api_key: str,
        wait_timeout: int = 600,
        worker_manager: Optional[Any] = None,
    ) -> None:
        try:
            from API.alldebrid import AllDebridClient

            client = AllDebridClient(api_key)
            log(f"[Worker {worker_id}] Submitting magnet to AllDebrid...")
            magnet_info = client.magnet_add(magnet_url)
            magnet_id = int(magnet_info.get("id", 0))
            if magnet_id <= 0:
                log(f"[Worker {worker_id}] Magnet add failed", file=sys.stderr)
                return
            log(f"[Worker {worker_id}] ✓ Magnet added (ID: {magnet_id})")
            # Poll for ready status (simplified)
            import time

            elapsed = 0
            while elapsed < wait_timeout:
                status = client.magnet_status(magnet_id)
                if status.get("ready"):
                    break
                time.sleep(5)
                elapsed += 5
            if elapsed >= wait_timeout:
                log(f"[Worker {worker_id}] Timeout waiting for magnet", file=sys.stderr)
                return
            files_result = client.magnet_links([magnet_id])
            magnet_files = files_result.get(str(magnet_id),
                                            {})
            files_array = magnet_files.get("files", [])
            if not files_array:
                log(f"[Worker {worker_id}] No files found", file=sys.stderr)
                return
            for file_info in files_array:
                file_url = file_info.get("link")
                file_name = file_info.get("name")
                if file_url:
                    Download_Torrent._download_file(file_url, output_dir / file_name)
                    log(f"[Worker {worker_id}] ✓ Downloaded {file_name}")
        except Exception as e:
            log(f"[Worker {worker_id}] Torrent download failed: {e}", file=sys.stderr)

    @staticmethod
    def _download_file(url: str, dest: Path) -> None:
        try:
            import requests

            resp = requests.get(url, stream=True)
            with open(dest, "wb") as f:
                for chunk in resp.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
        except Exception as e:
            log(f"File download failed: {e}", file=sys.stderr)

    def _start_background_worker(
        self,
        magnet_url,
        output_dir,
        config,
        api_key,
        wait_timeout
    ):
        worker_id = f"torrent_{uuid.uuid4().hex[:6]}"
        thread = threading.Thread(
            target=self._download_torrent_worker,
            args=(worker_id,
                  magnet_url,
                  output_dir,
                  config,
                  api_key,
                  wait_timeout),
            daemon=True,
            name=f"TorrentWorker_{worker_id}",
        )
        thread.start()


CMDLET = Download_Torrent()