# /// script
# requirements = ["httpx", "pydantic"]
# ///
import dataclasses
import os
from pathlib import Path
import httpx
import argparse
import time
import pydantic


class MagnetStatus(pydantic.BaseModel):
    id: int
    filename: str
    size: int
    status: str
    downloaded: int


@dataclasses.dataclass()
class AlldebridError(Exception):
    code: str
    message: str


class AlldebridDownloader:
    def __init__(self, api_token: str):
        self.client = httpx.Client(
            timeout=10,
            headers={
                "Authorization": f"Bearer {api_token}",
                "User-Agent": "alldebrid downloader for abdusco",
            },
        )

    def _check_error(self, res: httpx.Response):
        res.raise_for_status()
        data = res.json()
        if data["status"] == "error":
            error = data["error"]
            raise AlldebridError(code=error["code"], message=error["message"])

    def unrestrict_url(self, url: str) -> str | None:
        res = self.client.get(
            "https://api.alldebrid.com/v4/link/unlock",
            params={
                "link": url,
            },
        )
        self._check_error(res)
        data = res.json()
        if data["status"] == "success":
            return data["data"].get("link")

    def upload_torrent(self, torrent_path: Path) -> int:
        res = self.client.post(
            "https://api.alldebrid.com/v4/magnet/upload/file",
            files={
                "files[]": (torrent_path.name, torrent_path.read_bytes(), "application/x-bittorrent"),
            },
        )
        self._check_error(res)
        data = res.json()
        if data["status"] == "success":
            file = data["data"]["files"][0]
            return file["id"]

    def upload_magnet(self, magnet_uri: str) -> int:
        res = self.client.post(
            "https://api.alldebrid.com/v4/magnet/upload",
            data={
                "magnets[]": magnet_uri,
            },
        )
        self._check_error(res)
        data = res.json()
        if data["status"] == "success":
            magnet = data["data"]["magnets"][0]
            return magnet["id"]

    def get_torrent_links(self, torrent_id: int) -> list[str] | None:
        res = self.client.post(
            f"https://api.alldebrid.com/v4/magnet/files",
            data={"id[]": torrent_id},
        )
        self._check_error(res)

        data = res.json()
        if data["status"] == "success":
            magnets = data["data"]["magnets"]
            if not magnets:
                return
            files = magnets[0]["files"][0]["e"]
            if not files:
                raise ValueError("torrent is empty")
            large_files = [m for m in files if m["s"] > 1_000_000]
            if not large_files:
                raise ValueError("torrent only contains small files")
            links = [it["l"] for it in large_files]
            return links

    def wait_for_download_link(self, magnet_id, timeout=600):
        start = time.time()
        while time.time() - start < timeout:
            links = self.get_torrent_links(magnet_id)
            if links:
                return self.unrestrict_url(links[0])
            time.sleep(5)
        return None


def parse_args():
    parser = argparse.ArgumentParser(description="Alldebrid Downloader")
    parser.add_argument("input", help="URL, magnet link, or torrent file")
    parser.add_argument("--token", default=os.getenv("ALLDEBRID_TOKEN"), help="Alldebrid API token")
    return parser.parse_args()


def main():
    args = parse_args()

    downloader = AlldebridDownloader(args.token)

    input_as_path = Path(args.input)
    if input_as_path.is_file() and input_as_path.suffix == ".torrent":
        torrent_id = downloader.upload_torrent(input_as_path)
        if not torrent_id:
            return
        result = downloader.wait_for_download_link(torrent_id)
        if result:
            print(result)
    elif args.input.startswith("magnet:"):
        magnet_data = downloader.upload_magnet(args.input)
        if not magnet_data:
            return
        result = downloader.wait_for_download_link(magnet_data["id"])
        if result:
            print(result)
    elif args.input.startswith("http:") or args.input.startswith("https:"):
        result = downloader.unrestrict_url(args.input)
        if not result:
            return
        print(result)
    else:
        raise ValueError("Invalid input")


if __name__ == "__main__":
    main()