"""Helpers to manage ~/.ssh/known_hosts."""

from datetime import datetime
from os import PathLike
from pathlib import Path
from typing import Annotated, Union

from msgspec import Meta, Struct
from msgspec.json import Decoder as JSONDecoder
from msgspec.json import encode as json_encode
from pytz import UTC
from slices_clientlib_bi.bi_resource import ComputeResource, SSHInfo
from ssh_known_hosts_edit import SSHKnownHostsEdit

from slices_cli_bi.core.const import CONFIG_DIR


class HistoryEntry(Struct, forbid_unknown_fields=True):
    """An entry in the history file we keep when editing ~/.ssh/known_hosts."""

    resource_id: str
    ssh_pub_key: str
    host: str
    expires_at: Annotated[datetime, Meta(tz=True)] | None

    def __hash__(self):
        """Compute the hash of a HistoryEntry."""
        return hash((self.ssh_pub_key, self.resource_id, self.host, self.expires_at))


_history_entry_decoder = JSONDecoder(HistoryEntry)


class SSHKnownHostsManager:
    """A helper class for managing ~/.ssh/known_hosts by adding keys, and removing old keys from it."""

    known_hosts_history_file = Path.home() / CONFIG_DIR / "known_hosts.history"

    def __init__(self, *, known_hosts_file: Union[str, PathLike[str], None] = None):
        self.known_hosts_file = Path.home() / ".ssh" / "known_hosts" if not known_hosts_file else Path(known_hosts_file)

        # TODO refactor with something more performant than list when needed.
        self.history: set[HistoryEntry] = self._load_history_file()
        self.edit = SSHKnownHostsEdit(known_hosts_file_location=str(self.known_hosts_file))

    def _load_history_file(self) -> set[HistoryEntry]:
        if not self.known_hosts_history_file.is_file():
            return set()
        with self.known_hosts_history_file.open() as f:
            return {_history_entry_decoder.decode(e) for e in f.read().splitlines()}

    def _write_history_file(self, history: set[HistoryEntry]):
        if not self.known_hosts_history_file.parent.exists():
            self.known_hosts_history_file.parent.mkdir(parents=True, exist_ok=True)

        with self.known_hosts_history_file.open("w") as f:
            for entry in history:
                f.write(json_encode(entry).decode().replace("\n", "") + "\n")

    def _append_to_history_file(self, entries: list[HistoryEntry]):
        if not self.known_hosts_history_file.parent.exists():
            self.known_hosts_history_file.parent.mkdir(parents=True, exist_ok=True)

        with self.known_hosts_history_file.open("a") as f:
            for entry in entries:
                f.write(json_encode(entry).decode().replace("\n", "") + "\n")

    def _remove_from_history(self, entries: list[HistoryEntry]):
        for entry in entries:
            self.history.remove(entry)
        self._write_history_file(self.history)

    def _find_key_in_history(self, ssh_pub_key: str) -> list[HistoryEntry]:
        return [h for h in self.history if h.ssh_pub_key == ssh_pub_key]

    def _add_internal(
        self, resource: ComputeResource, ssh_login: SSHInfo, public_key: str, extra_history_entries: list[HistoryEntry]
    ) -> bool:
        changed = self.edit.add_to_known_hosts(ssh_login.host, public_key)  # ssh_login.port is not used here
        history_entry = HistoryEntry(
            resource_id=resource.id, ssh_pub_key=public_key, host=ssh_login.host, expires_at=resource.expires_at
        )
        extra_history_entries.append(history_entry)
        self.history.add(history_entry)
        return changed

    def add_login(self, resource: ComputeResource, ssh_login: SSHInfo) -> bool:
        """Add the public keys for a single ssh_login of a resource to ~/.ssh/known_hosts.

        @param resource: the resource.
        @param ssh_login: the ssh_login to add all SSH login public keys for.
        """
        extra_history_entries = []
        changed = False
        if ssh_login.host_keys:
            for public_key in ssh_login.host_keys:
                changed = self._add_internal(resource, ssh_login, public_key, extra_history_entries) or changed

        if ssh_login.jump_proxy and ssh_login.jump_proxy.host_keys:
            for public_key in ssh_login.jump_proxy.host_keys:
                # We add duplicates, but that's OK. They will be removed and expired as expected,
                # but their key will only be removed from known_hosts if the last entry is removed.
                # self._add_internal(resource, ssh_login.jump_proxy, public_key, history_entries)

                # I don't like that bastion is deleted. It's inconvenient...
                changed = self.edit.add_to_known_hosts(ssh_login.jump_proxy.host, public_key) or changed

        self._append_to_history_file(extra_history_entries)

        return changed

    def cleanup(self, inactive_resource_id: str) -> bool:
        """Cleanup the previously stored SSH keys for a single inactive resources.

        @param inactive_resource_id: the ID of the inactive resource to remove from ~/.ssh/known_hosts
        :returns: True if at least one known_host entries was removed, False otherwise.
        """
        entries_to_remove = [
            history_entry for history_entry in self.history if history_entry.resource_id == inactive_resource_id
        ]
        keys_to_remove = [history_entry.ssh_pub_key for history_entry in entries_to_remove]
        if keys_to_remove:
            self._remove_from_history(entries_to_remove)
            # Make sure we don't remove keys from known_hosts if there are still other entries with them
            keys_to_remove = [k for k in keys_to_remove if not self._find_key_in_history(k)]
            return self.edit.remove_from_known_hosts_by_pubkey(keys_to_remove)
        return False

    def cleanup_inactive(self, all_active_resource_ids: list[str]) -> bool:
        """Cleanup previously stored SSH keys for all inactive resources.

        @param all_active_resource_ids: a list of IDs of _all_ currently *active* resources
        """
        entries_to_remove = [
            history_entry for history_entry in self.history if history_entry.resource_id not in all_active_resource_ids
        ]
        keys_to_remove = [history_entry.ssh_pub_key for history_entry in entries_to_remove]
        if keys_to_remove:
            self._remove_from_history(entries_to_remove)
            # Make sure we don't remove keys from known_hosts if there are still other entries with them
            keys_to_remove = [k for k in keys_to_remove if not self._find_key_in_history(k)]
            return self.edit.remove_from_known_hosts_by_pubkey(keys_to_remove)
        return False

    def cleanup_expired(self) -> bool:
        """Cleanup previously stored SSH keys for all expired resources."""
        now = datetime.now(UTC)
        entries_to_remove = [
            history_entry
            for history_entry in self.history
            if history_entry.expires_at and history_entry.expires_at < now
        ]
        keys_to_remove = [history_entry.ssh_pub_key for history_entry in entries_to_remove]
        if keys_to_remove:
            self._remove_from_history(entries_to_remove)
            # Make sure we don't remove keys from known_hosts if there are still other entries with them
            keys_to_remove = [k for k in keys_to_remove if not self._find_key_in_history(k)]
            return self.edit.remove_from_known_hosts_by_pubkey(keys_to_remove)
        return False
