"""
ia_download.py

'ia' subcommand for downloading files from archive.org.
"""

# Copyright (C) 2012-2026 Internet Archive
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import annotations

import argparse
import re
import sys
from typing import TextIO

from internetarchive.cli.cli_utils import (
    QueryStringAction,
    validate_dir_path,
)
from internetarchive.files import File
from internetarchive.search import Search


def setup(subparsers):
    """
    Setup args for download command.

    Args:
        subparsers: subparser object passed from ia.py
    """
    parser = subparsers.add_parser("download",
                                   aliases=["do"],
                                   help="Download files from archive.org",)

    # Main options
    parser.add_argument("identifier",
                        nargs="?",
                        type=str,
                        help="Identifier of the item to download")
    parser.add_argument("file",
                        nargs="*",
                        help="Files to download (only allowed with identifier)")

    # Additional options
    parser.add_argument("-q", "--quiet",
                        action="store_true",
                        help="Turn off ia's output")
    parser.add_argument("-d", "--dry-run",
                        action="store_true",
                        help="Print URLs to stdout and exit")
    parser.add_argument("-i", "--ignore-existing",
                        action="store_true",
                        help="Clobber files already downloaded")
    parser.add_argument("-C", "--checksum",
                        action="store_true",
                        help="Skip files based on checksum")
    parser.add_argument("--checksum-archive",
                        action="store_true",
                        help="Skip files based on _checksum_archive.txt file")
    parser.add_argument("-R", "--retries",
                        type=int,
                        default=5,
                        help="Set number of retries to <retries> (default: 5)")
    parser.add_argument("-I", "--itemlist",
                        type=argparse.FileType("r"),
                        help=("Download items from a specified file. "
                             "Itemlists should be a plain text file with one "
                             "identifier per line"))
    parser.add_argument("-S", "--search",
                        help="Download items returned from a specified search query")
    parser.add_argument("-P", "--search-parameters",
                        nargs=1,
                        action=QueryStringAction,
                        metavar="KEY:VALUE",
                        help="Parameters to send with your --search query. "
                             "Can be specified multiple times.")
    parser.add_argument("-g", "--glob",
                        nargs=1,
                        action="extend",
                        help=("Only download files matching the given glob "
                             "pattern. Can be specified multiple times."))
    parser.add_argument("-e", "--exclude",
                        nargs=1,
                        action="extend",
                        help=("Exclude files matching the given glob pattern. "
                             "Can be specified multiple times. Can only be "
                             "used in conjunction with --glob."))
    parser.add_argument("-f", "--format",
                        nargs=1,
                        action="extend",
                        help=("Only download files of the specified format. "
                             "Can be specified multiple times. You can use the following "
                             "command to retrieve a list of file formats contained within "
                             "a given item: ia metadata --formats <identifier>"))
    parser.add_argument("--on-the-fly",
                        action="store_true",
                        help=("Download on-the-fly files, as well as other "
                             "matching files. on-the-fly files include derivative "
                             "EPUB, MOBI and DAISY files [default: False]"))
    parser.add_argument("--no-directories",
                        action="store_true",
                        help=("Download files into working directory. "
                             "Do not create item directories"))
    parser.add_argument("--destdir",
                        type=validate_dir_path,
                        help=("The destination directory to download files "
                             "and item directories to"))
    parser.add_argument("-s", "--stdout",
                        action="store_true",
                        help="Write file contents to stdout")
    parser.add_argument("--range",
                        dest="ranges",
                        action="append",
                        metavar="[FILE:]START-END",
                        help="Download only the given byte range(s). Requires "
                             "--stdout and can be specified multiple times. A bare "
                             "range (e.g. `0-1023`, `bytes=0-1023`, open-ended "
                             "`1024-`, or suffix `-1024` for the last 1024 bytes) "
                             "applies to the named file; use one file with "
                             "multiple ranges, or one range with multiple files. "
                             "Comma-separated ranges (e.g. `0-1023,4096-8191`) are "
                             "fetched in order, equivalent to repeating --range. To "
                             "pull ranges from different files, bind each explicitly "
                             "as FILE:START-END (e.g. `--range a.warc.gz:0-9 --range "
                             "b.warc.gz:50-99`). Segments are written back-to-back "
                             "with no separator, so e.g. selected .warc.gz records "
                             "can be piped straight to zcat. Disables resume and "
                             "full-file checksum validation for partial fetches.")
    parser.add_argument("--no-change-timestamp",
                        action="store_true",
                        help=("Don't change the timestamp of downloaded files to reflect "
                             "the source material"))
    parser.add_argument("-p", "--parameters",
                        nargs=1,
                        action=QueryStringAction,
                        metavar="KEY:VALUE",
                        help="Parameters to send with your download request. "
                             "Can be specified multiple times.")
    parser.add_argument("--count-views",
                        action="store_true",
                        help="Count this download toward archive.org view "
                             "counts. By default `ia download` opts out via "
                             "`cnt=0`; this flag omits that parameter.")
    parser.add_argument("-a", "--download-history",
                        action="store_true",
                        help="Also download files from the history directory")
    parser.add_argument("--source",
                        nargs=1,
                        action="extend",
                        help=("Filter files based on their source value in files.xml "
                             "(i.e. `original`, `derivative`, `metadata`). "
                             "Can be specified multiple times."))
    parser.add_argument("--exclude-source",
                        nargs=1,
                        action="extend",
                        help=("Exclude files based on their source value in files.xml "
                             "(i.e. `original`, `derivative`, `metadata`). "
                             "Can be specified multiple times."))
    parser.add_argument("-t", "--timeout",
                        type=float,
                        help=("Set a timeout for download requests. "
                             "This sets both connect and read timeout"))

    parser.set_defaults(func=lambda args: main(args, parser))


