import asyncio
import json
from hashlib import sha256
from os import utime
from pathlib import Path
from queue import Queue

import click
import httpx
from httpx import HTTPStatusError
from rich.console import Console
from rich import filesize

from bulk_restore_tool.__about__ import __version__
from bulk_restore_tool.client import Client, DetailType
from bulk_restore_tool.models import (
    RestoreJob,
    FileSelection,
    RestorableFile,
    Archive,
)
from bulk_restore_tool.db import (
    save_archive_metadata_sqlite,
    restore_queue,
)
from bulk_restore_tool.progress import ProgressBar

console = Console()


def store_in_context(ctx, param, arg):
    """
    Callback for `click.option`s that stores option value in the `ctx.obj` (also ensures this is a dict).

    Useful for enabling "global" options that can be passed at any stage of the command or sub-commands.
    """
    if arg:
        ctx.ensure_object(dict)
        ctx.obj.update({param.name: arg})
    return arg


log_file_option = click.option(
    "--log-file",
    help="Specify location where logging gets written. Can also be set by BRT_LOG_FILE environment variable. "
    "Default='bulk_restore.log' in current directory.",
    callback=store_in_context,
)
log_level_option = click.option(
    "--log-level",
    type=click.Choice(["WARN", "INFO", "DEBUG"]),
    help="Specify logging level. Can also be set by BRT_LOG_LEVEL environment variable. Defaults to INFO.",
    callback=store_in_context,
)


def logging_options(f):
    f = log_file_option(f)
    f = log_level_option(f)
    return f


@click.group(
    context_settings={
        "help_option_names": ["-h", "--help"],
        "max_content_width": 200,
    },
    invoke_without_command=True,
)
@click.option(
    "--version", is_flag=True, help="Print the currently installed version of the tool."
)
@logging_options
def brt(version, **kwargs):
    """
    \b
    Code42 Bulk Restore Tool

    To authenticate, set the following variables either in your shell environment or an .env file in your
    current working directory:

    \b
        BRT_URL
        BRT_API_CLIENT_ID
        BRT_API_CLIENT_SECRET

    For details on API Client authentication, see: https://support.code42.com/Incydr/Admin/Code42_console_reference/API_clients
    """
    if version:
        console.print(__version__)
        return


@brt.command()
@click.argument("id")
@click.option(
    "--name",
    help="Specify a name for the restore job. This name will be visible in audit log entries for the restore session. "
    "Default = A generated UUID.",
)
@click.option(
    "--type",
    "type_",
    type=click.Choice([d.value for d in DetailType]),
    default="archiveId",
    help="Specify what type of identifier to gather archive metadata for. Default = archiveId",
)
@click.option(
    "--target",
    type=click.Path(),
    help="Specify target folder where archive metadata will be stored (will be created if it doesn't exist). "
    "Default = Directory `brt` is executed in.",
)
@click.option(
    "--include-pattern",
    multiple=True,
    type=str,
    help="Specify regex pattern to select which file paths to include in the restore. Can be passed multiple times.",
)
@click.option(
    "--exclude-pattern",
    multiple=True,
    type=str,
    help="Specify regex pattern to select which file paths to exclude in the restore. Can be passed multiple times. A "
    "match on an exclude pattern will take precedence over any matching include patterns.",
)
@click.option(
    "--include-deleted",
    is_flag=True,
    help="Include files that have been marked as deleted in the archive file metadata. Default = false",
)
@click.option(
    "--max-deleted-date",
    type=click.DateTime(),
    help="Specify date for maximum age of deleted files to restore. Any files that were marked as deleted before this "
    "date will be excluded from restore. Default = None",
)
@click.option(
    "--as-of-date",
    type=click.DateTime(),
    help="Restore files as of specified date. Default = None",
)
@click.option(
    "--backup-set-ids",
    type=int,
    multiple=True,
    help="Specify which backup sets from the archive should be included, can be passed multiple times to include "
    "files from multiple sets. Default = Include files from all backup sets.",
)
@click.option(
    "--no-calculate",
    is_flag=True,
    help="Skip gathering of archive metadata and calculation of restore totals and only create job file. "
    "Default = false",
)
@logging_options
def create_job(
    id,
    target,
    type_,
    include_pattern,
    exclude_pattern,
    include_deleted,
    backup_set_ids,
    max_deleted_date,
    as_of_date,
    log_level,
    log_file,
    name,
    no_calculate,
):
    """
    Creates a restore job for a set of archives.

    Archives can be included in the job by specifying the identifier type (archiveId, deviceId, username, or legalHoldId)
    via the --type option (defaults to archiveId) and providing a comma-separated list of identifiers as the main command-line
    argument. Some examples:

        \b
        brt create-job --type archiveId id1,id2,id3
        brt create-job --type username user@example.com
        brt create-job --type deviceId 1234,4567
        brt create-job --type legalHoldId 7890

    If the --name option is not provided, a new UUID will be created as the job name and the parameters for the job will
    be saved in restore_<job_name>.json

    A <job_name> folder will also be created in the --target path to store the archive metadata and restored files (if
    --target option is not provided, it defaults to the current working directory).

    In the <job_name> folder, a human-readable .json file will be generated for each archive containing the file
    metadata from that archive. Additionally, a sqlite .db file will also be generated for each device in the job to
    track restore progress per device. Because a device can have multiple (destination) archives and/or multiple backup
    sets within a single archive, it's possible that some file paths will be duplicated within the full set of a single
    device's metadata. While duplicates will remain in the individual archive metadata files, they will be de-duplicated
    when restoring and when calculating the total file count and byte count of the restore job (the fileId with the most
    recent modification timestamp will be restored).
    """
    target = Path(target) if target else Path(".")
    params = {
        "includeDeleted": include_deleted,
        "backupSetIds": backup_set_ids,
        "deletedMaxAge": max_deleted_date,
        "restoreAsOf": as_of_date,
        "includePatterns": include_pattern,
        "excludePatterns": exclude_pattern,
    }
    client = Client(log_level=log_level, log_file=log_file)
    ids = id.split(",")
    archives = []
    for id in ids:
        archives.extend(client.get_archive_details(id, id_type=type_))
    job = RestoreJob(
        type=type_,
        identifier=id,
        selection=FileSelection(**params),
        archives=archives,
        targetFolder=target.absolute(),
    )
    if name:
        job.jobId = name
    if no_calculate:
        with open(f"restore_{job.jobId}.json", "w") as f:
            f.write(job.json(indent=2))
    else:
        _prepare_job(job, client)


