import asyncio
import logging
import sqlite3
import time
from asyncio import CancelledError
from pathlib import Path
from queue import Queue
from typing import Union, Optional

import httpx
import ijson
from httpx import HTTPStatusError
from rich import filesize

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


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.
        """
        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
        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 = {
            archive.deviceId: Device(
                username=archive.username,
                name=archive.deviceName,
                deviceId=archive.deviceId,
                db_file=self._definition.db_folder / archive.db_filename,
            )
            for archive in self._definition.archives
        }
        self._quiet = quiet
        self._progress_bar = None
        self._log_util = LogUtil(self._client.settings.logger, quiet)
        self._sessions_per_archive = self.settings.downloads_per_archive

    @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.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,
        }
        if self._definition.selection.endTimeMillis:
            params["endTimeMillis"] = self._definition.selection.endTimeMillis
        url = f"{archive.storageNodeWebUrl}/api/v3/archive-metadata/{archive_id}"
        with open(
                self._definition.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=1024):
                    metadata_file.write(chunk)
        self._log_util.print_and_log(
            f"Done fetching metadata for archive {archive_id}: {self._definition.folder / archive.json_filename}.",
            highlight=False,
        )
        return archive

    async def prepare_async(self):
        """
        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

        # 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)
                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)}"
        )

    @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
        )

        with open(str(target_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(log_util_, 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
                )

        return len(files)

    @staticmethod
    def _add_entry(log_util_: LogUtil, 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

    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._clear_in_progress_statuses()
        with self.progress_bar:
            tasks = []
            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)
            )
            archive_active_session_count = {}
            for archive in self._definition.archives:
                archive_active_session_count[archive.archiveId] = 0
                db = self.devices[archive.deviceId].db

                ensure_tables_sqlite(archive, self._definition.db_folder)

                with self._log_util.timer(f"Built file list to restore archive {archive.archiveId}"):
                    file_queue = restore_queue(db, archive.archiveId)
                for _ in range(self._sessions_per_archive):
                    task = self._restore_archive(
                        archive=archive,
                        file_queue=file_queue,
                    )
                    tasks.append(task)
                    archive_active_session_count[archive.archiveId] += 1  # New session to track
            progress_task = asyncio.create_task(self._progress_logger())

            for completed in asyncio.as_completed(tasks):
                archive = await completed
                archive_active_session_count[archive.archiveId] -= 1  # Session has finished. stop tracking
                if archive_active_session_count[archive.archiveId] == 0:  # All sessions complete, close bulk_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)
            await self._writer.job_completed()
            progress_task.cancel()
            self._log_util.print_and_log(f"Job {self._definition.jobId} completed!")

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

        # 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

        if not has_file_versions(db, archive.archiveId):
            self._log_util.print_and_log(f"Archive has no file versions, skipping. archiveId: {archive.archiveIdId}")
            return archive

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

            # This is where we find the next file to work on.
            # Instead of grabbing a single file, grab many files until we hit a batch size (by file count or bytes)
            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 archive={archive.archiveId} fileId={file.fileId} sourcePath={file.path} targetPath={self._definition.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 = 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_logger.error(
                        f"{max_attempts} restore attempts failed for {file.path} (fileId={file.fileId}) from archiveId={archive.archiveId} on server={archive.storageNodeWebUrl}"
                    )

            self.progress_bar.complete_file_progress(file)

        self.progress_bar.complete_archive_queue(archive.deviceId)
        return archive

    async def _restore_file(self, archive: Archive, file: RestorableFile):
        params = {
            "bulkRestoreSessionId": archive.sessionId,
            "archiveGuid": archive.archiveId,
            "fileId": file.fileId,
            "versionTimestamp": file.versionTimeMillis,
        }
        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 = self.devices[archive.deviceId]
                    device_logger, _ = self.settings.get_device_loggers(device, self._definition.folder)
                    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()
                async for chunk in response.aiter_bytes(chunk_size=1024):
                    await self._writer.write_chunk(file, chunk, self._definition.forceWinPathSanitization)  # type: ignore

        # 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

    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_restore_sessions(self):
        for archive in self._definition.archives:
            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}"
                )
        with open(self._definition.filename, "w") as f:
            f.write(self._definition.json(indent=2))

    def prepare(self):
        """
        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())

    async def start_async(self):
        """
        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()
        try:
            self._calculate_stats()
        except sqlite3.OperationalError:
            await self.prepare_async()
            self._calculate_stats()
        self._create_restore_sessions()
        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())
