#!/usr/bin/env python3
"""
Tool to send a download to an aria2c daemon via RPC.
"""

import argparse
import json
import logging
import sys
import os
import glob
import shutil
import re
import shlex
import time
from urllib.parse import urljoin

try:
    import aria2p
    from aria2p.utils import human_readable_bytes
except ImportError:
    print("aria2p is not installed. Please install it with: pip install aria2p", file=sys.stderr)
    sys.exit(1)

logger = logging.getLogger('download_aria_tool')

class TimeoutError(Exception):
    pass


def add_download_aria_parser(subparsers):
    """Add the parser for the 'download aria-rpc' command."""
    parser = subparsers.add_parser(
        'aria-rpc',
        description='Send a download to an aria2c daemon via RPC, using an info.json from stdin or a file.',
        formatter_class=argparse.RawTextHelpFormatter,
        help='Download a specific format using aria2c RPC.',
        epilog="""
Usage Notes for Fragmented Downloads (e.g., DASH):

To download and automatically merge fragmented formats, you must:
1. Use '--wait' to make the operation synchronous.
2. Use '--auto-merge-fragments' to enable the merge logic.
3. Ensure this script has access to the directory where aria2c saves files.

Example for a remote aria2c daemon:
  - The remote daemon saves files to '/srv/downloads' on its machine.
  - This directory is mounted locally at '/mnt/remote_aria2_downloads'.

cat latest-info.json | yt-ops-client download aria-rpc -f "299/137" \\
  --wait --auto-merge-fragments \\
  --remote-dir /srv/downloads \\
  --fragments-dir /mnt/remote_aria2_downloads
"""
    )
    parser.add_argument('--load-info-json', type=argparse.FileType('r', encoding='utf-8'), help="Path to the info.json file. If not provided, reads from stdin.")
    parser.add_argument('-f', '--format', required=True, help='The format ID to download. Supports yt-dlp style format selectors (e.g., "137/136,140").')
    parser.add_argument('--output-dir', help='Local directory to save the final merged file. Defaults to the current directory.')
    parser.add_argument('--fragments-dir', help='The local path where this script should look for downloaded fragments. If the aria2c daemon is remote, this should be a local mount point corresponding to --remote-dir. Defaults to --output-dir.')
    parser.add_argument('--remote-dir', help='The absolute path to the download directory on the remote aria2c host. This is passed via RPC.')
    parser.add_argument('--aria-host', default='localhost', help='The host of the aria2c RPC server. Default: localhost.')
    parser.add_argument('--aria-port', type=int, default=6800, help='The port of the aria2c RPC server. Default: 6800.')
    parser.add_argument('--aria-secret', help='The secret token for the aria2c RPC server (often required, e.g., "SQGCQPLVFQIASMPNPOJYLVGJYLMIDIXDXAIXOTX").')
    parser.add_argument('--proxy', help='Proxy to use for the download, e.g., "socks5://127.0.0.1:1080".')
    parser.add_argument('--downloader-args', help='Arguments for aria2c, in yt-dlp format (e.g., "aria2c:[-x 8, -k 1M]").')
    parser.add_argument('--wait', action='store_true', help='Wait for the download to complete and report its status. Note: This makes the operation synchronous and will block until the download finishes.')
    parser.add_argument('--wait-timeout', help='Timeout in seconds for waiting on downloads. Use "auto" to calculate based on a minimum speed of 200KiB/s. Requires --wait. Default: no timeout.')
    parser.add_argument('--auto-merge-fragments', action='store_true', help='Automatically merge fragments after download. Requires --wait and assumes the script has filesystem access to the aria2c host.')
    parser.add_argument('--remove-fragments-after-merge', action='store_true', help='Delete individual fragment files after a successful merge. Requires --auto-merge-fragments.')
    parser.add_argument('--cleanup', action='store_true', help='After a successful download, remove the final file(s) from the filesystem. For fragmented downloads, this implies --remove-fragments-after-merge.')
    parser.add_argument('--remove-on-complete', action=argparse.BooleanOptionalAction, default=True, help='Remove the download from aria2c history on successful completion. Use --no-remove-on-complete to disable. May fail on older aria2c daemons.')
    parser.add_argument('--purge-on-complete', action='store_true', help='Use aria2.purgeDownloadResult to clear ALL completed/failed downloads from history on success. Use as a workaround for older daemons.')
    parser.add_argument('--verbose', action='store_true', help='Enable verbose output for this script.')
    return parser

