import io
import warnings
from abc import ABC
from abc import abstractmethod
from csv import DictWriter
from hashlib import sha256, md5
from multiprocessing import Process
from multiprocessing import SimpleQueue
from os import utime
from pathlib import Path
from shutil import rmtree
from traceback import format_exception
from typing import Dict, Tuple, Union, List
from zipfile import ZipFile
from zipfile import ZIP_LZMA
import logging

import aiofiles
from aiofiles.threadpool.binary import AsyncBufferedIOBase
from rich.progress import Progress

from bulk_restore_tool.models import (
    RestorableFile,
    LinuxMetadata,
    MacMetadata,
    WindowsMetadata,
)


class RestoreWriter(ABC):
    """
    Interface for enabling flexible restore targets.
    """

    @abstractmethod
    async def write_chunk(self, file: RestorableFile, chunk: bytes, forceWinPathSanitization: bool):
        """
        The restore job will asynchronously download multiple files, so the writer needs to be able to asynchronously
        handle/process individual chunks. The `file` param provides any necessary file metadata (path, checksum, size,
        etc.) needed for processing the chunk.
        """
        ...

    @abstractmethod
    async def write_metadata(
        self,
        file: RestorableFile,
        metadata: Union[LinuxMetadata, WindowsMetadata, MacMetadata],
    ):
        """
        Provides detailed file metadata that can either be applied directly to the restored files on disk, or saved
        elsewhere as a reference. This method is called after all `.write_chunk()` calls are made for the file contents.
        """
        ...

    @abstractmethod
    async def file_completed(self, file: RestorableFile, forceWinPathSanitization: bool):
        """
        Once the download is complete, the restore job will call this method indicating to the writer that there
        are no more bytes to write for the given file.
        """
        ...

    @abstractmethod
    async def job_completed(self):
        """
        This method is called once all files have been downloaded by the job, allowing the writer to complete
        any final processing/cleanup before the BRT exits.
        """
        ...


class ADictWriter:
    """
    Wrapper that enables async interface to csv.DictWriter.
    """

    def __init__(self, aio_file: AsyncBufferedIOBase, fieldnames: List[str]):
        self._aio_file = aio_file
        self._buffer = io.StringIO(newline="")
        self._csv_writer = DictWriter(self._buffer, fieldnames)

    async def _flush(self) -> None:
        await self._aio_file.write(self._buffer.getvalue())
        self._buffer.seek(0)
        self._buffer.truncate(0)

    async def writeheader(self) -> None:
        self._csv_writer.writeheader()
        await self._flush()

    async def writerow(self, row: dict) -> None:
        self._csv_writer.writerow(row)
        await self._flush()

    async def close(self):
        await self._aio_file.close()


class FileWriter(RestoreWriter):
    """Restores files into folder structure they were backed up from."""

    def __init__(self, target: Path):
        self._target = target
        self._handles: dict[str, AsyncBufferedIOBase] = {}
        self._csv_writers: dict[str, ADictWriter] = {}

    async def write_chunk(self, file: RestorableFile, chunk: bytes, forceWinPathSanitization: bool = False):
        if self._handles.get(file.fileId) is None:
            file_path = self._target / file.get_restore_path(forceWinPathSanitization)
            file_path.parent.mkdir(parents=True, exist_ok=True)
            self._handles[file.fileId] = await aiofiles.open(file_path, "wb")
        await self._handles[file.fileId].write(chunk)

    async def write_metadata(
        self,
        file: RestorableFile,
        metadata: Union[LinuxMetadata, WindowsMetadata, MacMetadata],
    ):
        deviceId_os = f"{file.deviceId}_{metadata.os}"
        if self._csv_writers.get(deviceId_os) is None:
            csv_path = self._target / f"{deviceId_os}_metadata.csv"
            aio_file = await aiofiles.open(csv_path, "w", encoding="utf-8")
            fields = ["fileId", "path"] + list(metadata.dict(exclude={"os"}))
            writer = ADictWriter(aio_file, fieldnames=fields)
            await writer.writeheader()
            self._csv_writers[deviceId_os] = writer
        writer = self._csv_writers[deviceId_os]
        m_dict = metadata.dict(exclude={"os"})
        m_dict["fileId"] = file.fileId
        m_dict["path"] = file.path
        await writer.writerow(m_dict)

    async def file_completed(self, file: RestorableFile, forceWinPathSanitization: bool = False) -> None:
        await self._handles[file.fileId].close()
        del self._handles[file.fileId]
        await self._validate_checksum(file, forceWinPathSanitization)
        # apply last modified/access metadata
        sourceLastModifiedNanoSeconds = file.sourceLastModifiedMillis * 1_000_000
        path = self._target / file.get_restore_path(forceWinPathSanitization)
        utime(path, ns=(sourceLastModifiedNanoSeconds, sourceLastModifiedNanoSeconds))

    async def _validate_checksum(self, file: RestorableFile, forceWinPathSanitization: bool = False):
        path = self._target / file.get_restore_path(forceWinPathSanitization)
        async with aiofiles.open(path, "rb") as written_file:
            file_bytes = await written_file.read()
        if file.sourceSHA256Checksum is not None:
            checksum = sha256(file_bytes).hexdigest()
            if checksum != file.sourceSHA256Checksum:
                path.unlink()
                raise Exception(
                    f"Invalid SHA256 checksum for fileId={file.fileId}, expected {file.sourceSHA256Checksum}, got {checksum}."
                )
        else:
            checksum = md5(file_bytes).hexdigest()
            if checksum != file.sourceMD5Checksum:
                path.unlink()
                raise Exception(
                    f"Invalid MD5 checksum for fileId={file.fileId}, expected {file.sourceMD5Checksum}, got {checksum}."
                )

    async def job_completed(self):
        for writer in self._csv_writers.values():
            await writer.close()


