from . import config, crypto, utils

from typing import Optional, Sequence
from datetime import timedelta
import os
import logging

logger = logging.getLogger("ssh-ca")


class SSHCA:
    def __init__(self, config: config.Config) -> None:
        self.config = config

    def do_sign(
        self,
        public_key_bytes: bytes,
        certificate_type: crypto.SSHCertificateType,
        *,
        key_id: Optional[str] = None,  # type: ignore
        principals: Sequence[str],
        valid_time: Optional[timedelta] = None,
    ) -> crypto.SSHCertificate:
        """Sign a public key with a CA key."""
        public_key = crypto.loads_pub_key(public_key_bytes)

        if certificate_type == crypto.SSHCertificateType.HOST:
            ca_key = self._load_host_key()
            # Load default key_id and valid_time from config
            valid_time = valid_time or utils.parse_td(
                self.config.lookup("signing.host.valid_time")
            )

            key_id: str = key_id or self.config.lookup("signing.host.key_id")
        else:
            ca_key = self._load_user_key()

            # Load default key_id and valid_time from config
            valid_time = valid_time or utils.parse_td(
                self.config.lookup("signing.user.valid_time")
            )

            key_id: str = key_id or self.config.lookup("signing.user.key_id")

        logger.info(f"Signing public key with {key_id=} {principals=} {valid_time=}")

        cert = crypto.sign_key(
            ca_key,
            public_key,
            certificate_type,
            key_id=key_id.encode("utf-8"),
            principals=[p.encode("utf-8") for p in principals],
            valid_time=valid_time,
        )

        return cert

    def do_init(self):
        """
        Perform initial setup.

        This reads the CA key paths from the configuration and generates the
        keys if they do not already exist.

        If setup is already done, this is a no-op.
        """
        host_private_key_path = self.config.lookup("signing.ca-keys.host")
        host_public_key_path = f"{host_private_key_path}.pub"

        user_private_key_path = self.config.lookup("signing.ca-keys.user")
        user_public_key_path = f"{user_private_key_path}.pub"

        host_key_dir = os.path.dirname(host_private_key_path)
        user_key_dir = os.path.dirname(user_private_key_path)

        # Create any intermediate directories if needed
        os.makedirs(host_key_dir, exist_ok=True)
        os.makedirs(user_key_dir, exist_ok=True)

        all_exist = all(
            os.path.exists(p)
            for p in (
                host_private_key_path,
                host_public_key_path,
                user_private_key_path,
                user_public_key_path,
            )
        )

        if all_exist:
            logger.info("No need to re-initialise CA keys; they already exist.")
            return

        if os.path.exists(user_private_key_path) or os.path.exists(
            user_public_key_path
        ):
            raise ValueError(
                "User CA key already exists - cannot re-initialise safely."
            )

        if os.path.exists(host_private_key_path) or os.path.exists(
            host_public_key_path
        ):
            raise ValueError(
                "Host CA key already exists - cannot re-initialise safely."
            )

        logger.info("Performing initialisation...")

        for priv, pub in (
            (host_private_key_path, host_public_key_path),
            (user_private_key_path, user_public_key_path),
        ):
            key_pair = crypto.generate_key()
            crypto.save_key(key_pair, priv, pub)
            logger.info(f"Generated key pair ({priv!r}, {pub!r})")

    def _load_user_key(self) -> crypto.KeyPair:
        user_key_path = self.config.lookup("signing.ca-keys.user")
        user_key = crypto.load_key(user_key_path, f"{user_key_path}.pub")
        return user_key

    def _load_host_key(self) -> crypto.KeyPair:
        host_key_path = self.config.lookup("signing.ca-keys.host")
        host_key = crypto.load_key(host_key_path, f"{host_key_path}.pub")
        return host_key
