"""Local backup support for Core and Container installations."""

import asyncio
from collections.abc import AsyncIterator, Callable, Coroutine
import copy
from dataclasses import dataclass, replace
from io import BytesIO
import json
from pathlib import Path, PurePath, PureWindowsPath
from queue import SimpleQueue
import tarfile
import threading
from typing import IO, Any, cast

import aiohttp
from securetar import (
    InvalidPasswordError,
    SecureTarArchive,
    SecureTarError,
    SecureTarFile,
    SecureTarReadError,
    SecureTarRootKeyContext,
    get_archive_max_ciphertext_size,
)

from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.util import dt as dt_util
from homeassistant.util.async_iterator import (
    Abort,
    AsyncIteratorReader,
    AsyncIteratorWriter,
)
from homeassistant.util.json import JsonObjectType, json_loads_object

from .const import BUF_SIZE, LOGGER, SECURETAR_CREATE_VERSION
from .models import AddonInfo, AgentBackup, Folder, InvalidBackupFilename


class DecryptError(HomeAssistantError):
    """Error during decryption."""

    _message = "Unexpected error during decryption."


class EncryptError(HomeAssistantError):
    """Error during encryption."""

    _message = "Unexpected error during encryption."


class UnsupportedSecureTarVersion(DecryptError):
    """Unsupported securetar version."""

    _message = "Unsupported securetar version."


class IncorrectPassword(DecryptError):
    """Invalid password or corrupted backup."""

    _message = "Invalid password or corrupted backup."


class BackupEmpty(DecryptError):
    """No tar files found in the backup."""

    _message = "No tar files found in the backup."


def make_backup_dir(path: Path) -> None:
    """Create a backup directory if it does not exist."""
    path.mkdir(exist_ok=True)


def read_backup(backup_path: Path) -> AgentBackup:
    """Read a backup from disk."""

    with tarfile.open(backup_path, "r:", bufsize=BUF_SIZE) as backup_file:
        if not (data_file := backup_file.extractfile("./backup.json")):
            raise KeyError("backup.json not found in tar file")
        data = json_loads_object(data_file.read())
        addons = [
            AddonInfo(
                name=cast(str, addon["name"]),
                slug=cast(str, addon["slug"]),
                version=cast(str, addon["version"]),
            )
            for addon in cast(list[JsonObjectType], data.get("addons", []))
        ]

        folders = [
            Folder(folder)
            for folder in cast(list[str], data.get("folders", []))
            if folder != "homeassistant"
        ]

        homeassistant_included = False
        homeassistant_version: str | None = None
        database_included = False
        if (
            homeassistant := cast(JsonObjectType, data.get("homeassistant"))
        ) and "version" in homeassistant:
            homeassistant_included = True
            homeassistant_version = cast(str, homeassistant["version"])
            database_included = not cast(
                bool, homeassistant.get("exclude_database", False)
            )

        extra_metadata = cast(dict[str, bool | str], data.get("extra", {}))
        date = extra_metadata.get("supervisor.backup_request_date", data["date"])

        name = cast(str, data["name"])
        # The name is used to derive the on-disk filename via suggested_filename;
        # reject anything that could escape the backup directory.
        safe_name = PureWindowsPath(name).name
        if safe_name != name or name in ("", ".", ".."):
            raise InvalidBackupFilename(f"Invalid backup name: {name!r}")

        return AgentBackup(
            addons=addons,
            backup_id=cast(str, data["slug"]),
            database_included=database_included,
            date=cast(str, date),
            extra_metadata=extra_metadata,
            folders=folders,
            homeassistant_included=homeassistant_included,
            homeassistant_version=homeassistant_version,
            name=name,
            protected=cast(bool, data.get("protected", False)),
            size=backup_path.stat().st_size,
        )


def suggested_filename_from_name_date(name: str, date_str: str) -> str:
    """Suggest a filename for the backup."""
    date = dt_util.parse_datetime(date_str, raise_on_error=True)
    return "_".join(f"{name} {date.strftime('%Y-%m-%d %H.%M %S%f')}.tar".split())


def suggested_filename(backup: AgentBackup) -> str:
    """Suggest a filename for the backup."""
    return suggested_filename_from_name_date(backup.name, backup.date)