def normalize_byte_range(value: str) -> str:
    """Normalize a user-supplied ``--range`` value into an HTTP Range header value.

    Accepts ``START-END``, open-ended ``START-``, suffix ``-N`` (the last ``N``
    bytes), or a value already prefixed with ``bytes=``. Returns a string of the
    form ``bytes=...``.

    :param value: The raw ``--range`` argument.
    :returns: A normalized ``bytes=...`` Range header value.
    :raises ValueError: If the range is not a valid single byte range.
    """
    spec = value.strip()
    if spec.lower().startswith("bytes="):
        spec = spec[len("bytes="):]
    if not re.fullmatch(r"\d+-\d*|-\d+", spec):
        raise ValueError(f"invalid byte range: {value!r}")
    start, _, end = spec.partition("-")
    if start and end and int(end) < int(start):
        raise ValueError(f"range end is before range start: {value!r}")
    return f"bytes={spec}"


def parse_byte_ranges(value: str) -> list[str]:
    """Parse a (possibly comma-separated) ``--range`` value into Range headers.

    Accepts the HTTP multi-range syntax: one or more comma-separated byte ranges
    with an optional single leading ``bytes=`` (e.g. ``50-99,123-180`` or
    ``bytes=0-9,-100``). Each range is normalized independently and returned as
    its own ``bytes=...`` value -- one per upstream request, since we expand a
    comma list client-side rather than sending a single multi-range request.

    :param value: The raw ``--range`` argument.
    :returns: One normalized ``bytes=...`` value per comma-separated range.
    :raises ValueError: If the value is empty or any segment is not a valid range.
    """
    spec = value.strip()
    if spec.lower().startswith("bytes="):
        spec = spec[len("bytes="):]
    parts = spec.split(",")
    if any(p.strip() == "" for p in parts):
        raise ValueError(f"invalid byte range: {value!r}")
    return [normalize_byte_range(p) for p in parts]


def positional_files(args: argparse.Namespace) -> list[str]:
    """Return the files named explicitly on the command line.

    Handles both the ``<id> FILE [FILE ...]`` and the ``<id>/path`` forms.

    :param args: Parsed CLI arguments.
    :returns: The list of explicitly-named files (empty if none).
    """
    if args.identifier and args.identifier != "-":
        if "/" in args.identifier:
            return ["/".join(args.identifier.split("/")[1:])]
        return list(args.file)
    return []


def build_range_jobs(
    args: argparse.Namespace,
    parser: argparse.ArgumentParser,
) -> list[tuple[str, str]]:
    """Resolve repeatable ``--range`` values into an ordered job list.

    Each ``--range`` value is either a bare ``[bytes=]START-END`` (which binds to
    the explicitly-named file) or a file-bound ``FILE:START-END``. A value may
    carry several comma-separated ranges (HTTP multi-range syntax, e.g.
    ``0-9,50-99``), which expand to one job each, in order. Returns an ordered
    list of ``(filename, "bytes=...")`` jobs. Any rule violation is reported via
    ``parser.error`` (which exits with status 2).

    :param args: Parsed CLI arguments (``args.ranges`` holds the raw values).
    :param parser: The argument parser, used to emit errors.
    :returns: Ordered ``(filename, range_header_value)`` jobs.
    """
    err = ("expected [bytes=]START-END[,START-END...] or "
           "FILE:START-END[,...]")
    parsed: list[tuple[str | None, str]] = []
    for raw in args.ranges:
        # A bare range never contains a colon (only digits, '-', ',', and an
        # optional leading 'bytes='), so a ':' unambiguously marks the FILE:RANGES
        # form. Split on the last colon so filenames containing a colon still work.
        file_part, sep, range_part = raw.rpartition(":")
        if sep and file_part:
            target, spec = file_part, range_part
        else:
            target, spec = None, raw
        try:
            for rng in parse_byte_ranges(spec):
                parsed.append((target, rng))
        except ValueError:
            parser.error(f"could not parse --range {raw!r}: {err}")

    has_bare = any(f is None for f, _ in parsed)
    has_bound = any(f is not None for f, _ in parsed)
    if has_bare and has_bound:
        parser.error("--range: cannot mix bare ranges (START-END) and file-bound "
                     "ranges (FILE:START-END) in the same command")

    selectors = [name for name, val in (
        ("--glob", args.glob),
        ("--format", args.format),
        ("--source", args.source),
        ("--exclude-source", args.exclude_source),
        ("--search", args.search),
        ("--itemlist", args.itemlist),
    ) if val]
    if selectors:
        parser.error(f"--range cannot be combined with {', '.join(selectors)}")

    if has_bound:
        if positional_files(args):
            parser.error("--range FILE:START-END cannot be combined with positional "
                         "file arguments; the files come from the --range values")
        return [(str(f), b) for f, b in parsed]

    # Bare form: bind ranges to the positional file(s).
    files = positional_files(args)
    if not files:
        first = args.ranges[0]
        parser.error(f"--range {first} needs a file: name one positionally "
                     f"(ia download ID FILE --range {first}) or bind it "
                     f"(--range FILE:{first})")
    if len(files) > 1 and len(parsed) > 1:
        parser.error(f"ambiguous --range: {len(parsed)} ranges and {len(files)} files "
                     "given. With more than one file, bind each range to its file "
                     "explicitly, e.g. `--range FILE1:0-9 --range FILE2:55-99`.")
    if len(parsed) == 1:
        rng = parsed[0][1]
        return [(f, rng) for f in files]
    return [(files[0], b) for _, b in parsed]


