#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p "python3.withPackages (ps: [ ps.natsort ps.requests ])"
# vim: set filetype=python :
"""
usage: sane-bt-search [options] <query_string>

searches Jackett for torrent files matching the title.
returns select results and magnet links.
"""

# about Jackett
# - source: <https://github.com/Jackett/Jackett>
# - can be queried via APIs:
#   - Torznab: <https://torznab.github.io/spec-1.3-draft/index.html>
#   - TorrentPotato: <https://github.com/RuudBurger/CouchPotatoServer/wiki/Couchpotato-torrent-provider>
#   - its own JSON-based API
#
# about torrent jargon:
# - <https://en.wikipedia.org/wiki/Pirated_movie_release_types>

from dataclasses import dataclass
from datetime import datetime
import argparse
import logging
import json
import natsort
import requests
import sys
import time

SERVICE = "https://jackett.uninsane.org"
ENDPOINTS = dict(
    results="api/v2.0/indexers/all/results"
    # results_torznab="api/v2.0/indexers/all/results/torznab"
)

DEFAULT_RESULT_COUNT = 12

epoch = datetime(1970, 1, 1)

logger = logging.getLogger(__name__)

class BadCliArgs(Exception):
    def __init__(self, msg: str = None):
        helpstr = __doc__
        if msg:
            super().__init__(f"{msg}\n\n{helpstr}")
        else:
            super().__init__(helpstr)


def try_parse_time(t: str):
    try:
        return datetime.fromisoformat(t)
    except ValueError: pass

    if len(t) > len('YYYY-MM-DD'):
        # sometimes these timestamps are encoded with e.g. too many digits in the milliseconds field.
        # so just keep chomping until we get something that parses as a timestamp
        return try_parse_time(t[:-1])

def parse_time(t: str) -> datetime:
    return try_parse_time(t).astimezone() or epoch


# preference, best to worst
TRACKER_RANKS = [
    'bakabt',
    'subsplease',
    'nyaa.si',
    'miobt',
    'yts',
    'internet archive',
    'the pirate bay',
    # haven't sorted these
    '1337x',
    'kickasstorrents.to',
]
def tracker_rank(tracker: str):
    tracker_ = tracker.lower()
    if tracker_ in TRACKER_RANKS:
        return TRACKER_RANKS.index(tracker_)
    logger.warn(f"unknown tracker: {tracker_!r}")
    return len(TRACKER_RANKS)

DROP_CATS = { "dvd", "hd", "misc", "other", "sd", "uhd" }
BOOK_CATS = { "audio", "books", "ebook" }
MANGA_CATS = { "books", "comics", "ebook" }
VIDEO_CATS = { "anime", "movies", "tv" }
KNOWN_CATS = frozenset(
    list(BOOK_CATS) +
    list(MANGA_CATS) +
    list(VIDEO_CATS) +
    ["xxx"]
)
def clean_cat(c: str) -> str | None:
    if c in DROP_CATS: return None
    return c

def is_cat(cats: list[str], wanted_cats: list[str], default: bool = False) -> bool:
    """
    return True if any of the `cats` is in `wanted_cats`.
    in the event there no category is recognized, assume `default`
    """
    if not any(c in KNOWN_CATS for c in cats):
        return default
    else:
        return any(c in wanted_cats for c in cats)

class Filter:
    def __init__(self, book: bool=False, h265: bool=False, manga: bool=False, video: bool=False):
        self.book = book
        self.h265 = h265
        self.manga = manga
        self.video = video

    def filter(self, torrents: list['Torrent']) -> list['Torrent']:
        return [t for t in torrents if self.is_match(t)]

    def is_match(self, t: 'Torrent', default: bool = False) -> bool:
        valid = True
        valid = valid and (not self.book or self.is_book(t, default))
        valid = valid and (not self.h265 or self.is_h265(t))
        valid = valid and (not self.manga or self.is_manga(t, default))
        valid = valid and (not self.video or self.is_video(t, default))
        return valid

    @staticmethod
    def is_h265(t: 'Torrent') -> bool:
        meta = t.title.lower()
        return "h265" in meta \
            or "x265" in meta \
            or "HEVC" in meta

    @staticmethod
    def is_book(t: 'Torrent', default: bool = False) -> bool:
        return is_cat(t.categories, BOOK_CATS, default)

    @staticmethod
    def is_manga(t: 'Torrent', default: bool = False) -> bool:
        return is_cat(t.categories, MANGA_CATS, default)

    @staticmethod
    def is_video(t: 'Torrent', default: bool = False) -> bool:
        return is_cat(t.categories, VIDEO_CATS, default)


