import base64
import shutil
import subprocess
from datetime import datetime
from pathlib import Path
from typing import Optional

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa

from mailiness import g

from . import settings

debug = getattr(g, "debug", False)


def get_default_selector():
    return datetime.now().strftime("%Y%m%d")


class DKIM:
    def __init__(self, domain: str, selector: Optional[str] = None):
        self.dkim_maps_path = g.config["spam"]["dkim_maps_path"]
        self.dkim_private_key_dir = g.config["spam"]["dkim_private_key_directory"]
        dkim_map = self.load_from_dkim_map_file()
        self.private_key = None
        if domain in dkim_map.keys():
            self.selector = dkim_map[domain]
            with (
                Path(self.dkim_private_key_dir) / Path(f"{domain}.{self.selector}.key")
            ).open("rb") as fp:
                self.private_key = serialization.load_pem_private_key(
                    fp.read(), password=None
                )
        else:
            self.generate_key()
            if selector is None:
                self.selector = get_default_selector()
            else:
                self.selector = selector

        self.domain = domain

    def generate_key(self):
        self.private_key = rsa.generate_private_key(
            public_exponent=settings.RSA_PUBLIC_EXPONENT,
            key_size=settings.DKIM_KEY_SIZE,
        )

    def private_key_as_pem(self) -> str:
        as_bytes = self.private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
        return as_bytes.decode("utf-8")

    def public_key_as_der(self) -> bytes:
        return self.private_key.public_key().public_bytes(
            encoding=serialization.Encoding.DER,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )

    def load_dkim_map(self, data: str) -> dict:
        result = {}
        for line in data.splitlines():
            domain, selector = line.split(" ")
            result[domain] = selector
        return result

    def write_dkim_map(self, data: dict):
        shutil.copy2(self.dkim_maps_path, self.dkim_maps_path + ".bak")

        new_map = ""
        for domain, selector in data.items():
            new_map += f"{domain} {selector}\n"

        with Path(self.dkim_maps_path).open("w", encoding="utf-8") as fp:
            fp.write(new_map)

    def dns_txt_record(self) -> str:
        b64_der_pubkey = base64.b64encode(self.public_key_as_der())
        txt_record = ""
        txt_record += f"DNS TXT Record for {self.domain}\n\n\n\n"
        txt_record += f"Name:\n\n{self.selector}._domainkey\n\n\n"
        txt_record += f"Content:\n\nv=DKIM1;k=rsa;p={b64_der_pubkey.decode('utf-8')}"
        return txt_record

    def delete_key(self):
        dkim_map = self.load_from_dkim_map_file()
        del dkim_map[self.domain]
        self.write_dkim_map(dkim_map)

        key_file = Path(self.dkim_private_key_dir) / Path(
            f"{self.domain}.{self.selector}.key"
        )

        key_file.unlink(missing_ok=True)

        if not debug:
            self._run_system_hooks()

    def load_from_dkim_map_file(self) -> dict:
        dkim_map_file = Path(self.dkim_maps_path)

        if dkim_map_file.exists():
            with dkim_map_file.open("r", encoding="utf-8") as fp:
                return self.load_dkim_map(fp.read())
        return {}

    def save_private_key(self):
        dest = Path(self.dkim_private_key_dir) / Path(
            f"{self.domain}.{self.selector}.key"
        )

        with dest.open("w", encoding="utf-8") as fp:
            fp.write(self.private_key_as_pem())

        dkim_map = self.load_from_dkim_map_file()

        dkim_map[self.domain] = self.selector

        self.write_dkim_map(dkim_map)

        if not debug:
            self._run_system_hooks(dest)

    def _run_system_hooks(self, keyfile=None):
        subprocess.run(["systemctl", "reload", "rspamd"])
        if keyfile:
            shutil.chown(str(keyfile), user="_rspamd", group="_rspamd")