def validate_password(path: Path, password: str | None) -> bool:
    """Validate the password.

    This assumes every inner tar is encrypted with the same secure tar version and
    same password.
    """
    with SecureTarArchive(
        path, "r", bufsize=BUF_SIZE, password=password
    ) as backup_file:
        compressed = False
        ha_tar_name = "homeassistant.tar"
        try:
            ha_tar = backup_file.tar.extractfile(ha_tar_name)
        except KeyError:
            compressed = True
            ha_tar_name = "homeassistant.tar.gz"
            try:
                ha_tar = backup_file.tar.extractfile(ha_tar_name)
            except KeyError:
                LOGGER.error("No homeassistant.tar or homeassistant.tar.gz found")
                return False
        try:
            with SecureTarFile(
                path,  # Not used
                gzip=compressed,
                password=password,
                fileobj=ha_tar,
            ):
                # If we can read the tar file, the password is correct
                return True
        except tarfile.ReadError, InvalidPasswordError, SecureTarReadError:
            LOGGER.debug("Invalid password")
            return False
        except Exception:  # noqa: BLE001
            LOGGER.exception("Unexpected error validating password")
    return False


def validate_password_stream(
    input_stream: IO[bytes],
    password: str | None,
) -> None:
    """Validate the password.

    This assumes every inner tar is encrypted with the same secure tar version and
    same password.
    """
    with SecureTarArchive(
        fileobj=input_stream,
        mode="r",
        bufsize=BUF_SIZE,
        streaming=True,
        password=password,
    ) as input_archive:
        for obj in input_archive.tar:
            if not obj.name.endswith((".tar", ".tgz", ".tar.gz")):
                continue
            try:
                with input_archive.extract_tar(obj) as decrypted:
                    if decrypted.plaintext_size is None:
                        raise UnsupportedSecureTarVersion
                    decrypted.read(1)  # Read a single byte to trigger the decryption
            except (InvalidPasswordError, SecureTarReadError) as err:
                raise IncorrectPassword from err
            else:
                return
    raise BackupEmpty


def _get_expected_archives(backup: AgentBackup) -> set[str]:
    """Get the expected archives in the backup."""
    expected_archives = set()
    if backup.homeassistant_included:
        expected_archives.add("homeassistant")
    for addon in backup.addons:
        expected_archives.add(addon.slug)
    for folder in backup.folders:
        expected_archives.add(folder.value)
    return expected_archives


def decrypt_backup(
    backup: AgentBackup,
    input_stream: IO[bytes],
    output_stream: IO[bytes],
    password: str | None,
    on_done: Callable[[Exception | None], None],
    minimum_size: int,
    key_context: SecureTarRootKeyContext,
) -> None:
    """Decrypt a backup."""
    error: Exception | None = None
    try:
        try:
            with (
                SecureTarArchive(
                    fileobj=input_stream,
                    mode="r",
                    bufsize=BUF_SIZE,
                    streaming=True,
                    password=password,
                ) as input_archive,
                tarfile.open(
                    fileobj=output_stream, mode="w|", bufsize=BUF_SIZE
                ) as output_tar,
            ):
                _decrypt_backup(backup, input_archive, output_tar)
        except (DecryptError, SecureTarError, tarfile.TarError) as err:
            LOGGER.warning("Error decrypting backup: %s", err)
            error = err
        except Abort:
            raise
        except Exception as err:  # noqa: BLE001
            LOGGER.exception("Unexpected error when decrypting backup: %s", err)
            error = err
        else:
            # Pad the output stream to the requested minimum size
            padding = max(minimum_size - output_stream.tell(), 0)
            output_stream.write(b"\0" * padding)
        finally:
            # Write an empty chunk to signal the end of the stream
            output_stream.write(b"")
    except Abort:
        LOGGER.debug("Cipher operation aborted")
    finally:
        on_done(error)


