import asyncio
import copy
import logging
import sqlite3
import time
from asyncio import CancelledError
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from queue import Queue, Empty
from typing import Union, Optional

import httpx
import ijson
from httpx import HTTPStatusError
from rich import filesize
from sqlite_utils import Database

from bulk_restore_tool._client import Client
from bulk_restore_tool._db import (
    add_indices,
    restore_directory_queue,
    load_directories_into_queue,
    set_dir_status,
)
from bulk_restore_tool._db import get_file_status
from bulk_restore_tool._db import restore_queue
from bulk_restore_tool._db import (
    save_archive_metadata_sqlite,
    load_restorable_files_into_queue,
    ensure_tables_sqlite,
    has_file_versions,
    close_database,
)
from bulk_restore_tool._db import set_file_status
from bulk_restore_tool._log_util import LogUtil
from bulk_restore_tool._progress import ProgressBar
from bulk_restore_tool._writers import FileWriter
from bulk_restore_tool._writers import ProcessZipWriter
from bulk_restore_tool._writers import RestoreWriter
from bulk_restore_tool.models import Archive
from bulk_restore_tool.models import Device
from bulk_restore_tool.models import FileStatus
from bulk_restore_tool.models import JobDefinition
from bulk_restore_tool.models import RestorableFile
from bulk_restore_tool.models import metadata_map


class InvalidBulkRestoreSession(Exception): ...


