#!/usr/bin/env -S uv run
# /// script
# dependencies = [
#     "httpx",
# ]
# ///
from concurrent.futures import ThreadPoolExecutor
import dataclasses
import itertools
import os
from pathlib import Path
import httpx
import argparse
import time
import urllib.parse


class AlldebridError(Exception):
    def __init__(self, code: str, message: str):
        self.code = code
        self.message = message
        super().__init__(f"{code}: {message}")


@dataclasses.dataclass
class Link:
    filename: str
    url: str
    size: int
    download_url: str | None = None

    @property
    def size_mb(self) -> float:
        return self.size / 1_048_576


def flatten_tree(data: dict | list) -> list[dict]:
    if isinstance(data, list):
        return itertools.chain.from_iterable(flatten_tree(item) for item in data)
    elif isinstance(data, dict):
        if "e" in data:  # It's a directory
            return flatten_tree(data["e"])
        elif "l" in data and "s" in data:  # It's a file
            return [data]
    return []


pool = ThreadPoolExecutor(max_workers=10)


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) -> Link | 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 Link(
                filename=data["data"]["filename"],
                url=url,
                size=data["data"]["filesize"],
                download_url=data["data"]["link"],
            )

    def _unrestrict_link(self, link: Link) -> Link | None:
        link.download_url = self.unrestrict_url(link.url).download_url
        return 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"]
            if not files:
                return
            files = [Link(url=it["l"], filename=it["n"], size=it["s"]) for it in flatten_tree(files)]
            large_files = [f for f in files if f.size_mb > 5]
            if not large_files:
                raise ValueError("torrent only contains small files")
            return large_files

    def wait_for_download_links(self, magnet_id, timeout=600) -> list[Link]:
        start = time.time()
        while time.time() - start < timeout:
            links = self.get_torrent_links(magnet_id)
            if links:
                return list(pool.map(self._unrestrict_link, links))
            time.sleep(10)
        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")
    parser.add_argument("--html", action="store_true", help="Print links as HTML")
    return parser.parse_args()


def print_links(links: list[Link]):
    if not links:
        return
    for link in links:
        print(link.download_url)


def print_links_as_html(links: list[Link]):
    if not links:
        return
    css = """
        body {
            margin: 0;
            padding: 1rem;
            font-family: consolas, menlo, monospace;
            font-size: 14px;
        }
        table {
            border-collapse: collapse;
            width: 100%;
        }
        th, td {
            border: 1px solid black;
            padding: 8px;
            text-align: left;
        }
        td.numeric {
            text-align: right;
        }
    """
    row_htmls = []
    for it in links:
        encoded_link = urllib.parse.quote(it.download_url)
        alfred_link = f"alfred://runtrigger/piracy/direct_link/?argument={encoded_link}"
        row_html = f"""
            <tr>
                <td><a href='{it.download_url}'>{it.filename}</a></td>
                <td class='numeric'>{it.size_mb:.1f}</td>
                <td><a href='{alfred_link}'>Alfred</a></td>
            </tr>
        """
        row_htmls.append(row_html)

    table_html = f"""
        <style>{css}</style>
        <script>window.app.setFloating(false)</script>
        <table>
            <thead><tr>
                <th>Filename</th>
                <th>Size MB</th>
                <th>Action</th>
            </tr></thead>
            <tbody>
            {"".join(row_htmls)}
            </tbody>
        </table>
    """
    print(table_html)


def main():
    args = parse_args()
    if not args.token:
        raise ValueError("No token provided")

    downloader = AlldebridDownloader(args.token)

    printer = print_links_as_html if args.html else print_links

    input_as_path = Path(args.input)
    if args.input.startswith("magnet:"):
        magnet_id = downloader.upload_magnet(args.input)
        if not magnet_id:
            return
        links = downloader.wait_for_download_links(magnet_id)
        printer(links)
    elif args.input.startswith("http:") or args.input.startswith("https:"):
        link = downloader.unrestrict_url(args.input)
        printer([link])
    elif input_as_path.is_file() and input_as_path.suffix == ".torrent":
        torrent_id = downloader.upload_torrent(input_as_path)
        if not torrent_id:
            return
        links = downloader.wait_for_download_links(torrent_id)
        printer(links)
    else:
        raise ValueError("Invalid input")


if __name__ == "__main__":
    main()