import os
import base64
from sys import stdout
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ed25519
from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes, PublicKeyTypes

# TODO guard against bad file access
def create(path: str) -> None:
    private_key = ed25519.Ed25519PrivateKey.generate()
    public_key = private_key.public_key()
    enc = serialization.NoEncryption()

    os.mkdir(path)

    data = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=enc,
    )
    with open(os.path.join(path, "private.pem"), "wb") as f:
        f.write(data)

    data = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    with open(os.path.join(path, "public.pub"), "wb") as f:
        f.write(data)

def load(identities_dir: str, identity: str) -> tuple[PrivateKeyTypes, PublicKeyTypes]:
    if not os.path.isdir(identities_dir):
        os.makedirs(identities_dir)

    path = os.path.join(identities_dir, identity)

    if not os.path.isdir(path):
        stdout.write(f"Generating new identity '{identity}'...\n")
        create(path)
        stdout.write(f"Created new identity '{identity}' at: {path}\n")

    with open(os.path.join(path, "private.pem"), "rb") as f:
        private_key_data = f.read()
    with open(os.path.join(path, "public.pub"), "rb") as f:
        public_key_data = f.read()

    return (
        serialization.load_pem_private_key(private_key_data, password=None),
        serialization.load_pem_public_key(public_key_data)
    )

def id_from_pubkey(pubkey: PublicKeyTypes) -> bytes:
    der = pubkey.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    b64 = base64.b64encode(der)
    return b64[16:24]