def cleanup_aria_download(api, downloads):
    """Pause and remove downloads from aria2c."""
    if not downloads:
        return
    try:
        logger.info(f"Attempting to clean up {len(downloads)} download(s) from aria2c...")
        # Filter out downloads that might already be gone
        valid_downloads = [d for d in downloads if hasattr(d, 'gid')]
        if not valid_downloads:
            logger.info("No valid downloads to clean up.")
            return
        api.pause(valid_downloads)
        # Give aria2c a moment to process the pause command before removing
        time.sleep(0.5)
        api.remove(valid_downloads)
        logger.info("Cleanup successful.")
    except Exception as e:
        logger.warning(f"An error occurred during aria2c cleanup: {e}")


def parse_aria_error(download):
    """Parses an aria2p Download object to get a detailed error message."""
    error_code = download.error_code
    error_message = download.error_message

    if not error_message:
        return f"Unknown aria2c error (Code: {error_code})"

    # Check for common HTTP errors in the message
    http_status_match = re.search(r'HTTP status (\d+)', error_message)
    if http_status_match:
        status_code = int(http_status_match.group(1))
        if status_code == 403:
            return f"HTTP Error 403: Forbidden. The URL may have expired or requires valid cookies/headers."
        elif status_code == 404:
            return f"HTTP Error 404: Not Found. The resource is unavailable."
        else:
            return f"HTTP Error {status_code}."

    if "Timeout" in error_message or "timed out" in error_message.lower():
        return "Download timed out."

    # Fallback to the raw error message
    return f"Aria2c error (Code: {error_code}): {error_message}"


def parse_aria_args_to_options(args_str):
    """
    Parses yt-dlp style downloader args for aria2c.
    Example: "aria2c:[-x 8, -k 1M]" or just "-x 8 -k 1M"
    Returns a dictionary of options for aria2p.
    """
    if not args_str or not args_str.strip():
        return {}

    inner_args_str = args_str.strip()
    match = re.match(r'aria2c:\s*\[(.*)\]', inner_args_str)
    if match:
        # Handle yt-dlp's format
        inner_args_str = match.group(1).replace(',', ' ')
    else:
        # If it doesn't match, assume the whole string is a set of arguments.
        logger.debug(f"Downloader args '{args_str}' does not match 'aria2c:[...]' format. Parsing as a raw argument string.")

    arg_list = shlex.split(inner_args_str)

    # Use a mini-parser to handle CLI-style args
    parser = argparse.ArgumentParser(add_help=False, prog="aria2c_args_parser")
    parser.add_argument('-x', '--max-connection-per-server')
    parser.add_argument('-k', '--min-split-size')
    parser.add_argument('-s', '--split')
    parser.add_argument('--all-proxy')
    
    try:
        # We only care about known arguments
        known_args, unknown_args = parser.parse_known_args(arg_list)
        if unknown_args:
            logger.warning(f"Ignoring unknown arguments in --downloader-args: {unknown_args}")
        # Convert to dict, removing None values
        return {k: v for k, v in vars(known_args).items() if v is not None}
    except Exception:
        logger.warning(f"Failed to parse arguments inside --downloader-args: '{inner_args_str}'")
        return {}