class RestoreJob:
    def __init__(
        self,
        definition: JobDefinition,
        url: Optional[str] = None,
        api_client_id: Optional[str] = None,
        api_client_secret: Optional[str] = None,
        logger: Optional[logging.Logger] = None,
        log_file: Union[str, Path] = None,
        log_level: Union[int, str] = None,
        device_log_level: Union[int, str] = None,
        writer: Optional[RestoreWriter] = None,
        quiet: bool = False,
    ):
        """
        :param definition: A `JobDefinition` object specifying what is to be restored.
        :param url: Optional url of your Code42 console (if not provided in constructor, must be present in BRT_URL
            variable either in your shell or .env file in current working directory).
        :param api_client_id: Optional identifier of your Code42 API client for authentication (if not provided in
            constructor, must be present in BRT_API_CLIENT_ID variable either in your shell or .env file in current
            working directory).
        :param api_client_secret: Optional identifier of your Code42 API client secret for authentication (if not
            provided in constructor, must be present in BRT_API_CLIENT_ID variable either in your shell or .env file in
            current working directory).
        :param logger: Optional custom logger to use instead of default file logger. If provided
        :param log_file: Optional path to write BRT primary logs to. Defaults to 'bulk_restore.log' in CWD if not
            provided.
        :param log_level: Set logging level of primary BRT logs. Defaults to INFO.
        :param device_log_level: Set level for logging restore progress for each device.
            - INFO logs each file restored to main log file and writes a log file per device in the target directory.
            - WARN will skip logging every successful restore, but will log each failed attempt (the BRT will attempt each file up to 3 times).
            - ERROR will only log files that failed to restore to the main log file and the <device>_error.log in the target directory.
            - CRITICAL will cause no logging to happen for individual file restores, and no log files will be written to target folder.
        :param writer: Optional implementation of the `RestoreWriter` interface, which will be used to write restored
            file contents to target.
        :param quiet: Flag to disable writing of restore progress bar to stderr. Defaults to False.
        """
        if log_file is None:
            log_file = definition.log_file
        if log_level is None:
            log_level = definition.log_level
        if device_log_level is None:
            device_log_level = definition.device_log_level

        self._client = Client(
            url=url,
            api_client_id=api_client_id,
            api_client_secret=api_client_secret,
            logger=logger,
            log_file=log_file,
            log_level=log_level,
            device_log_level=device_log_level,
        )

        self._aclient = httpx.AsyncClient(auth=self._client.config["auth"], timeout=60)  # type: ignore
        self.settings = self._client.settings
        self._quiet = quiet
        self._log_util = LogUtil(self._client.settings.logger, quiet)
        definition.folder.mkdir(parents=True, exist_ok=True)
        definition.db_folder.mkdir(parents=True, exist_ok=True)
        self._definition = definition
        self._ensure_archive_details()
        self._writer = writer
        self.totalFileCount = 0
        self.totalByteCount = 0
        self._loop = asyncio.get_event_loop()
        self._prepared = False

        self.devices = {}
        for archive in self._definition.archives:
            if archive.deviceId not in self.devices:
                db_file = self._definition.db_folder / archive.db_filename
                self._log_util.log(
                    "Adding device to track: "
                    + archive.deviceId
                    + " archive: "
                    + archive.archiveId
                    + " db_file: "
                    + str(db_file),
                    logging.DEBUG,
                )

                if db_file.is_file():
                    self._log_util.log(
                        "DB for device already exists. device=" + archive.deviceId
                    )

                self.devices[archive.deviceId] = Device(
                    username=archive.username,
                    name=archive.deviceName,
                    deviceId=archive.deviceId,
                    db_file=db_file,
                )
            else:
                self._log_util.log(
                    "Duplicate device found, already being tracked: "
                    + archive.deviceId
                    + " archive: "
                    + archive.archiveId,
                    logging.DEBUG,
                )

        self._progress_bar = None
        self._concurrent_archive_download_limit = (
            self.settings.concurrent_archive_download_limit
        )
        self._sessions_per_archive = self.settings.downloads_per_archive
        self._restore_stats = RestoreStatus()
        self._log_util.log("Completed building RestoreJob.", logging.DEBUG)

    def get_job_status(self):
        """Returns a copy of the current state of the Job. This is a snapshot in time."""
        return copy.deepcopy(self._restore_stats)

    @property
    def prepared(self):
        """If archive metadata files exist for all archives, we assume the job is ready to run."""
        for archive in self._definition.archives:
            if not (self._definition.db_folder / archive.json_filename).exists():
                return False
        return True

    @property
    def progress_bar(self):
        if not self.prepared:
            raise Exception(
                "Progress bar can't be initialized without a prepared job. Run RestoreJob.prepare() first."
            )
        if self._progress_bar is None:
            self._progress_bar = ProgressBar(job=self)
        return self._progress_bar

    def _ensure_archive_details(self):
        """Make sure archive objects have all necessary data."""
        updated = []
        for archive in self._definition.archives:
            if not all((archive.deviceId, archive.storageNodeWebUrl, archive.username)):
                updated.extend(
                    self._client.get_archive_details(
                        identifier=archive.archiveId, id_type="archive"
                    )
                )
            else:
                updated.append(archive)
        self._definition.archives = updated

    async def _get_archive_metadata(self, archive: Archive):
        """Downloads the metadata for a given archive and writes to a json file."""
        archive_id = archive.archiveId
        params = {
            "includeDeleted": self._definition.selection.includeDeleted
            or self._definition.selection.includeOnlyDeleted,
        }
        if self._definition.selection.endTimeMillis:
            params["endTimeMillis"] = self._definition.selection.endTimeMillis
        url = f"{archive.storageNodeWebUrl}/api/v3/archive-metadata/{archive_id}"
        self._log_util.log(
            f"Creating archive metadata json file. "
            f"path={str(self._definition.db_folder / archive.json_filename)}"
        )
        metadata_file_size_bytes = 0
        with open(
            self._definition.db_folder / archive.json_filename, "wb"
        ) as metadata_file:
            async with self._aclient.stream("GET", url, params=params) as response:
                response.raise_for_status()
                async for chunk in response.aiter_bytes(chunk_size=1000000):
                    metadata_file_size_bytes += metadata_file.write(chunk)
        self._log_util.print_and_log(
            f"Done fetching metadata for archive {archive_id}: {self._definition.folder / archive.json_filename}. size={metadata_file_size_bytes}",
            highlight=False,
        )
        return archive

    async def prepare_async(self, close_db_when_done: bool = False):
        """
        Async entrypoint for downloading metadata for all archives in the job.

        Creates a sqlite db for each device and applies marks files as SELECTED based on job definition's
        selection criteria.
        """
        msg = f"Preparing job metadata with the following selection criteria: {self._definition.selection}"
        self._log_util.print_and_log(msg)
        archives = self._definition.archives
        target_folder = self._definition.folder
        database_folder = self._definition.db_folder
        self._restore_stats.status = JobState.METADATA_DOWNLOAD

        # download metadata async
        self._log_util.print_and_log(
            f"Fetching metadata for archives: {[a.archiveId for a in archives]}."
        )
        with self._log_util.timer("Metadata for all archives retrieved", logging.INFO):
            for task in asyncio.as_completed(
                [self._get_archive_metadata(a) for a in archives]
            ):
                await task

        with self._log_util.timer(
            "Metadata loaded and inserted into DB for all archives.", logging.INFO
        ):
            for archive in archives:
                try:
                    count = RestoreJob.load_and_insert_into_database_archive_metadata(
                        self._log_util,
                        target_folder,
                        database_folder,
                        archive,
                        self._definition.selection,
                    )
                    self._log_util.print_and_log(
                        f"Loaded {count} files for archive {archive.archiveId}",
                        logging.DEBUG,
                    )
                    self._get_archive_status_tracker(
                        archive.archiveId
                    ).archive_file_count = count
                except Exception as exc:
                    self._log_util.print_and_log(
                        f"Problem loading metadata for archive {archive.archiveId}",
                        logging.ERROR,
                    )
                    raise exc

        self._add_database_indices(self._log_util)

        self._log_util.print_and_log("Calculating job size...")
        self._calculate_stats()
        self._log_util.print_and_log(
            f"Job size is {self.totalFileCount} files totaling {filesize.decimal(self.totalByteCount)}"
        )
        if close_db_when_done:
            self.close_all_database_connections()
        self._restore_stats.status = JobState.METADATA_DOWNLOAD_COMPLETE

    @staticmethod
    def load_and_insert_into_database_archive_metadata(
        log_util_: LogUtil,
        target_folder,
        db_folder,
        archive,
        selection,
        batch_size=10000,
    ):
        count = 0

        # In the case where no file versions were returned, we still need to
        # ensure the database has tables.
        ensure_tables_sqlite(archive=archive, db_target=db_folder, log_util_=log_util_)

        with open(
            str(db_folder) + "/" + str(archive.json_filename), "rb"
        ) as input_file:
            with log_util_.timer(
                f"Parsing and loading file versions for {archive.archiveId}",
                logging.INFO,
            ):
                files = []
                parser = ijson.parse(input_file)

                prefix = ""
                state = "seek_array"
                for json_path, event, value in parser:
                    # State machine !
                    #####################################################################
                    if state == "seek_array":
                        if event == "start_array":
                            state = "seek_entry"
                        else:
                            pass

                    #####################################################################
                    elif state == "seek_entry":
                        if event == "start_map":
                            state = "build_entry"
                            entry = {}
                            entry_key = ""
                            deleted_seen = False

                        elif event == "end_array":
                            state = "exit"

                    #####################################################################
                    elif state == "build_entry":
                        if event == "start_map":
                            if (
                                json_path
                                == "data.archiveMetadataEntries.item.lastNonDeletedEntry"
                            ):
                                # previous non-deleted file versions held in a sub-array.
                                # give those fields a special "previous_" name, so they can
                                # be extracted during Pydantic parsing.
                                prefix = "previous_"
                                deleted_seen = True
                                entry_key = ""

                        elif event == "end_map":
                            if prefix != "":
                                # we have reached the end of a lastNonDeletedEntry structure.
                                prefix = ""

                            else:
                                RestoreJob._add_entry(files, entry, deleted_seen)
                                deleted_seen = False
                                if len(files) >= batch_size:
                                    count += RestoreJob._write_entries_to_database(
                                        log_util_, files, archive, db_folder, selection
                                    )
                                    files = []  # clear the array

                                state = "seek_entry"

                        elif event == "map_key":
                            entry_key = value

                        else:
                            if event == "string":
                                entry_value = value

                            elif event == "number":
                                entry_value = int(value)

                            elif event == "boolean":
                                entry_value = value

                            elif event == "null":
                                entry_value = None

                            entry[prefix + entry_key] = entry_value

                count += RestoreJob._write_entries_to_database(
                    log_util_, files, archive, db_folder, selection
                )

            return count

    @staticmethod
    def _write_entries_to_database(
        log_util_: LogUtil, files, archive, job_folder, selection
    ):
        if len(files) > 0:
            with log_util_.timer(
                f"Saving metadata batch to db: {archive.archiveId}", logging.DEBUG
            ):
                save_archive_metadata_sqlite(
                    archive=archive,
                    files=files,
                    db_target=job_folder,
                    selection=selection,
                    log_util_=log_util_,
                )

        return len(files)

    @staticmethod
    def _add_entry(files, entry: dict, deleted_seen):
        if deleted_seen:
            deleted_block = {}
            keys = {
                "previous_versionTimeMillis": "versionTimeMillis",
                "previous_sourceLastModifiedMillis": "sourceLastModifiedMillis",
                "previous_sourceLengthBytes": "sourceLengthBytes",
                "previous_sourceMD5Checksum": "sourceMD5Checksum",
                "previous_sourceSHA256Checksum": "sourceSHA256Checksum",
            }
            for k in keys:
                if entry.get(k) is not None:
                    deleted_block[keys[k]] = entry[k]
                    entry.pop(k)

            entry["lastNonDeletedEntry"] = deleted_block

        files.append(entry)

    def _add_database_indices(self, log_util_: LogUtil):
        for device in self.devices.values():
            with log_util_.timer(
                f"Adding indices for device {device.deviceId}", logging.INFO
            ):
                add_indices(device.db)

    def _calculate_stats(self):
        self.totalFileCount = self.totalByteCount = 0
        for device in self.devices.values():
            self.totalFileCount += (
                device.totalFiles if device.totalFiles is not None else 0
            )
            self.totalByteCount += (
                device.totalBytes if device.totalBytes is not None else 0
            )

        self._restore_stats.expected_total_job_file_count = self.totalFileCount
        self._restore_stats.expected_total_job_file_size_bytes = self.totalByteCount

    def _clear_in_progress_statuses(self):
        # reset any IN_PROGRESS file statuses to SELECTED (if tool was previously cancelled mid-download)
        for device in self.devices.values():
            with device.db.conn:
                device.db.execute(
                    "UPDATE file set status = 'SELECTED' where status = 'IN_PROGRESS'"
                )

    async def _restore_main(self):
        self._restore_stats.status = JobState.RESTORE_IN_PROCESS
        self._restore_stats.time_started = datetime.now()
        self._calculate_stats()

        self._clear_in_progress_statuses()
        with self.progress_bar:
            self._writer: RestoreWriter = self._writer or (
                ProcessZipWriter(
                    target=self._definition.folder, progress=self.progress_bar.extra
                )
                if self._definition.zipResults
                else FileWriter(target=self._definition.folder)
            )

            # Ensure all the archives for a given device are processed together
            archive_queue = Queue()
            archives = self._definition.archives.copy()
            archives.sort(key=lambda a: a.deviceId, reverse=True)
            for archive in archives:
                archive_queue.put(archive)

            tasks = []
            for _ in range(self._concurrent_archive_download_limit):
                t = asyncio.create_task(self._restore_all_archives(archive_queue))
                tasks.append(t)

            progress_task = asyncio.create_task(self._progress_logger())
            await asyncio.gather(*tasks)

            await self._writer.job_completed()
            progress_task.cancel()
            self.close_all_database_connections()

            self._log_util.print_and_log(f"Job {self._definition.jobId} completed!")

        self._restore_stats.status = JobState.RESTORE_COMPLETE
        self._restore_stats.calculate_elapsed_time()
        self._restore_stats.calculate_restore_stats()

    def close_all_database_connections(self):
        for device in self.devices.values():
            close_database(device.db)

    def validate_archive_storage_node(self, archive):
        details = self._client.get_archive_details(archive.archiveId, "archive")
        storage_node_url = details[0].storageNodeWebUrl

        if storage_node_url != archive.storageNodeWebUrl:
            self._log_util.log(
                f"Storage node address in job file does not match actual storage node location. "
                f"archive was balanced. redirecting to new archive location. "
                f"{archive.archiveId=} old_url={archive.storageNodeWebUrl} new_url={storage_node_url}",
                logging.WARN,
            )
            archive.storageNodeWebUrl = details[0].storageNodeWebUrl

    async def _restore_all_archives(self, archive_queue: Queue):
        try:
            while True:
                archive = archive_queue.get(block=False)
                archive_status = self._get_archive_status_tracker(archive.archiveId)
                archive_status.archive_status = ArchiveRestoreState.ARCHIVE_STARTED
                self.validate_archive_storage_node(archive)
                self._create_restore_session(archive)

                db = self.devices[archive.deviceId].db
                ensure_tables_sqlite(
                    archive, self._definition.db_folder, self._log_util
                )

                with self._log_util.timer(
                    f"Built file list to restore archive {archive.archiveId}"
                ):
                    file_queue = restore_queue(db, archive.archiveId)

                # Kick of N "worker threads" to restore all the files.
                # Each "thread" will restore a different file in parallel
                archive_restore_sessions = []
                for _ in range(self._sessions_per_archive):
                    worker = self._restore_archive(
                        archive=archive,
                        file_queue=file_queue,
                    )
                    archive_restore_sessions.append(worker)

                # Wait for all the "worker threads" to complete for this archive.
                await asyncio.gather(*archive_restore_sessions)
                archive_status.archive_status = ArchiveRestoreState.ARCHIVE_DONE

                # Restore the empty directories if requested by the user.
                if self._definition.restore_all_directories:
                    await self._restore_folders(archive, db)

                # only complete the archive progress, once all tasks for that archive have completed.
                self.progress_bar.complete_archive_queue(
                    archive.deviceId, archive.archiveId
                )

                # close the restore session
                try:
                    self._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}."
                self._log_util.print_and_log(msg)

        except Empty:
            pass  # all archives have been processed. Allow this task to exit.

    async def _restore_archive(self, archive: Archive, file_queue: Queue):
        device = self.devices[archive.deviceId]
        db = device.db
        archive_status = self._get_archive_status_tracker(archive.archiveId)

        # 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 = self.settings.get_device_loggers(
            device, self._definition.folder
        )
        main_logger = self.settings.logger

        # FYI. this also prints if all the files to be restored were grabbed by a parallel task.
        if not has_file_versions(db, archive.archiveId):
            self._log_util.print_and_log(
                f"Archive has no file versions, skipping. archiveId: {archive.archiveId}"
            )
            return archive

        # while there is work left to do.  or we have loaded more work -> work on batches of files
        while not file_queue.empty() or load_restorable_files_into_queue(
            db, archive.archiveId, file_queue
        ):
            file: RestorableFile = file_queue.get()

            if file.status != FileStatus.RETRY:
                # skip if file is IN_PROGRESS or COMPLETED by other archive queue
                if get_file_status(file, db) != FileStatus.SELECTED:
                    continue
                else:
                    set_file_status(file, FileStatus.IN_PROGRESS, db)
            self.progress_bar.start_file_progress(file)
            try:
                await self._restore_file(archive, file)
                set_file_status(file, FileStatus.COMPLETED, db)
                file.status = FileStatus.COMPLETED
                device_logger.info(
                    f"Restored file. archive={archive.archiveId} fileId={file.fileId} sourcePath={file.path} targetPath={self._definition.folder / file.restore_path} bytes={file.sourceLengthBytes} ({file.size})"
                )
                archive_status.files_restored += 1
                archive_status.bytes_restored += file.sourceLengthBytes

            except InvalidBulkRestoreSession:
                main_logger.warning(
                    f"Invalid bulk restore sessionId={archive.sessionId}, creating new session."
                )
                archive.sessionId = self._client.create_bulk_restore_session(
                    archive, name=self._definition.jobId
                )
                main_logger.debug(f"Updating job file: {self._definition.json()}")
                with open(self._definition.filename, "w") as f:
                    f.write(self._definition.json(indent=2))
                file.status = FileStatus.RETRY
                file_queue.queue.appendleft(file)
                self.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}"
                )
                self.settings.logger.debug(err, exc_info=True)

                if file.failure_count < max_attempts:
                    file.status = FileStatus.RETRY
                    file_queue.queue.appendleft(file)
                else:
                    error_message = (
                        f"{max_attempts} restore attempts failed for {file.path} (fileId={file.fileId}) "
                        f"from archiveId={archive.archiveId} on server={archive.storageNodeWebUrl} "
                        f"[restorePath={file.get_restore_path(self._definition.forceWinPathSanitization)} "
                        f"sizeBytes={file.sourceLengthBytes} "
                        f"fileType={file.fileType}]"
                    )
                    error_logger.error(error_message)
                    self._log_util.log(error_message, level=logging.WARN)
                    archive_status.files_failed += 1

            self.progress_bar.complete_file_progress(file)

        return archive

    async def _restore_file(self, archive: Archive, file: RestorableFile):
        params = {
            "bulkRestoreSessionId": archive.sessionId,
            "archiveGuid": archive.archiveId,
            "fileId": file.fileId,
            "versionTimestamp": file.versionTimeMillis,
        }
        device = self.devices[archive.deviceId]
        device_logger, _ = self.settings.get_device_loggers(
            device, self._definition.folder
        )

        self._log_util.log(
            f"Starting file restore: archive={archive.archiveId} fileId={file.fileId} size={file.sourceLengthBytes} sourcePath={file.path}"
        )

        streamed_bytes = 0
        if file.sourceLengthBytes == 0:
            # don't make network requests for 0 byte files
            await self._writer.write_chunk(file, b"", self._definition.forceWinPathSanitization)  # type: ignore
        else:
            # restore file
            url = f"{archive.storageNodeWebUrl}/api/v3/bulk-restore-file-data"
            async with self._aclient.stream("GET", url, params=params) as response:
                if response.status_code // 100 != 2:  # if not a 2XX code
                    await response.aread()
                    device_logger.warn(
                        f"A problem was encountered downloading a file. path={file.path} Response={response.status_code}"
                    )

                    if response.status_code == 404:
                        if "Session ID not found" in response.text:
                            raise InvalidBulkRestoreSession()

                    elif response.status_code == 401:
                        device_logger.debug(
                            f"401 Response. rsp={response.text}  headers={response.headers}"
                        )
                response.raise_for_status()

                byte_counter = 0
                async for chunk in response.aiter_bytes(chunk_size=1_000_000):
                    try:
                        write_result = await self._writer.write_chunk(file, chunk, self._definition.forceWinPathSanitization)  # type: ignore
                        if write_result is None:
                            # It is possible for the write operation to return None. However, why it returned None
                            # is unclear. https://docs.python.org/3/library/io.html#io.RawIOBase.write/
                            self._log_util.log(
                                f"Attempt to write to file returned 'None'. chunk_size={len(chunk)} bytes_completed={byte_counter} bytes_expected={file.sourceLengthBytes} path={file.path}",
                                logging.WARN,
                            )
                            raise RuntimeError(
                                f"Error writing data to file. path={file.path}"
                            )

                        streamed_bytes += write_result
                        if streamed_bytes >= byte_counter + 100_000_000:  # 100MB
                            byte_counter = streamed_bytes
                            self._log_util.log(
                                f"streaming data. bytes_completed={byte_counter} bytes_expected={file.sourceLengthBytes} path={file.path}"
                            )
                    except IsADirectoryError:
                        # cannot restore the file, because a directory already exists where we want to place the file.
                        # resolve this by renaming the current file.
                        before_path = file.get_restore_path(
                            self._definition.forceWinPathSanitization
                        )
                        file.set_directory_name_collision()
                        self._log_util.print_and_log(
                            f"Destination filename is already taken by a directory. Renaming the file. "
                            f"previousName={before_path} "
                            f"updatedName={file.get_restore_path(self._definition.forceWinPathSanitization)}",
                            level=logging.WARN,
                        )
                        streamed_bytes += await self._writer.write_chunk(
                            file, chunk, self._definition.forceWinPathSanitization
                        )  # type: ignore

        if file.sourceLengthBytes != streamed_bytes:
            device_logger.warn(
                f"Expected file length did not match actual number of bytes downloaded. "
                f"file_size={file.sourceLengthBytes} downloaded_size={streamed_bytes} fileId={file.fileId} path={file.path}"
            )

        # restore metadata
        metadata_url = f"{archive.storageNodeWebUrl}/api/v3/bulk-restore-file-metadata"
        metadata_response = await self._aclient.get(metadata_url, params=params)
        if metadata_response.status_code == 200:
            data = metadata_response.json()["data"]
            metadata = metadata_map[data["os"]](**data)
            await self._writer.write_metadata(file, metadata)
        else:
            device = self.devices[archive.deviceId]
            _, error_logger = self.settings.get_device_loggers(
                device, self._definition.folder
            )
            error_logger.error(
                f"Failed to fetch metadata for fileId={file.fileId} path={file.path}. Response: {metadata_response.text}"
            )
        await self._writer.file_completed(file, self._definition.forceWinPathSanitization)  # type: ignore

        self._log_util.log(
            f"Completed file restore: archive={archive.archiveId} size={file.sourceLengthBytes} sourcePath={file.path} restorePath={file.get_restore_path(self._definition.forceWinPathSanitization)}"
        )

    async def _restore_folders(self, archive: Archive, db: Database):
        self._log_util.print_and_log(
            f"Starting restore of empty folders for archive={archive.archiveId}"
        )

        dir_queue = restore_directory_queue(db, archive)
        while not dir_queue.empty() or load_directories_into_queue(
            db,
            archive.archiveId,
            archive.deviceId,
            archive.storageNodeWebUrl,
            dir_queue,
        ):
            directory: RestorableFile = dir_queue.get()

            if directory.path == directory.planId:
                self._log_util.print_and_log(f"Skipping planId path={directory.path}")
            else:
                directory_path_str = (
                    self._definition.folder / directory.get_restore_path()
                )
                directory_path = Path(directory_path_str)
                directory_path.mkdir(parents=True, exist_ok=True)
            set_dir_status(directory.path, FileStatus.COMPLETED, db)

        self._log_util.print_and_log(
            f"Completed restore of empty folders for archive={archive.archiveId}"
        )

    async def _progress_logger(self, seconds=30):
        """Log completed progress every n seconds."""
        last_time = time.time()
        try:
            while True:
                await asyncio.sleep(1)
                if time.time() - last_time > seconds:
                    prog_string = self.progress_bar.get_total_progress()
                    self.settings.logger.info(f"Restore Progress: {prog_string}")
                    last_time = time.time()
        except CancelledError:
            prog_string = self.progress_bar.get_total_progress()
            self.settings.logger.info(f"Restore Completed: {prog_string}")

    def _create_definition_file(self):
        with open(self._definition.filename, "w") as f:
            f.write(self._definition.json(indent=2))

    def _create_restore_session(self, archive):
        if archive.sessionId is None:
            archive.sessionId = self._client.create_bulk_restore_session(
                archive, name=self._definition.jobId
            )
            self.settings.logger.info(
                f"Created bulk restore sessionId={archive.sessionId} for archiveId={archive.archiveId} on server={archive.storageNodeWebUrl}"
            )

    def prepare(self, close_db_when_done: bool):
        """
        Sync entrypoint for downloading metadata for all archives in the job.

        Creates a sqlite db for each device and applies marks files as SELECTED based on job definition's
        selection criteria.
        """
        self._loop.run_until_complete(self.prepare_async(close_db_when_done))

    async def start_async(self, close_db_when_done: bool = False):
        """
        Async entrypoint for starting the restore job. If metadata has not yet been prepared, will automatically do so.
        """
        if not self.prepared:
            await self.prepare_async(close_db_when_done)
        try:
            self._calculate_stats()
        except sqlite3.OperationalError:
            await self.prepare_async(close_db_when_done)
            self._calculate_stats()
        self._create_definition_file()
        await self._restore_main()

    def start(self):
        """
        Sync entrypoint for starting the restore job. If metadata has not yet been prepared, will automatically do so.
        """
        self._loop.run_until_complete(self.start_async())

    def _get_archive_status_tracker(self, archive_id):
        archive_status = self._restore_stats.archive_status.get(archive_id)

        # if an archive status tracker does not yet exist, create one.
        if archive_status is None:
            archive = next(
                (x for x in self._definition.archives if x.archiveId == archive_id),
                None,
            )
            archive_status = ArchiveStatus()
            archive_status.archive_id = archive_id
            archive_status.device_id = archive.deviceId
            archive_status.device_name = archive.deviceName
            self._restore_stats.archive_status[archive_id] = archive_status

        return archive_status