@dataclass(eq=True, order=True, unsafe_hash=True)
class Torrent:
    seeders: int
    pub_date: datetime
    size: int
    tracker: str
    title: str
    magnet: str | None
    http_dl_uri: str | None  # probably a .torrent file but it COULD be a referral to a magnet:// URI
    tracker_uri: str | None
    categories: frozenset[str]  # human-friendly list of categories, lowercase. e.g. ["Books", "Anime"]

    def __str__(self) -> str:
        cats = "/".join(self.categories) if self.categories else "?"
        rows = []
        rows.append(f"{self.seeders}[S]\t{cats}\t{self.tracker}\t{self.pub_date}\t{self.mib}M\t{self.title}")
        if self.tracker_uri:
            rows.append(f"\t{self.tracker_uri}")
        rows.append(f"\t{self.dl_uri}")
        return "\n".join(rows)

    @property
    def dl_uri(self) -> str | None:
        return self.magnet or self.http_dl_uri

    @property
    def mib(self) -> int:
        return int(round(self.size / 1024 / 1024))

    @staticmethod
    def from_dict(d: dict) -> 'Torrent':
        logger.debug(f"Torrent.from_dict: fields: { ' '.join(d.keys()) }")
        for k, v in d.items():
            if k not in ("CategoryDesc", "Seeders", "PublishDate", "Size", "Tracker", "Title", "MagnetUri", "Guid", "Link", "Details") and \
                    v != None and v != "" and v != [] and v != {}:
                logger.debug(f"  {k} = {v}")

        seeders = d.get("Seeders")
        pub_date = d.get("PublishDate")
        size = d.get("Size")
        tracker = d.get("Tracker")
        title = d.get("Title")
        magnet = d.get("MagnetUri") or d.get("Guid")
        http_dl_uri = d.get("Link")
        tracker_uri = d.get("Details")
        categories = d.get("CategoryDesc", "").replace("/", ",").split(",")
        categories = (c.strip().lower() for c in categories)
        categories = frozenset(clean_cat(c) for c in categories if clean_cat(c))

        if magnet and not magnet.startswith("magnet:"):
            logger.info(f"invalid magnet: {magnet}")
            magnet = None

        if http_dl_uri:
            # quirk: jackett responds with port number even if it's a default port number
            # it still works, just adds noise to the output
            http_dl_uri = http_dl_uri.replace('https://jackett.uninsane.org:443/', 'https://jackett.uninsane.org/')

        if tracker_uri:
            # quirk: bakabt's post-login redirection fails if the slashes aren't just right
            tracker_uri = tracker_uri.replace('https://bakabt.me//', 'https://bakabt.me/')

        if seeders is not None and pub_date is not None and title is not None and (magnet is not None or http_dl_uri is not None):
            pub_date = parse_time(pub_date)
            return Torrent(seeders, pub_date, size, tracker, title, magnet, http_dl_uri, tracker_uri, categories=categories)

    def to_dict(self) -> dict:
        # N.B.: not all fields: needs to be kept in sync with consumers like mx-sanebot
        return dict(
            seeders=self.seeders,
            pub_date=self.pub_date.strftime("%Y-%m-%d"),
            size=self.size,
            tracker=self.tracker,
            title=self.title,
            magnet=self.magnet,
        )

class Client:
    def __init__(self):
        self.apikey = open("/run/secrets/jackett_apikey").read().strip()

    def api_call(self, method: str, params: dict) -> dict:
        endpoint = ENDPOINTS[method]
        url = f"{SERVICE}/{endpoint}"
        params = params.copy()
        params.update(apikey=self.apikey, _=str(int(time.time())))
        resp = requests.get(url, params=params)
        text = str(resp.text)
        try:
            return json.loads(text)
        except Exception as e:
            logger.error(f"failed to parse json API response from {url}: {text}")
            raise e

    def query(self, q: str) -> list[Torrent]:
        torrents = set()
        api_res = self.api_call("results", dict(Query=q))
        for r in api_res["Results"]:
            t = Torrent.from_dict(r)
            if t is not None:
                torrents.add(t)

        return sorted(torrents, reverse=True)

def sort_results(torrents: list[Torrent], by: str) -> list[Torrent]:
    if by == 'seeders':
        return sorted(torrents, key=lambda t: (t.seeders, t), reverse=True)
    elif by == 'tracker':
        return sorted(torrents, key=lambda t: (-tracker_rank(t.tracker), t), reverse=True)
    assert False, f"unknown sort method: {by}"

def format_results(all_results: list[Torrent], filtered_results: list[Torrent], json: bool):
    if json:
        dumpable = [t.to_dict() for t in filtered_results]
        print(json.dumps(dumpable))
    else:
        num_total = len(all_results)
        num_filtered = len(filtered_results)
        print(f"found {num_total} result(s) filtered to {num_filtered}")
        for r in filtered_results:
            print(r)

def main(args: list[str]):
    logging.basicConfig()
    logging.getLogger().setLevel(logging.WARNING)

    parser = argparse.ArgumentParser(description='search torrent trackers')
    parser.add_argument('--full', action='store_true', help='show all results')
    parser.add_argument('--top', help=f'how many results to show (default: {DEFAULT_RESULT_COUNT})')
    parser.add_argument('--sort-by', default='seeders', help='how to rank matches (seeders, tracker)')
    parser.add_argument('--json', action='store_true', help='output results in json')
    parser.add_argument('--verbose', action='store_true')
    parser.add_argument('--book', action='store_true', help='show only book (ebook or audiobook) results')
    parser.add_argument('--h265', action='store_true', help='show only H.265/HEVC results (might cause false negatives)')
    parser.add_argument('--manga', action='store_true', help='show only manga results')
    parser.add_argument('--video', action='store_true', help='show only video (tv or film)')
    parser.add_argument("query", help="text to search for in torrent titles/descriptions")

    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    client = Client()
    all_results = client.query(args.query)

    filter = Filter(book=args.book, h265=args.h265, manga=args.manga, video=args.video)
    filtered_results = filter.filter(all_results)
    ordered_results = sort_results(filtered_results, args.sort_by)

    if not args.full:
        ordered_results = ordered_results[:int(args.top or str(DEFAULT_RESULT_COUNT))]

    format_results(all_results, ordered_results, args.json)

if __name__ == "__main__":
    main(sys.argv[1:])