def main_download_aria(args):
    """Main logic for the 'download-aria' command."""
    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=log_level, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', stream=sys.stderr)

    if args.remove_fragments_after_merge and not args.auto_merge_fragments:
        logger.error("--remove-fragments-after-merge requires --auto-merge-fragments.")
        return 1
    if args.auto_merge_fragments and not args.wait:
        logger.error("--auto-merge-fragments requires --wait.")
        return 1
    if args.wait_timeout and not args.wait:
        logger.error("--wait-timeout requires --wait.")
        return 1

    if args.wait:
        logger.info("Will wait for download to complete and report status. This is a synchronous operation.")
    else:
        logger.info("Will submit download and exit immediately (asynchronous).")

    info_json_content = ""
    input_source_name = ""
    if args.load_info_json:
        info_json_content = args.load_info_json.read()
        input_source_name = args.load_info_json.name
    else:
        info_json_content = sys.stdin.read()
        input_source_name = "stdin"

    if not info_json_content.strip():
        logger.error(f"Failed to read info.json from {input_source_name}. Input is empty.")
        return 1

    try:
        info_data = json.loads(info_json_content)
        logger.info(f"Successfully loaded info.json from {input_source_name}.")
    except json.JSONDecodeError:
        logger.error(f"Failed to parse info.json from {input_source_name}. Is the input valid JSON?")
        return 1

    # Find the requested format, supporting yt-dlp style selectors
    target_format = None
    # A format selector can be a comma-separated list of preferences,
    # where each preference can be a slash-separated list of format_ids.
    # e.g., "299/137/136,140" means try 299, then 137, then 136, then 140.
    format_preferences = [item.strip() for sublist in (i.split('/') for i in args.format.split(',')) for item in sublist if item.strip()]

    available_formats_map = {f['format_id']: f for f in info_data.get('formats', []) if 'format_id' in f}

    for format_id in format_preferences:
        if format_id in available_formats_map:
            target_format = available_formats_map[format_id]
            logger.info(f"Selected format ID '{format_id}' from selector '{args.format}'.")
            break

    if not target_format:
        logger.error(f"No suitable format found for selector '{args.format}' in info.json.")
        return 1

    # Get file size for auto-timeout and dynamic options
    total_filesize = target_format.get('filesize') or target_format.get('filesize_approx')

    # Construct filename
    video_id = info_data.get('id', 'unknown_video_id')
    title = info_data.get('title', 'unknown_title')
    ext = target_format.get('ext', 'mp4')
    # Sanitize title for filename
    safe_title = "".join([c for c in title if c.isalpha() or c.isdigit() or c in (' ', '-', '_')]).rstrip()
    filename = f"{safe_title} [{video_id}].f{target_format['format_id']}.{ext}"

    # Prepare options for aria2
    aria_options = {
        # Options from yt-dlp's aria2c integration for performance and reliability
        'max-connection-per-server': 16,
        'split': 16,
        'min-split-size': '1M',
        'http-accept-gzip': 'true',
        'file-allocation': 'none',
    }

    if args.proxy:
        aria_options['all-proxy'] = args.proxy

    custom_options = parse_aria_args_to_options(args.downloader_args)

    # Dynamically set min-split-size if not overridden by user
    if 'min_split_size' not in custom_options and total_filesize:
        if total_filesize > 100 * 1024 * 1024:  # 100 MiB
            aria_options['min-split-size'] = '5M'
            logger.info("File is > 100MiB, dynamically setting min-split-size to 5M.")

    if custom_options:
        aria_options.update(custom_options)
        logger.info(f"Applied custom aria2c options from --downloader-args: {custom_options}")

    aria_options['out'] = filename

    # Add headers from info.json, mimicking yt-dlp's behavior for aria2c
    headers = target_format.get('http_headers')
    if headers:
        header_list = [f'{key}: {value}' for key, value in headers.items()]
        aria_options['header'] = header_list
        logger.info(f"Adding {len(header_list)} HTTP headers to the download.")
        if args.verbose:
            for h in header_list:
                if h.lower().startswith('cookie:'):
                    logger.debug(f"  Header: Cookie: [REDACTED]")
                else:
                    logger.debug(f"  Header: {h}")

    is_fragmented = 'fragments' in target_format
    if not is_fragmented:
        url = target_format.get('url')
        if not url:
            logger.error(f"Format ID '{args.format}' has neither a URL nor fragments.")
            return 1

    try:
        logger.info(f"Connecting to aria2c RPC at http://{args.aria_host}:{args.aria_port}")
        client = aria2p.Client(
            host=f"http://{args.aria_host}",
            port=args.aria_port,
            secret=args.aria_secret or ""
        )
        api = aria2p.API(client)

        timeout_seconds = None
        if args.wait_timeout:
            if args.wait_timeout.lower() == 'auto':
                if total_filesize:
                    # Min speed: 200 KiB/s. Min timeout: 30s.
                    min_speed = 200 * 1024
                    calculated_timeout = int(total_filesize / min_speed)
                    timeout_seconds = max(30, calculated_timeout)
                    total_filesize_hr, _ = human_readable_bytes(total_filesize)
                    logger.info(f"Auto-calculated timeout: {timeout_seconds}s (based on {total_filesize_hr} at 200KiB/s).")
                else:
                    logger.warning("Cannot use 'auto' timeout: file size not available in info.json. Timeout disabled.")
            else:
                try:
                    timeout_seconds = int(args.wait_timeout)
                    if timeout_seconds <= 0:
                        raise ValueError
                except ValueError:
                    logger.error(f"Invalid --wait-timeout value: '{args.wait_timeout}'. Must be a positive integer or 'auto'.")
                    return 1

        if is_fragmented:
            return download_fragments_aria(args, api, target_format, filename, aria_options, timeout_seconds, remote_dir=args.remote_dir)
        else:
            return download_url_aria(args, api, url, filename, aria_options, timeout_seconds, remote_dir=args.remote_dir)

    except Exception as e:
        logger.error(f"An error occurred while communicating with aria2c: {e}", exc_info=args.verbose)
        return 1