def _prepare_job(job: RestoreJob, client: Client):
    console.print("Creating job folders...")
    job.folder.mkdir(parents=True, exist_ok=True)

    selection_params = {
        "endTime": job.selection.endTimeMillis,
        "includeDeleted": job.selection.includeDeleted,
    }
    msg = f"Fetching archive metadata with the following selection criteria: {job.selection}"
    console.print(msg)
    client.settings.logger.info(msg)
    for archive in job.archives:
        console.print(
            f"  user={archive.username}, device_name={archive.deviceName}, device_guid={archive.deviceId}, archive={archive.archiveId}"
        )
        manifest = client.get_archive_metadata(archive, **selection_params).json()
        with open(job.folder / archive.json_filename, "w") as f:
            json.dump(manifest, f, indent=2)
        files = manifest["data"]["archiveMetadataEntries"]
        save_archive_metadata_sqlite(archive, files, job.folder, job.selection)

    console.print("Calculating restore size...")
    job.calculate_stats()
    msg = f"Job size is {job.totalFileCount} files totaling {filesize.decimal(job.totalByteCount)}"
    console.print(msg)
    client.settings.logger.info(msg)
    with open(job.filename, "w") as f:
        f.write(job.json(indent=2))
        msg = f"Job definition saved: {job.filename}"
        console.print(msg)
        client.settings.logger.info(msg)


@brt.command()
@click.argument("JOB_FILE", type=click.File())
def start_job(job_file):
    """Begins restoring from archives defined in a given job file."""
    client = Client()
    job = RestoreJob.parse_raw(job_file.read())
    if not job.folder.exists():
        _prepare_job(job, client)
    else:
        refresh = False
        for archive in job.archives:
            if not (job.folder / archive.json_filename).exists():
                msg = f"Metadata files missing for archiveId {archive.archiveId}..."
                console.print(msg)
                client.settings.logger.warning(msg)
                refresh = True
        if refresh:
            msg = "Re-fetching job metadata..."
            console.print(msg)
            client.settings.logger.warning(msg)
            _prepare_job(job, client)
    for archive in job.archives:
        if archive.sessionId is None:
            archive.sessionId = client.create_bulk_restore_session(
                archive, name=job.jobId
            )
            client.settings.logger.info(
                f"Created bulk restore sessionId={archive.sessionId} for archiveId={archive.archiveId} on server={archive.storageNodeWebUrl}"
            )
    with open(job.filename, "w") as f:
        f.write(job.json(indent=2))
    asyncio.run(restore_main(job, client))


SESSIONS_PER_ARCHIVE = 1