class ArchiveRestoreState(Enum):
    ARCHIVE_NOT_STARTED = "not_started"
    ARCHIVE_STARTED = "archive_started"
    ARCHIVE_DONE = "archive_done"


@dataclass
class ArchiveStatus:
    archive_status: ArchiveRestoreState = ArchiveRestoreState.ARCHIVE_NOT_STARTED
    device_id: str = ""
    device_name: str = ""
    archive_id: int = -1
    bytes_restored: int = 0
    files_restored: int = 0
    files_failed: int = 0


class JobState(Enum):
    NOT_STARTED = "not_started"
    METADATA_DOWNLOAD = "metadata_download"
    METADATA_DOWNLOAD_COMPLETE = "metadata_download_complete"
    RESTORE_IN_PROCESS = "restore_in_process"
    RESTORE_COMPLETE = "restore_complete"


@dataclass
class RestoreStatus:
    status: JobState = JobState.NOT_STARTED
    archive_status: {int, ArchiveStatus} = field(default_factory=dict)
    expected_total_job_file_count: int = -1
    expected_total_job_file_size_bytes: int = -1
    total_files_restored: int = 0
    total_bytes_restored: int = 0
    total_files_failed: int = 0
    time_started: datetime = None
    time_ended: datetime = None
    time_delta: timedelta = None

    def calculate_restore_stats(self):
        for _, archive_status in self.archive_status.items():
            self.total_files_restored += archive_status.files_restored
            self.total_bytes_restored += archive_status.bytes_restored
            self.total_files_failed += archive_status.files_failed

    def calculate_elapsed_time(self):
        self.time_ended = datetime.now()
        self.time_delta = self.time_ended - self.time_started
