#!/usr/bin/env python3
import json
import os
import sys
from urllib.parse import urlparse
import requests


def get_base_url(url):
    parsed = urlparse(url)
    return f"{parsed.scheme}://{parsed.netloc}"


def collect_at2400_urls(data, base_url):
    urls = []

    def walk(obj):
        if isinstance(obj, dict):
            images = obj.get("images")
            if isinstance(images, dict):
                at2400 = images.get("at2400")
                if isinstance(at2400, str):
                    if at2400.startswith("http"):
                        urls.append(at2400)
                    else:
                        urls.append(base_url + at2400)
            for value in obj.values():
                walk(value)
        elif isinstance(obj, list):
            for item in obj:
                walk(item)

    walk(data)
    return urls


def download(url, path, session):
    r = session.get(url, timeout=60)
    r.raise_for_status()
    with open(path, "wb") as f:
        f.write(r.content)


def is_invalid_page_response(r, data):
    """
    Detect the site's 'no more pages' signal. Publitas can express this
    either via a non-200 status, or via a 200 response whose JSON body
    carries an error/message field. We check both so we stop cleanly
    regardless of which form it takes.
    """
    if r.status_code != 200:
        return True
    if isinstance(data, dict):
        for key in ("error", "message", "status"):
            val = data.get(key)
            if isinstance(val, str) and "invalid" in val.lower():
                return True
    return False


def fetch_spread_pages(book_base_url, session):
    """
    Given a viewer URL like:
        https://view.publitas.com/malmberg/605393-01_nvna_6vg_lob_bladerboek
    repeatedly fetch:
        <book_base_url>/spreads.json?page=1
        <book_base_url>/spreads.json?page=2
        ...
    until the site signals there is no such page, yielding parsed JSON
    payloads (and the base_url to resolve relative image paths against).
    """
    book_base_url = book_base_url.rstrip("/")
    page = 1
    while True:
        spreads_url = f"{book_base_url}/spreads.json?page={page}"
        try:
            r = session.get(spreads_url, timeout=60)
        except Exception as e:
            print(f"    Request failed for page {page}: {e}")
            break

        try:
            data = r.json()
        except ValueError:
            data = None

        if is_invalid_page_response(r, data):
            print(f"  Stopping at page {page} (no more pages).")
            break

        print(f"  Reading {spreads_url}")
        yield data

        page += 1


def derive_book_name(url):
    """
    Derive a book title from the URL's path, using the last path segment.
    E.g. https://view.publitas.com/malmberg/608941-01_nova_nask_12bk_lwb_a_bladerboek
      -> 608941-01_nova_nask_12bk_lwb_a_bladerboek
    Works for a plain viewer URL as well as a full spreads.json URL (the
    "spreads.json" segment itself is stripped off first).
    """
    path = urlparse(url).path.rstrip("/")
    segments = [s for s in path.split("/") if s]
    if segments and segments[-1].lower() == "spreads.json":
        segments = segments[:-1]
    if not segments:
        raise ValueError(f"Could not derive a book name from URL: {url}")
    return segments[-1]


def read_jobs(filename):
    """
    Parses a jobs file that is simply a list of URLs, one per line
    (blank lines ignored):

        https://view.publitas.com/malmberg/608941-01_nova_nask_12bk_lwb_a_bladerboek
        https://view.publitas.com/malmberg/605393-01_nvna_6vg_lob_bladerboek

    Each URL becomes its own job/book. The book's folder name is derived
    from the last path segment of the URL (see derive_book_name).

    Each URL may be either:
      - a plain viewer URL (e.g. .../605393-01_nvna_6vg_lob_bladerboek),
        in which case pages are auto-discovered by incrementing
        ?page=N until the site reports no more pages, or
      - a full spreads.json URL (e.g. .../spreads.json?page=1), which is
        fetched as-is (no auto-pagination), for backwards compatibility.

    If you want to group multiple URLs under the same book, list them on
    consecutive lines that derive to the same name -- they'll be merged
    automatically.
    """
    jobs = []  # list of (book_name, [urls])
    index_by_name = {}

    with open(filename, "r", encoding="utf-8") as f:
        for raw_line in f:
            line = raw_line.strip()
            if not line:
                continue
            if not line.startswith(("http://", "https://")):
                print(f"Skipping non-URL line: {line}")
                continue

            name = derive_book_name(line)
            if name in index_by_name:
                jobs[index_by_name[name]][1].append(line)
            else:
                index_by_name[name] = len(jobs)
                jobs.append((name, [line]))

    return jobs


def main():
    if len(sys.argv) < 2:
        print("Usage:")
        print("  python3 extract_publitas.py urls.txt [output_dir]")
        sys.exit(1)

    urls_file = sys.argv[1]
    output_dir = sys.argv[2] if len(sys.argv) > 2 else "out"
    os.makedirs(output_dir, exist_ok=True)

    session = requests.Session()
    jobs = read_jobs(urls_file)

    total_pages = 0
    for book_name, book_urls in jobs:
        print(f"\n=== {book_name} ===")
        book_dir = os.path.join(output_dir, book_name)
        os.makedirs(book_dir, exist_ok=True)

        page_num = 0
        for url in book_urls:
            base_url = get_base_url(url)

            if "spreads.json" in url:
                # Backwards-compatible path: explicit single spreads.json URL
                print(f"Reading {url}")
                try:
                    r = session.get(url, timeout=60)
                    r.raise_for_status()
                    data_iter = [r.json()]
                except Exception as e:
                    print(f"Failed {url}: {e}")
                    continue
            else:
                # New path: auto-paginate a viewer base URL
                print(f"Auto-paginating {url}")
                data_iter = fetch_spread_pages(url, session)

            for data in data_iter:
                image_urls = collect_at2400_urls(data, base_url)
                print(f"    Found {len(image_urls)} images on this page")
                for img_url in image_urls:
                    filename = f"{page_num:04d}.jpg"
                    path = os.path.join(book_dir, filename)
                    try:
                        download(img_url, path, session)
                        print(f"      {filename}")
                        page_num += 1
                    except Exception as e:
                        print(f"      Failed {img_url}: {e}")

        print(f"  -> {page_num} pages saved to {book_dir}")
        total_pages += page_num

    print(f"\nDone. Downloaded {total_pages} pages across {len(jobs)} book(s).")


if __name__ == "__main__":
    main()