async def restore_main(job: RestoreJob, client: Client):
    with ProgressBar(job) as progress_bar:
        tasks = []
        for archive in job.archives:
            db = job.devices[archive.deviceId].db
            file_queue = restore_queue(db, archive.archiveId)
            for _ in range(SESSIONS_PER_ARCHIVE):
                task = restore_archive(
                    client=client,
                    archive=archive,
                    job=job,
                    file_queue=file_queue,
                    progress_bar=progress_bar,
                )
                tasks.append(task)
        for completed in asyncio.as_completed(tasks):
            archive = await completed
            try:
                client.complete_bulk_restore_session(archive, archive.sessionId)  # type: ignore
            except HTTPStatusError:
                ...
            msg = f"Completed all queued files for archive {archive.archiveId} of device {archive.deviceId}."
            client.settings.logger.info(msg)


async def restore_archive(
    client: Client,
    archive: Archive,
    job: RestoreJob,
    file_queue: Queue,
    progress_bar: ProgressBar,
):
    device = job.devices[archive.deviceId]

    # A device/error logger for each device writes to its own file for easy audit of the restore record for each device.
    # The error logger will only log files that failed all 3 restore attempts so an end-user can quickly identify
    # which files are not in the target folder that _should_ be.
    # The device loggers also write to the main log file (if the --log-level option is INFO or higher)
    device_logger, error_logger = client.settings.get_device_loggers(device, job.folder)
    main_logger = client.settings.logger

    file_table = device.db["file"]
    storage_client = httpx.AsyncClient(
        auth=client.config["auth"], base_url=archive.storageNodeWebUrl, timeout=600  # type: ignore
    )
    while not file_queue.empty():
        file: RestorableFile = file_queue.get()
        progress_bar.start_file_progress(file)
        try:
            await restore_file(storage_client, archive, file, job.folder)
            file_table.update(
                pk_values=(file.archiveId, file.fileId), updates={"completed": True}
            )
            file.completed = True
            device_logger.info(
                f"Restored archive={archive.archiveId} fileId={file.fileId} sourcePath={file.path} targetPath={job.folder / file.restore_path} bytes={file.sourceLengthBytes} ({file.size})"
            )

        except InvalidBulkRestoreSession:
            main_logger.warning(
                f"Invalid bulk restore sessionId={archive.sessionId}, creating new session."
            )
            archive.sessionId = client.create_bulk_restore_session(
                archive, name=job.jobId
            )
            main_logger.debug(f"Updating job file: {job.json()}")
            with open(job.filename, "w") as f:
                f.write(job.json(indent=2))
            file_queue.queue.appendleft(file)
            progress_bar.complete_file_progress(file)
            continue

        except Exception as err:
            # track failed attempts, log each error and retry up to `max_attempts` times
            max_attempts = 3
            file.failure_count += 1
            device_logger.warning(
                f"Failed attempt {file.failure_count}/{max_attempts} to restore fileId={file.fileId} path={file.path}. Error: {err}"
            )
            client.settings.logger.debug(err, exc_info=True)

            if file.failure_count < max_attempts:
                file_queue.queue.appendleft(file)
            else:
                error_logger.warning(
                    f"{max_attempts} restore attempts failed for {file.path} (fileId={file.fileId}) from archiveId={archive.archiveId} on server={archive.storageNodeWebUrl}"
                )

        progress_bar.complete_file_progress(file)
    return archive


async def restore_file(
    client: httpx.AsyncClient, archive: Archive, file: RestorableFile, target: Path
):
    file_path = target / file.restore_path
    file_path.parent.mkdir(parents=True, exist_ok=True)
    if file.sourceLengthBytes == 0:
        # don't make network requests for 0byte files
        file_path.touch(exist_ok=True)
        return
    params = {
        "bulkRestoreSessionId": archive.sessionId,
        "archiveGuid": archive.archiveId,
        "fileId": file.fileId,
        "versionTimestamp": file.versionTimeMillis,
    }
    with open(file_path, "wb") as target_file:
        async with client.stream(
            method="GET", url="/api/v3/bulk-restore-file-data", params=params
        ) as response:
            if (
                response.status_code == 404
                and "Bulk Restore Session ID not found" in response.text
            ):
                raise InvalidBulkRestoreSession()
            response.raise_for_status()
            async for chunk in response.aiter_bytes(chunk_size=1024):
                target_file.write(chunk)
    # apply last modified/access metadata
    sourceLastModifiedNanoSeconds = file.sourceLastModifiedMillis * 1_000_000
    utime(file_path, ns=(sourceLastModifiedNanoSeconds, sourceLastModifiedNanoSeconds))

    checksum = sha256(file_path.read_bytes()).hexdigest()
    if checksum != file.sourceSHA256Checksum:
        raise Exception(
            f"Invalid SHA256 checksum for fileId={file.fileId}, expected {file.sourceSHA256Checksum}, got {checksum}."
        )


class InvalidBulkRestoreSession(Exception):
    ...
