#! /usr/bin/env python

# coomer_dl - download Coomer creators, posts and single files
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import requests
import re
import os
import argparse
import time
from urllib.parse import urlparse, parse_qs

DEFAULT_OUTPUT = "/data/coomer"  # Base download folder, override with -o/--output

headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/85.0.4183.121 Safari/537.36'
}
# DDoS-Guard answers API calls with 403 unless the client asks for text/css
api_headers = {**headers, 'Accept': 'text/css'}

PAGE_SIZE = 50  # Posts returned per /posts call


def api_get(base, path, **params):
    """GET <base>/api/v1<path> and return the decoded JSON, or None on failure"""
    try:
        r = requests.get(f"{base}/api/v1{path}", params=params, headers=api_headers, timeout=60)
        r.raise_for_status()
        return r.json()
    except (requests.exceptions.RequestException, ValueError) as e:
        print(f"Failed to fetch {base}/api/v1{path}: {e}")
        return None


def post_files(post):
    """Yield (path, name) for the main file and attachments of a post, without duplicates"""
    seen = set()
    for f in [post.get('file') or {}] + (post.get('attachments') or []):
        if f.get('path') and f['path'] not in seen:
            seen.add(f['path'])
            yield f['path'], f.get('name') or f['path'].split('/')[-1]


def download_file(file_url, file_name, save_dir, retries=5):
    """Function to download the file with retries and resume functionality"""
    local_filename = os.path.join(save_dir, file_name.replace('/', '_'))

    attempt = 0
    while attempt < retries:
        # Recompute each attempt so a retry resumes from what was written
        req_headers = dict(headers)
        if os.path.exists(local_filename):
            req_headers["Range"] = f"bytes={os.path.getsize(local_filename)}-"
        try:
            with requests.get(file_url, headers=req_headers, stream=True, timeout=(15, 60)) as response:  # (connect, read)
                if response.status_code == 416:  # Range past end: already complete
                    print(f"Already downloaded: {local_filename}")
                    return local_filename
                response.raise_for_status()
                # Check if it's a resumable download
                if response.status_code == 206:
                    print(f"Resuming download for: {local_filename}")
                    mode = 'ab'
                else:  # Full body: start over instead of appending to a partial file
                    print(f"Downloading: {local_filename}")
                    mode = 'wb'

                with open(local_filename, mode) as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        if chunk:  # Filter out keep-alive chunks
                            f.write(chunk)
            print(f"Download complete: {local_filename}")
            return local_filename
        except (requests.exceptions.ConnectionError, requests.exceptions.ChunkedEncodingError, requests.exceptions.Timeout) as e:
            attempt += 1
            print(f"Error during download: {e}. Retrying {attempt}/{retries}...")
            time.sleep(5)  # Wait before retrying
        except Exception as e:
            print(f"Failed to download {file_url}: {e}")
            break

    print(f"Failed to download {file_url} after {retries} attempts.")
    return None


def download_posts(base, posts, save_dir):
    """Download every file of the given posts into save_dir"""
    os.makedirs(save_dir, exist_ok=True)
    for post in posts:
        for path, name in post_files(post):
            # Prefix the post id: keeps a post's files together and names unique
            download_file(f"{base}/data{path}", f"{post['id']}_{name}", save_dir)


def creator_name(base, service, user):
    """Display name of a creator, falling back to the id in the URL"""
    profile = api_get(base, f"/{service}/user/{user}/profile")
    return (profile or {}).get('name') or user


def process_creator(base, service, user, base_save_directory):
    """Download every post of a creator, following ?o= pagination"""
    posts = []
    while True:
        page = api_get(base, f"/{service}/user/{user}/posts", o=len(posts))
        if not page:  # Empty list past the last page, None on error
            break
        posts += page
        print(f"Found {len(posts)} posts so far")
        if len(page) < PAGE_SIZE:
            break

    if not posts:
        print(f"No posts found for {service}/{user}")
        return
    download_posts(base, posts, os.path.join(base_save_directory, creator_name(base, service, user)))


def process_post(base, service, user, post_id, base_save_directory):
    """Download every file of a single post"""
    data = api_get(base, f"/{service}/user/{user}/post/{post_id}")
    if not data:
        return
    post = data.get('post', data)  # Newer API wraps the post, older returns it bare
    download_posts(base, [post], os.path.join(base_save_directory, creator_name(base, service, user)))


def main():
    parser = argparse.ArgumentParser(
        description="Download every file of a Coomer creator, of a single post, or a single\n"
                    "file. Partially downloaded files are resumed and finished ones skipped\n"
                    "when run again against the same folder.",
        epilog="Creator and post files go to OUTPUT/<creator>/<post id>_<file name>.\n"
               "Single files go to OUTPUT/<file name>.\n\n"
               "Examples:\n"
               "  %(prog)s https://coomer.st/onlyfans/user/NAME\n"
               "  %(prog)s https://coomer.st/onlyfans/user/NAME/post/123456\n"
               "  %(prog)s -o ~/Downloads 'https://coomer.st/data/ab/cd/HASH.mp4?f=name.mp4'",
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('url', help="creator URL (https://<domain>/<service>/user/<id>), post URL "
                                    "(.../post/<id>) or file URL (.../data/...)")
    parser.add_argument('-o', '--output', default=DEFAULT_OUTPUT, metavar='DIR',
                        help=f"base download folder (default: {DEFAULT_OUTPUT})")
    args = parser.parse_args()

    base_save_directory = os.path.expanduser(args.output)
    u = urlparse(args.url)
    # File nodes (n1.coomer.st, ...) serve /data too, the API lives on the main domain
    base = f"{u.scheme}://{re.sub(r'^n\d+\.', '', u.netloc)}"

    if u.path.startswith('/data/'):
        os.makedirs(base_save_directory, exist_ok=True)
        name = parse_qs(u.query).get('f', [u.path.split('/')[-1]])[0]
        download_file(args.url, name, base_save_directory)
    elif m := re.match(r'/([^/]+)/user/([^/]+)/post/([^/]+)', u.path):
        print(f"Processing post {args.url}")
        process_post(base, *m.groups(), base_save_directory)
    elif m := re.match(r'/([^/]+)/user/([^/]+)/?$', u.path):
        print(f"Processing creator {args.url}")
        process_creator(base, *m.groups(), base_save_directory)
    else:
        parser.error(f"unrecognised URL: {args.url}")


if __name__ == "__main__":
    main()
