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

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

from bulk_restore_tool.models import JobDefinition
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
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:
    _SESSIONS_PER_ARCHIVE = 1

    def __init__(
        self,
        definition: JobDefinition,
        url: Optional[str] = None,
        username: Optional[str] = None,
        password: Optional[str] = None,
        api_client_id: Optional[str] = None,
        api_client_secret: Optional[str] = None,
        log_file: Union[str, Path] = "bulk_restore.log",
        log_level: str = "INFO",
        writer: Optional[RestoreWriter] = None,
        quiet: bool = False,
    ):
        self._client = Client(
            url=url,
            username=username,
            password=password,
            api_client_id=api_client_id,
            api_client_secret=api_client_secret,
            log_file=log_file,
            log_level=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)
        self._definition = definition
        self._ensure_archive_details()
        self._writer = writer
        self.totalFileCount = 0
        self.totalByteCount = 0
        self._loop = asyncio.new_event_loop()
        self._prepared = False
        self.devices = {
            archive.deviceId: Device(
                username=archive.username,
                name=archive.deviceName,
                deviceId=archive.deviceId,
                db_file=self._definition.folder / archive.db_filename,
            )
            for archive in self._definition.archives
        }
        self._quiet = quiet
        self._progress_bar = None
        self._console = Console(file=sys.stderr)

    def print_and_log(self, msg: str, **kwargs):
        """Prints msg to console (if `quiet` param is False) and logs msg to main log file."""
        if not self._quiet:
            self._console.print(msg, **kwargs)
        self.settings.logger.info(msg)

    @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."""
        archiveId = 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/{archiveId}"
        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.print_and_log(
            f"Done fetching metadata for archive {archiveId}: {self._definition.folder / archive.json_filename}.",
            highlight=False,
        )
        return archive

    async def _prepare_job(self):
        """Fetches metadata for all archives in the job definition and builds sqlite databases for each device."""
        msg = f"Preparing job metadata with the following selection criteria: {self._definition.selection}"
        self.print_and_log(msg)
        archives = self._definition.archives
        job_folder = self._definition.folder

        # download metadata async
        self.print_and_log(
            f"Fetching metadata for archives: {[a.archiveId for a in archives]}."
        )
        for task in asyncio.as_completed(
            [self._get_archive_metadata(a) for a in archives]
        ):
            await task

        # format json files and create dbs
        for archive in archives:
            with open(job_folder / archive.json_filename, "r+") as file:
                metadata = json.load(file)
                file.seek(0)
                file.truncate()
                json.dump(metadata, file, indent=2)
            files = metadata["data"]["archiveMetadataEntries"]
            self.print_and_log(f"Saving metadata to db for {archive.archiveId}.")
            save_archive_metadata_sqlite(
                archive=archive,
                files=files,
                target=job_folder,
                selection=self._definition.selection,
            )
        self.print_and_log("Calculating job size...")
        self._calculate_stats()
        self.print_and_log(
            f"Job size is {self.totalFileCount} files totaling {filesize.decimal(self.totalByteCount)}"
        )

    def _calculate_stats(self):
        self.totalFileCount = self.totalByteCount = 0
        for device in self.devices.values():
            self.totalFileCount += device.totalFiles
            self.totalByteCount += device.totalBytes

    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)
            )
            for archive in self._definition.archives:
                db = self.devices[archive.deviceId].db
                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)
            progress_task = asyncio.create_task(self._progress_logger())
            for completed in asyncio.as_completed(tasks):
                archive = await completed
                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.print_and_log(msg)
            self._writer.job_completed()
            progress_task.cancel()
            self.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

        while not file_queue.empty():
            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 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.warning(
                        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):
        if file.sourceLengthBytes == 0:
            # don't make network requests for 0byte files
            await self._writer.write_chunk(file, b"")  # type: ignore
            await self._writer.file_completed(file)  # type: ignore
            return
        params = {
            "bulkRestoreSessionId": archive.sessionId,
            "archiveGuid": archive.archiveId,
            "fileId": file.fileId,
            "versionTimestamp": file.versionTimeMillis,
        }
        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 == 404:
                await response.aread()
                if "Session ID not found" in response.text:
                    raise InvalidBulkRestoreSession()
            response.raise_for_status()
            async for chunk in response.aiter_bytes(chunk_size=1024):
                await self._writer.write_chunk(file, chunk)  # type: ignore
            await self._writer.file_completed(file)  # 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):
        """
        Downloads 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_job())

    def start(self):
        """Starts the restore job. If metadata has not yet been prepared, will automatically do so."""
        if not self.prepared:
            self.prepare()
        try:
            self._calculate_stats()
        except sqlite3.OperationalError:
            self.prepare()
            self._calculate_stats()
        self._create_restore_sessions()
        self._loop.run_until_complete(self._restore_main())