def _decrypt_backup(
    backup: AgentBackup,
    input_archive: SecureTarArchive,
    output_tar: tarfile.TarFile,
) -> None:
    """Decrypt a backup."""
    expected_archives = _get_expected_archives(backup)
    for obj in input_archive.tar:
        # We compare with PurePath to avoid issues with different path separators,
        # for example when backup.json is added as "./backup.json"
        object_path = PurePath(obj.name)
        if object_path == PurePath("backup.json"):
            # Rewrite the backup.json file to indicate that the backup is decrypted
            if not (reader := input_archive.tar.extractfile(obj)):
                raise DecryptError
            metadata = json_loads_object(reader.read())
            metadata["protected"] = False
            updated_metadata_b = json.dumps(metadata).encode()
            metadata_obj = copy.deepcopy(obj)
            metadata_obj.size = len(updated_metadata_b)
            output_tar.addfile(metadata_obj, BytesIO(updated_metadata_b))
            continue
        prefix, _, suffix = object_path.name.partition(".")
        if suffix not in ("tar", "tgz", "tar.gz"):
            LOGGER.debug("Unknown file %s will not be decrypted", obj.name)
            output_tar.addfile(obj, input_archive.tar.extractfile(obj))
            continue
        if prefix not in expected_archives:
            LOGGER.debug("Unknown inner tar file %s will not be decrypted", obj.name)
            output_tar.addfile(obj, input_archive.tar.extractfile(obj))
            continue
        with input_archive.extract_tar(obj) as decrypted:
            # Guard against SecureTar v1 which doesn't store plaintext size
            if (plaintext_size := decrypted.plaintext_size) is None:
                raise UnsupportedSecureTarVersion
            decrypted_obj = copy.deepcopy(obj)
            decrypted_obj.size = plaintext_size
            output_tar.addfile(decrypted_obj, decrypted)


def encrypt_backup(
    backup: AgentBackup,
    input_stream: IO[bytes],
    output_stream: IO[bytes],
    password: str | None,
    on_done: Callable[[Exception | None], None],
    minimum_size: int,
    key_context: SecureTarRootKeyContext,
) -> None:
    """Encrypt a backup."""
    error: Exception | None = None
    try:
        try:
            with (
                tarfile.open(
                    fileobj=input_stream, mode="r|", bufsize=BUF_SIZE
                ) as input_tar,
                SecureTarArchive(
                    fileobj=output_stream,
                    mode="w",
                    bufsize=BUF_SIZE,
                    streaming=True,
                    root_key_context=key_context,
                    create_version=SECURETAR_CREATE_VERSION,
                ) as output_archive,
            ):
                _encrypt_backup(backup, input_tar, output_archive)
        except (EncryptError, SecureTarError, tarfile.TarError) as err:
            LOGGER.warning("Error encrypting backup: %s", err)
            error = err
        except Abort:
            raise
        except Exception as err:  # noqa: BLE001
            LOGGER.exception("Unexpected error when encrypting backup: %s", err)
            error = err
        else:
            # Pad the output stream to the requested minimum size
            padding = max(minimum_size - output_stream.tell(), 0)
            output_stream.write(b"\0" * padding)
        finally:
            # Write an empty chunk to signal the end of the stream
            output_stream.write(b"")
    except Abort:
        LOGGER.debug("Cipher operation aborted")
    finally:
        on_done(error)


def _encrypt_backup(
    backup: AgentBackup,
    input_tar: tarfile.TarFile,
    output_archive: SecureTarArchive,
) -> None:
    """Encrypt a backup."""
    inner_tar_idx = 0
    expected_archives = _get_expected_archives(backup)
    for obj in input_tar:
        # We compare with PurePath to avoid issues with different path separators,
        # for example when backup.json is added as "./backup.json"
        object_path = PurePath(obj.name)
        if object_path == PurePath("backup.json"):
            # Rewrite the backup.json file to indicate that the backup is encrypted
            if not (reader := input_tar.extractfile(obj)):
                raise EncryptError
            metadata = json_loads_object(reader.read())
            metadata["protected"] = True
            updated_metadata_b = json.dumps(metadata).encode()
            metadata_obj = copy.deepcopy(obj)
            metadata_obj.size = len(updated_metadata_b)
            output_archive.tar.addfile(metadata_obj, BytesIO(updated_metadata_b))
            continue
        prefix, _, suffix = object_path.name.partition(".")
        if suffix not in ("tar", "tgz", "tar.gz"):
            LOGGER.debug("Unknown file %s will not be encrypted", obj.name)
            output_archive.tar.addfile(obj, input_tar.extractfile(obj))
            continue
        if prefix not in expected_archives:
            LOGGER.debug("Unknown inner tar file %s will not be encrypted", obj.name)
            continue
        if (fileobj := input_tar.extractfile(obj)) is None:
            LOGGER.debug(
                "Non regular inner tar file %s will not be encrypted", obj.name
            )
            continue
        output_archive.import_tar(fileobj, obj, derived_key_id=inner_tar_idx)
        inner_tar_idx += 1