def download_url_aria(args, api, url, filename, aria_options, timeout_seconds, remote_dir=None):
    """Handle downloading a single URL with aria2c."""
    if remote_dir:
        aria_options['dir'] = remote_dir
    logger.info(f"Adding download for format '{args.format}' with URL: {url[:70]}...")
    downloads = api.add_uris([url], options=aria_options)

    if not downloads:
        logger.error("Failed to add download to aria2c. The API returned an empty result.")
        return 1

    # Handle older aria2p versions that return a single Download object instead of a list
    download = downloads[0] if isinstance(downloads, list) else downloads
    logger.info(f"Successfully added download to aria2c. GID: {download.gid}")

    if args.wait:
        logger.info(f"Waiting for download {download.gid} to complete...")
        start_time = time.time()
        try:
            while True:
                if timeout_seconds and (time.time() - start_time > timeout_seconds):
                    raise TimeoutError(f"Download did not complete within {timeout_seconds}s timeout.")

                # Re-fetch the download object to get the latest status
                download.update()
                # A download is no longer active if it's complete, errored, paused, or removed.
                if download.status not in ('active', 'waiting'):
                    break

                progress_info = (
                    f"\rGID {download.gid}: {download.status} "
                    f"{download.progress_string()} "
                    f"({download.download_speed_string()}) "
                    f"ETA: {download.eta_string()}"
                )
                sys.stdout.write(progress_info)
                sys.stdout.flush()
                time.sleep(0.5)
        except (KeyboardInterrupt, TimeoutError) as e:
            sys.stdout.write('\n')
            if isinstance(e, KeyboardInterrupt):
                logger.warning("Wait interrupted by user. Cleaning up download...")
                cleanup_aria_download(api, [download])
                return 130
            else:  # TimeoutError
                logger.error(f"Download timed out. Cleaning up... Error: {e}")
                cleanup_aria_download(api, [download])
                return 1
        except aria2p.ClientException as e:
            # This can happen if the download completes and is removed by aria2c
            # before we can check its final status. Assume success in this case.
            logger.warning(f"Could not get final status for GID {download.gid} (maybe removed on completion?): {e}. Assuming success.")
            print(f"Download for GID {download.gid} presumed successful.")
            return 0

        sys.stdout.write('\n')  # Newline after progress bar

        # Final status check (no need to update again, we have the latest status)
        if download.status == 'complete':
            logger.info(f"Download {download.gid} completed successfully.")
            
            downloaded_filepath_remote = None
            if download.files:
                downloaded_filepath_remote = download.files[0].path
                print(f"Download successful: {downloaded_filepath_remote}")
            else:
                print("Download successful, but no file path reported by aria2c.")

            if args.cleanup and downloaded_filepath_remote:
                local_filepath = None
                # To map remote path to local, we need remote_dir and a local equivalent.
                # We'll use fragments_dir as the local equivalent, which defaults to output_dir.
                local_base_dir = args.fragments_dir or args.output_dir or '.'
                if remote_dir:
                    if downloaded_filepath_remote.startswith(remote_dir):
                        relative_path = os.path.relpath(downloaded_filepath_remote, remote_dir)
                        local_filepath = os.path.join(local_base_dir, relative_path)
                    else:
                        logger.warning(f"Cleanup: Downloaded file path '{downloaded_filepath_remote}' does not start with remote-dir '{remote_dir}'. Cannot map to local path.")
                else:
                    logger.warning(f"Cleanup: --remote-dir not specified. Assuming download path is accessible locally as '{downloaded_filepath_remote}'.")
                    local_filepath = downloaded_filepath_remote

                if local_filepath:
                    try:
                        if os.path.exists(local_filepath):
                            os.remove(local_filepath)
                            logger.info(f"Cleanup: Removed downloaded file '{local_filepath}'")
                        else:
                            logger.warning(f"Cleanup: File not found at expected local path '{local_filepath}'. Skipping removal.")
                    except OSError as e:
                        logger.error(f"Cleanup failed: Could not remove file '{local_filepath}': {e}")
            elif args.cleanup:
                logger.warning("Cleanup requested, but no downloaded file path was reported by aria2c.")

            if args.purge_on_complete:
                try:
                    api.purge_download_result()
                    logger.info("Purged all completed/failed downloads from aria2c history.")
                except Exception as e:
                    logger.warning(f"Failed to purge download history: {e}")
            elif args.remove_on_complete:
                try:
                    api.remove_download_result(download)
                    logger.info(f"Removed download {download.gid} from aria2c history.")
                except Exception as e:
                    logger.warning(f"Failed to remove download {download.gid} from history: {e}")

            return 0
        else:
            detailed_error = parse_aria_error(download)
            logger.error(f"Download {download.gid} failed. Error: {detailed_error}")
            return 1
    else:
        print(f"Successfully added download. GID: {download.gid}")
        return 0