class ProcessZipWriter(FileWriter):
    """
    Restores files into per-device .zip archives, using a separate subprocess for each device to not block file download
    progress while zipping.

    The file is first written to disk and validated by FileWriter parent class, then the restored file path is passed
    to the queue shared by the zip subprocess.
    """

    def __init__(self, target: Path, progress: Progress):
        super().__init__(target)
        self._zip_processes: Dict[str, ZipProcess] = {}
        self._progress = progress
        self._tasks: Dict[str, int] = {}
        self._log_queue: SimpleQueue[Tuple[str, str]] = SimpleQueue()
        # if a job was cancelled, there could be files restored that haven't been zipped
        # we need to check for
        for existing_zipfile in target.glob("*.zip"):
            deviceId = existing_zipfile.with_suffix("").name
            self._init_zip_process(deviceId, check_folder=True)

    async def file_completed(self, file: RestorableFile, forceWinPathSanitization: bool = False):
        self._process_log_queue()

        # let FileWriter parent class write file to disk
        await super().file_completed(file, forceWinPathSanitization)

        deviceId = file.deviceId
        file_path = self._target / file.get_restore_path(forceWinPathSanitization)
        # start zip subprocess if this is the first file restored
        if self._zip_processes.get(deviceId) is None:
            self._init_zip_process(deviceId, check_folder=False)

        self._zip_processes[deviceId].queue.put((deviceId, file_path))

    def _init_zip_process(self, deviceId: str, check_folder: bool):
        zip_file = self._target / f"{deviceId}.zip"
        proc = ZipProcess(filename=zip_file, log_queue=self._log_queue)
        proc.start()
        # check for any files in restore folder that haven't been zipped yet (e.g. if user exited mid-zip)
        if check_folder:
            restore_folder = self._target / deviceId
            for file in restore_folder.rglob("*"):
                if file.is_file():
                    proc.queue.put((deviceId, file))
        self._zip_processes[deviceId] = proc
        self._tasks[deviceId] = self._progress.add_task(
            f"Compressing restored files to {zip_file}"
        )

    def _process_log_queue(self):
        while not self._log_queue.empty():
            deviceId, msg = self._log_queue.get()
            if msg.startswith("Error"):
                logger = logging.getLogger(f"{deviceId}-error")
            else:
                logger = logging.getLogger(deviceId)
            logger.info(msg)

    async def job_completed(self):
        for writer in self._csv_writers.values():
            await writer.close()

        # wait for zip processes to complete
        for deviceId, proc in self._zip_processes.items():

            # add metadata CSVs to zip file
            for csv_file in self._target.glob(f"{deviceId}*.csv"):
                proc.queue.put((deviceId, csv_file))

            # send sentinel `None` to queue to tell sub-process no more files are coming
            proc.queue.put(None)  # type: ignore
            proc.join()
            task_id = self._tasks[deviceId]
            self._progress.remove_task(task_id)
        # clean up restore folders
        for deviceId in self._zip_processes:
            device_folder = self._target / deviceId
            rmtree(device_folder, ignore_errors=True)


class ZipProcess(Process):
    """
    Manages creating a zip file in a subprocess.

    Filenames put on the ZipProcess.queue will be added to the zip file and then deleted from disk.
    """

    def __init__(self, filename: Path, log_queue: SimpleQueue):
        self._filename = filename
        self._job_folder = filename.parent
        self._log_queue = log_queue

        self.queue: SimpleQueue[Tuple[str, Path]] = SimpleQueue()
        super().__init__()

    def run(self):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            with ZipFile(file=self._filename, mode="a") as zf:
                while True:
                    item = self.queue.get()
                    if item is None:
                        break
                    else:
                        deviceId, file = item
                        self._log_queue.put(
                            (
                                deviceId,
                                f"Starting zip of file={file} to {self._filename}",
                            )
                        )
                        try:
                            zf.write(
                                file,
                                arcname=file.relative_to(self._job_folder),
                                compress_type=ZIP_LZMA,
                            )
                        except Exception as err:
                            exc = "".join(format_exception(err))
                            self._log_queue.put(
                                (
                                    deviceId,
                                    f"Error zipping file={file} to {self._filename}\n{exc}",
                                )
                            )
                        else:  # leave the restored file on disk if exception is raised, otherwise delete it
                            file.unlink()
                            self._log_queue.put(
                                (
                                    deviceId,
                                    f"Completed zip of file={file} to {self._filename}",
                                )
                            )