@dataclass(kw_only=True)
class _CipherWorkerStatus:
    done: asyncio.Event
    error: Exception | None = None
    reader: AsyncIteratorReader
    thread: threading.Thread
    writer: AsyncIteratorWriter


class _CipherBackupStreamer:
    """Encrypt or decrypt a backup."""

    _cipher_func: Callable[
        [
            AgentBackup,
            IO[bytes],
            IO[bytes],
            str | None,
            Callable[[Exception | None], None],
            int,
            SecureTarRootKeyContext,
        ],
        None,
    ]

    def __init__(
        self,
        hass: HomeAssistant,
        backup: AgentBackup,
        open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
        password: str,
    ) -> None:
        """Initialize."""
        self._workers: list[_CipherWorkerStatus] = []
        self._backup = backup
        self._hass = hass
        self._open_stream = open_stream
        self._password = password
        self._key_context = SecureTarRootKeyContext(password)

    def size(self) -> int:
        """Return the maximum size of the decrypted or encrypted backup."""
        return get_archive_max_ciphertext_size(
            self._backup.size, SECURETAR_CREATE_VERSION, self._num_tar_files()
        )

    def _num_tar_files(self) -> int:
        """Return the number of inner tar files."""
        b = self._backup
        return len(b.addons) + len(b.folders) + b.homeassistant_included + 1

    async def open_stream(self) -> AsyncIterator[bytes]:
        """Open a stream."""

        def on_done(error: Exception | None) -> None:
            """Call by the worker thread when it's done."""
            worker_status.error = error
            self._hass.loop.call_soon_threadsafe(worker_status.done.set)

        stream = await self._open_stream()
        reader = AsyncIteratorReader(self._hass.loop, stream)
        writer = AsyncIteratorWriter(self._hass.loop)
        worker = threading.Thread(
            target=self._cipher_func,
            args=[
                self._backup,
                reader,
                writer,
                self._password,
                on_done,
                self.size(),
                self._key_context,
            ],
        )
        worker_status = _CipherWorkerStatus(
            done=asyncio.Event(), reader=reader, thread=worker, writer=writer
        )
        self._workers.append(worker_status)
        worker.start()
        return writer

    async def wait(self) -> None:
        """Wait for the worker threads to finish."""
        for worker in self._workers:
            worker.reader.abort()
            worker.writer.abort()
        await asyncio.gather(*(worker.done.wait() for worker in self._workers))


class DecryptedBackupStreamer(_CipherBackupStreamer):
    """Decrypt a backup."""

    _cipher_func = staticmethod(decrypt_backup)

    def backup(self) -> AgentBackup:
        """Return the decrypted backup."""
        return replace(self._backup, protected=False, size=self.size())


class EncryptedBackupStreamer(_CipherBackupStreamer):
    """Encrypt a backup."""

    _cipher_func = staticmethod(encrypt_backup)

    def backup(self) -> AgentBackup:
        """Return the encrypted backup."""
        return replace(self._backup, protected=True, size=self.size())


async def receive_file(
    hass: HomeAssistant, contents: aiohttp.BodyPartReader, path: Path
) -> None:
    """Receive a file from a stream and write it to a file."""
    queue: SimpleQueue[tuple[bytes, asyncio.Future[None] | None] | None] = SimpleQueue()

    def _sync_queue_consumer() -> None:
        with path.open("wb") as file_handle:
            while True:
                if (_chunk_future := queue.get()) is None:
                    break
                _chunk, _future = _chunk_future
                if _future is not None:
                    hass.loop.call_soon_threadsafe(_future.set_result, None)
                file_handle.write(_chunk)

    fut: asyncio.Future[None] | None = None
    try:
        fut = hass.async_add_executor_job(_sync_queue_consumer)
        megabytes_sending = 0
        while chunk := await contents.read_chunk(BUF_SIZE):
            megabytes_sending += 1
            if megabytes_sending % 5 != 0:
                queue.put_nowait((chunk, None))
                continue

            chunk_future = hass.loop.create_future()
            queue.put_nowait((chunk, chunk_future))
            await asyncio.wait(
                (fut, chunk_future),
                return_when=asyncio.FIRST_COMPLETED,
            )
            if fut.done():
                # The executor job failed
                break

        queue.put_nowait(None)  # terminate queue consumer
    finally:
        if fut is not None:
            await fut
