"""Database helpers.

Allows pre-loading database for tests.
"""

from collections.abc import Iterator
from contextlib import contextmanager
import httpx
from sshecret.crypto import generate_private_key, generate_public_key_string


class DatabasePreloader:
    """Database preloader class."""

    def __init__(self, admin_url: str, username: str, password: str) -> None:
        """Instantiate class to populate database."""
        self.admin_url: str = admin_url
        self.username: str = username
        self.password: str = password

    @contextmanager
    def login(self) -> Iterator[httpx.Client]:
        """Login and yield client."""
        login_client = httpx.Client(base_url=self.admin_url)
        resp = login_client.post(
            "api/v1/token",
            data={"username": self.username, "password": self.password}
        )
        assert resp.status_code == 200
        data = resp.json()
        token = data["access_token"]
        headers = {"Authorization": f"Bearer {token}"}
        with httpx.Client(base_url=self.admin_url, headers=headers) as client:
            yield client

    def create_client(self, *names: str) -> None:
        """Create one or more clients."""
        with self.login() as http_client:
            for name in names:
                private_key = generate_private_key()
                public_key = generate_public_key_string(private_key.public_key())
                data = {
                    "name": name,
                    "description": "Test client",
                    "public_key": public_key,
                    "sources": ["0.0.0.0/0", "::/0"],
                }
                resp = http_client.post("api/v1/clients/", json=data)
                resp.raise_for_status()

    def create_secret(self, *secrets: tuple[str, list[str]]) -> None:
        """Create secret.

        Argument format is (secret_name, [client1, client2, ...])

        Clients must exist.
        """
        with self.login() as http_client:
            for name, clients in secrets:
                data = {
                    "name": name,
                    "clients": clients,
                    "value": {"auto_generate": True, "length": 32}
                }

                http_client.post("api/v1/secrets/", json=data)