def download_fragments_aria(args, api, target_format, filename, aria_options, timeout_seconds, remote_dir=None):
    """Handle downloading fragmented formats with aria2c."""
    logger.info(f"Format '{args.format}' is fragmented. Adding all fragments to download queue.")
    fragment_base_url = target_format.get('fragment_base_url')
    fragments = target_format['fragments']

    MAX_FRAGMENTS = 50000
    if len(fragments) > MAX_FRAGMENTS:
        logger.error(
            f"The number of fragments ({len(fragments)}) exceeds the safety limit of {MAX_FRAGMENTS}. "
            f"This is to prevent overwhelming the aria2c server. Aborting."
        )
        return 1

    # We need to set the 'dir' option for all fragments if specified.
    # The 'out' option will be set per-fragment.
    frag_aria_options = aria_options.copy()
    frag_aria_options.pop('out', None)  # Remove the main 'out' option

    if remote_dir:
        frag_aria_options['dir'] = remote_dir
        logger.info(f"Instructing remote aria2c to save fragments to: {remote_dir}")

    base_filename, file_ext = os.path.splitext(filename)

    calls = []
    for i, fragment in enumerate(fragments):
        frag_url = fragment.get('url')
        if not frag_url:
            if not fragment_base_url:
                logger.error(f"Fragment {i} has no URL and no fragment_base_url is available. Aborting.")
                return 1
            frag_url = urljoin(fragment_base_url, fragment['path'])

        # Use the base filename from the main file, but add fragment identifier
        fragment_filename = f"{base_filename}-Frag{i}{file_ext}"

        current_frag_options = frag_aria_options.copy()
        current_frag_options['out'] = os.path.basename(fragment_filename)

        # Prepare parameters for multicall in the format:
        # {"methodName": "aria2.addUri", "params": [["url"], {"out": "file.mp4"}]}
        # The secret token is automatically added by aria2p.
        params = [[frag_url], current_frag_options]
        call_struct = {
            "methodName": api.client.ADD_URI,
            "params": params
        }
        calls.append(call_struct)

    results = api.client.multicall(calls)
    if not results:
        logger.error("Failed to add fragments to aria2c. The API returned an empty result.")
        return 1

    # The result of a multicall of addUri is a list of lists, where each inner list
    # contains the GID of one download, e.g., [['gid1'], ['gid2']].
    # A failed call for a fragment may result in a fault struct dict instead of a list.
    # We extract GIDs from successful calls.
    gids = [result[0] for result in results if isinstance(result, list) and result]

    if len(gids) != len(fragments):
        failed_count = len(fragments) - len(gids)
        logger.warning(f"{failed_count} out of {len(fragments)} fragments failed to be added to aria2c.")

    if not gids:
        logger.error("Failed to add any fragments to aria2c. All submissions failed.")
        return 1

    logger.info(f"Successfully added {len(gids)} fragments to aria2c.")
    if args.verbose:
        logger.debug(f"GIDs: {gids}")

    if args.wait:
        logger.info(f"Waiting for {len(gids)} fragments to complete...")
        start_time = time.time()
        downloads_to_cleanup = []
        try:
            while True:
                if timeout_seconds and (time.time() - start_time > timeout_seconds):
                    raise TimeoutError(f"Fragment downloads did not complete within {timeout_seconds}s timeout.")

                downloads = api.get_downloads(gids)
                downloads_to_cleanup = downloads  # Store for potential cleanup
                # A download is considered "active" if it's currently downloading or waiting in the queue.
                # It is "not active" if it is complete, errored, paused, or removed.
                active_downloads = [d for d in downloads if d.status in ('active', 'waiting')]
                if not active_downloads:
                    break  # All downloads are complete or have stopped for other reasons

                for d in active_downloads:
                    d.update()

                completed_count = len(downloads) - len(active_downloads)
                total_bytes = sum(d.total_length for d in downloads)
                downloaded_bytes = sum(d.completed_length for d in downloads)
                total_speed = sum(d.download_speed for d in downloads)
                progress_percent = (downloaded_bytes / total_bytes * 100) if total_bytes > 0 else 0

                progress_info = (
                    f"\rProgress: {completed_count}/{len(downloads)} fragments | "
                    f"{progress_percent:.1f}% "
                    f"({human_readable_bytes(downloaded_bytes)}/{human_readable_bytes(total_bytes)}) "
                    f"Speed: {human_readable_bytes(total_speed)}/s"
                )
                sys.stdout.write(progress_info)
                sys.stdout.flush()
                time.sleep(0.5)
        except (KeyboardInterrupt, TimeoutError) as e:
            sys.stdout.write('\n')
            if isinstance(e, KeyboardInterrupt):
                logger.warning("Wait interrupted by user. Cleaning up fragments...")
                cleanup_aria_download(api, downloads_to_cleanup)
                return 130
            else:  # TimeoutError
                logger.error(f"Download timed out. Cleaning up fragments... Error: {e}")
                cleanup_aria_download(api, downloads_to_cleanup)
                return 1
        except aria2p.ClientException as e:
            # This can happen if downloads complete and are removed by aria2c
            # before we can check their final status. Assume success in this case.
            logger.warning(f"Could not get final status for some fragments (maybe removed on completion?): {e}. Assuming success.")

        sys.stdout.write('\n')

        # Final status check
        failed_downloads = []
        try:
            downloads = api.get_downloads(gids)
            failed_downloads = [d for d in downloads if d.status != 'complete']
        except aria2p.ClientException as e:
            logger.warning(f"Could not perform final status check for fragments (maybe removed on completion?): {e}. Assuming success.")
            # If we can't check, we assume success based on the earlier wait loop not failing catastrophically.
            failed_downloads = []

        if failed_downloads:
            logger.error(f"{len(failed_downloads)} fragments failed to download.")
            for d in failed_downloads:
                detailed_error = parse_aria_error(d)
                logger.error(f"  GID {d.gid}: {detailed_error}")
            return 1
        else:
            logger.info("All fragments downloaded successfully.")
            output_dir = args.output_dir or '.'
            final_filepath = os.path.join(output_dir, filename)
            fragments_lookup_dir = args.fragments_dir or output_dir

            if args.auto_merge_fragments:
                logger.info(f"Attempting to merge fragments into: {final_filepath}")
                logger.info(f"Searching for fragments in local directory: {os.path.abspath(fragments_lookup_dir)}")
                try:
                    # base_filename and file_ext are available from earlier in the function
                    # We must escape the base filename in case it contains glob special characters like [ or ].
                    escaped_base = glob.escape(base_filename)
                    search_path = os.path.join(fragments_lookup_dir, f"{escaped_base}-Frag*{file_ext}")
                    fragment_files = glob.glob(search_path)

                    if not fragment_files:
                        logger.error(f"No fragment files found with pattern: {search_path}")
                        return 1

                    def fragment_sort_key(f):
                        match = re.search(r'Frag(\d+)', os.path.basename(f))
                        return int(match.group(1)) if match else -1
                    fragment_files.sort(key=fragment_sort_key)

                    with open(final_filepath, 'wb') as dest_file:
                        for frag_path in fragment_files:
                            with open(frag_path, 'rb') as src_file:
                                shutil.copyfileobj(src_file, dest_file)

                    logger.info(f"Successfully merged {len(fragment_files)} fragments into {final_filepath}")

                    if args.remove_fragments_after_merge or args.cleanup:
                        logger.info("Removing fragment files...")
                        for frag_path in fragment_files:
                            os.remove(frag_path)
                        logger.info("Fragment files removed.")
                    
                    if args.cleanup:
                        try:
                            os.remove(final_filepath)
                            logger.info(f"Cleanup: Removed merged file '{final_filepath}'")
                        except OSError as e:
                            logger.error(f"Cleanup failed: Could not remove merged file '{final_filepath}': {e}")

                    print(f"Download and merge successful: {final_filepath}")

                    if args.purge_on_complete:
                        try:
                            api.purge_download_result()
                            logger.info("Purged all completed/failed downloads from aria2c history.")
                        except Exception as e:
                            logger.warning(f"Failed to purge download history: {e}")
                    elif args.remove_on_complete:
                        try:
                            # The `downloads` variable from the last status check should be valid here.
                            api.remove_download_result(downloads)
                            logger.info(f"Removed {len(downloads)} fragment downloads from aria2c history.")
                        except aria2p.ClientException as e:
                            logger.warning(f"Could not remove fragment downloads from history (maybe already gone?): {e}")
                        except Exception as e:
                            logger.warning(f"Failed to remove fragment downloads from history: {e}")
                    
                    return 0

                except Exception as e:
                    logger.error(f"An error occurred during merging: {e}", exc_info=args.verbose)
                    logger.error("Fragments were downloaded but not merged.")
                    return 1
            else:
                print("Download successful. Fragments now need to be merged manually.")
                print(f"The final merged file should be named: {final_filepath}")
                print("You can merge them with a command like:")
                print(f"  cat `ls -v '{os.path.join(fragments_lookup_dir, base_filename)}'-Frag*'{file_ext}'` > '{final_filepath}'")

                if args.cleanup:
                    logger.info("Cleanup requested. Removing downloaded fragments...")
                    try:
                        # base_filename and file_ext are available from earlier in the function
                        escaped_base = glob.escape(base_filename)
                        search_path = os.path.join(fragments_lookup_dir, f"{escaped_base}-Frag*{file_ext}")
                        fragment_files = glob.glob(search_path)

                        if not fragment_files:
                            logger.warning(f"Cleanup: No fragment files found with pattern: {search_path}")
                        else:
                            for frag_path in fragment_files:
                                os.remove(frag_path)
                            logger.info(f"Removed {len(fragment_files)} fragment files.")
                    except Exception as e:
                        logger.error(f"An error occurred during fragment cleanup: {e}", exc_info=args.verbose)

                if args.purge_on_complete:
                    try:
                        api.purge_download_result()
                        logger.info("Purged all completed/failed downloads from aria2c history.")
                    except Exception as e:
                        logger.warning(f"Failed to purge download history: {e}")
                elif args.remove_on_complete:
                    try:
                        # The `downloads` variable from the last status check should be valid here.
                        api.remove_download_result(downloads)
                        logger.info(f"Removed {len(downloads)} fragment downloads from aria2c history.")
                    except aria2p.ClientException as e:
                        logger.warning(f"Could not remove fragment downloads from history (maybe already gone?): {e}")
                    except Exception as e:
                        logger.warning(f"Failed to remove fragment downloads from history: {e}")

                return 0
    else:
        print(f"Successfully added {len(gids)} fragments. GIDs: {gids}")
        print("These fragments will need to be merged manually after download.")
        return 0