def validate_args(args: argparse.Namespace, parser: argparse.ArgumentParser) -> None:
    if args.itemlist and args.search:
        parser.error("--itemlist and --search cannot be used together")

    if args.itemlist or args.search:
        if args.identifier:
            parser.error("Cannot specify an identifier with --itemlist/--search")
        if args.file:
            parser.error("Cannot specify files with --itemlist/--search")
    else:
        if not args.identifier:
            parser.error("Identifier is required when not using --itemlist/--search")

    if args.ranges:
        if not args.stdout:
            parser.error("--range requires --stdout")
        if args.identifier == "-":
            parser.error("--range cannot be combined with reading identifiers "
                         "from stdin")
        args.range_jobs = build_range_jobs(args, parser)
    else:
        args.range_jobs = None


def main(args: argparse.Namespace, parser: argparse.ArgumentParser) -> None:
    """
    Main entry point for 'ia download'.
    """
    args.search_parameters = args.search_parameters or {}
    args.parameters = args.parameters or {}

    ids: list[File | str] | Search | TextIO
    validate_args(args, parser)

    if args.itemlist:
        ids = [x.strip() for x in args.itemlist if x.strip()]
        if not ids:
            parser.error("--itemlist file is empty or contains only whitespace")
        total_ids = len(ids)
    elif args.search:
        try:
            _search = args.session.search_items(args.search,
                                                params=args.search_parameters)
            total_ids = _search.num_found
            if total_ids == 0:
                print(f"error: the query '{args.search}' returned no results", file=sys.stderr)
                sys.exit(1)
            ids = _search
        except ValueError as e:
            print(f"error: {e}", file=sys.stderr)
            sys.exit(1)

    # Download specific files.
    if args.identifier and args.identifier != "-":
        if "/" in args.identifier:
            identifier = args.identifier.split("/")[0]
            files = ["/".join(args.identifier.split("/")[1:])]
        else:
            identifier = args.identifier
            files = args.file
        total_ids = 1
        ids = [identifier]
    elif args.identifier == "-":
        total_ids = 1
        ids = sys.stdin
        files = None
    else:
        files = None

    errors = []
    for i, identifier in enumerate(ids):
        try:
            identifier = identifier.strip()
        except AttributeError:
            identifier = identifier.get("identifier")
        if total_ids > 1:
            item_index = f"{i + 1}/{total_ids}"
        else:
            item_index = None

        try:
            item = args.session.get_item(identifier)
        except Exception as exc:
            print(f"{identifier}: failed to retrieve item metadata - errors", file=sys.stderr)
            if "You are attempting to make an HTTPS" in str(exc):
                print(f"\n{exc}", file=sys.stderr)
                sys.exit(1)
            else:
                continue

        # Otherwise, download the entire item.
        ignore_history_dir = not args.download_history
        _errors = item.download(
            files=files,
            formats=args.format,
            glob_pattern=args.glob,
            exclude_pattern=args.exclude,
            dry_run=args.dry_run,
            verbose=not args.quiet,
            ignore_existing=args.ignore_existing,
            checksum=args.checksum,
            checksum_archive=args.checksum_archive,
            destdir=args.destdir,
            no_directory=args.no_directories,
            retries=args.retries,
            item_index=item_index,
            ignore_errors=True,
            on_the_fly=args.on_the_fly,
            no_change_timestamp=args.no_change_timestamp,
            params=args.parameters,
            ignore_history_dir=ignore_history_dir,
            source=args.source,
            exclude_source=args.exclude_source,
            stdout=args.stdout,
            timeout=args.timeout,
            count_views=args.count_views,
            range_jobs=args.range_jobs,
        )
        if _errors:
            errors.append(_errors)
    if errors:
        # TODO: add option for a summary/report.
        sys.exit(1)
    else:
        sys.exit(0)
