#!/usr/bin/python3

import os, sys, re
import filecmp
from glob import glob
from mako.template import Template
from pathlib import Path
import subprocess
import yaml


def backupProj(project):
    database = None
    paths = []
    sqlite_db = None

    if project == 'postgres':
        database = 'all'
    elif project in secrets['postgres'].keys():
        database = project

    if project in env['backup'].keys():
        for path in env['backup'][project]:
            paths += [path]

    if project in env['backup_sqlite'].keys():
        sqlite_db = env['backup_sqlite'][project]

    if database is not None or sqlite_db is not None or len(paths) > 0:
        print(f"Running backup for project {project}.")

        if len(paths) == 0:
            borgCreate(project, database=database, sqlite_db=sqlite_db)
        else:
            borgCreate(project, path=paths[0], database=database, sqlite_db=sqlite_db)

        for path in paths[1:]:
            archiveName = re.sub('/', '-', path)
            borgCreate(f"{project}-{archiveName}", path=path)


def borgCreate(name, path=None, database=None, sqlite_db=None):
    if path is None and database is None and sqlite_db is None:
        print(f"Cannot create backup, you must pass at least one parameter amongst: path, database, database_path (archive name = {name}).", file=sys.stderr)
        return 1

    borgPaths = []

    if path is not None:
        absPath = Path(path).absolute()

        parentDir = absPath.parent.absolute()
        os.chdir(parentDir)

        borgPaths += [absPath.relative_to(parentDir)]

        borgInput = None

    if database is not None:
        print(f"Dumping database {database}.")

        if database == 'all':
            dumpProc = subprocess.run(["podman", "exec", "postgres", "pg_dumpall"], capture_output=True, text=True)
        else:
            dumpProc = subprocess.run(["podman", "exec", "postgres", "pg_dump", "-c", database], capture_output=True, text=True)

        if dumpProc.returncode != 0:
            print(f"Failed to dump database {database}.", file=sys.stderr)
            return 1

        borgPaths += ["-", "--stdin-name", f"{database}.sql"]

        borgInput = dumpProc.stdout

    if sqlite_db is not None:
        print(f"Dumping SQLite database {sqlite_db}.")

        sqlite_args = ["podman", "run", "--rm", "-v", f"{sqlite_db}:/db.sqlite", "docker.io/library/alpine:latest", "sh", "-c", "apk add sqlite &>/dev/null; sqlite3 /db.sqlite .dump"]
        dumpProc = subprocess.run(sqlite_args, capture_output=True, text=True)

        if dumpProc.returncode != 0:
            print(f"Failed to dump SQLite database {sqlite_db}.", file=sys.stderr)
            return 1

        database = Path(sqlite_db).stem

        borgPaths += ["-", "--stdin-name", f"{database}.sqlite"]

        borgInput = dumpProc.stdout

    if path is None:
        print(f"Creating archive '{name}' from database {database}.")
    elif database is None:
        print(f"Creating archive '{name}' from path {path}.")
    else:
        print(f"Creating archive '{name}' from database {database} and path {absPath}.")

    borgArgs = ["create", "--compression=lzma", f"{env['borg_repo']}::{name}_{{utcnow}}"]
    runBorg(borgArgs + borgPaths, input=borgInput)

    os.chdir(os.path.realpath(sys.path[0]))

    print(f"Pruning archives '{name}_*'.")
    runBorg(["prune", "--glob-archives", f"{name}_*"] + env['borg_prune_opts'] + [env['borg_repo']])


def getImageId (image):
    return runPodman("image", ["inspect", "--format", "{{.Id}}", image]).stdout.strip()


def getUid(service):
    if service in env['users'].keys() and env['users'][service] != 0:
        return env['users'][service] + env['uid_shift']
    else:
        return env['podman_uid']


def pullProj(project):
    print(f"Pulling images for project {project}.")
    
    with open(f"projects/{project}/compose.yaml.rendered", 'r') as composefile:
        images = re.findall('(?<=image:\\s).+', composefile.read())

    pulledImages = []
    for image in images:
        currentId = getImageId(image)
        if re.search('^localhost/', image):
            runPodman("compose", ["-f", f"projects/{project}/compose.yaml.rendered", "build", "--pull"])
        else:
            runPodman("pull", image)
        pulledId = getImageId(image)
        if currentId != pulledId:
            pulledImages += image
            print(f"Pulled new version of image {image}.")
        else:
            print(f"No update available for image {image}.")

    return pulledImages


def renderFile(templateFile):
    print(f"Rendering file {templateFile}.")

    renderedFilename = re.sub('\\.mako$', '.rendered', templateFile)

    template = Template(filename=templateFile)

    outputFile = open(renderedFilename, "w")
    outputFile.write(template.render(env=env, secrets=secrets))
    outputFile.close()


def runBorg(args, input=None):
    if isinstance(args, str):
        args = args.split(' ')

    borgEnv = {"BORG_PASSPHRASE": secrets['borg']}
    proc = subprocess.run(["borg"] + args, input=input, capture_output=True, text=True, env=borgEnv)

    stderr = proc.stderr.strip()
    if stderr != '':
        print(stderr, file=sys.stderr)

    return proc


def runPodman(cmd, args):
    if isinstance(args, str):
        args = args.split(' ')

    if cmd == 'compose':
        runArgs = ["podman-compose"] + args
    else:
        runArgs = ["podman", cmd] + args

    if env['rootless']:
        proc = subprocess.run(runArgs, capture_output=True, text=True)
        stderr = proc.stderr.strip()
    else:
        proc = runSudo(runArgs)
        stderr = proc.stderr.read().strip()

    if stderr != '':
        print(stderr, file=sys.stderr)

    return proc


def runSudo(args):
    if isinstance(args, str):
        args = args.split(' ')

    child = subprocess.Popen(["sudo"] + args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

    if re.search('^\\[sudo\\] password for .+', child.stdout.read().strip()):
        child.communicate(input=input().encode())[0]

    child.wait()

    stderr = child.stderr.read().strip()
    if stderr != '':
        print(stderr, file=sys.stderr)

    return child


def setCertPerms(service):
    dirs = ["/etc/letsencrypt", "/etc/letsencrypt/live", "/etc/letsencrypt/archive"]

    pkeyFile = env['certs'][service]['pkey']   

    domain_dir = re.search('.+(?=\\/.+$)', pkeyFile).group(0)
    dirs += [domain_dir, re.sub('live', 'archive', domain_dir)]

    for path in dirs:
        setOwner(path, 0, env['podman_uid'])
        setPerms(path, 711)

    setOwner(pkeyFile, 0, getUid(service))
    setPerms(pkeyFile, 640)


def setNftables():
    print()
    renderFile("nftables.conf.mako")
    if not filecmp.cmp("nftables.conf.rendered", "/etc/nftables.conf"):
        print("nftables.conf changed, copying new version.")
        runSudo(["cp", "nftables.conf.rendered", "/etc/nftables.conf"])
        runSudo(["systemctl", "restart", "nftables"])
    else:
        print("nftables.conf unchanged.")


def setOwner(path, uid=None, gid=None):
    stat = os.stat(path)
    if uid is None:
        uid = stat.st_uid
    if gid is None:
        gid = stat.st_gid    

    if stat.st_uid != uid or stat.st_gid != gid:
        print(f"Changing ownership of {path} to {uid}:{gid} from {stat.st_uid}:{stat.st_gid}.")
        runSudo(["chown", f"{uid}:{gid}", path])
    else:
        print(f"Ownership of {path} already set to {uid}:{gid}.")


def setPerms(path, mode):
    mode = str(mode)
    stat = os.stat(path)
    curMode = oct(stat.st_mode)[-3:]
    if mode != curMode:
        print(f"Changing permissions of {path} to {mode} from {curMode}.")
        if stat.st_uid == os.getuid():
            subprocess.run(["chmod", mode, path])
        else:
            runSudo(["chmod", mode, path])
    else:
        print(f"Permissions of {path} already set to {mode}.")


def setupProj(project):
    print(f"Running setup for project {project}.")

    if os.path.isfile(f"projects/{project}/compose.yaml.rendered"):
        backupProj(project)

    if project == 'diun':
        if not os.path.isfile(env['socket']):
            if env['rootless']:
                subprocess.run(["systemctl", "--user", "restart", "podman.socket"])
            else:
                runSudo("systemctl restart podman.socket")
        setPerms(env['socket'], 660)
        setOwner(env['socket'], env['podman_uid'], getUid('diun'))

    if project in env['certs'].keys():
        setCertPerms(project)
        
    for templateFile in glob(f"projects/{project}/*.mako", include_hidden=True):
        renderFile(templateFile)
        renderedFilename = re.sub('\\.mako$', '.rendered', templateFile)
        setPerms(renderedFilename, 640)
        setOwner(renderedFilename, os.getuid(), getUid(project))

    if project in env['volumes']:
        for volume in env['volumes'][project].values():
            if not os.path.exists(volume):
                os.makedirs(volume)
            setPerms(volume, 750)
            setOwner(volume, getUid(project), env['podman_uid'])

    upProj(project)


def upProj(project):
    if runPodman("pod", ["exists", f"pod_{project}"]).returncode == 0:
        print(f"Tearing down stack for project {project}.")
        runPodman("compose", ["-f", f"projects/{project}/compose.yaml.rendered", "down"])

    print(f"Creating & starting stack for project {project}.")
    runPodman("compose", ["-f", f"projects/{project}/compose.yaml.rendered", "up", "-d"])


def updateProj(project):
    if not os.path.isfile(f"projects/{project}/compose.yaml.rendered"):
        setupProj(project)
            
    print(f"Running update for project {project}.")
    
    if len(pullProj(project)) > 0:
        backupProj(project)
        upProj(project)


def main():
    envFile = "pyenv.yml"
    secretsFile = "pysecrets.yml"
    
    os.chdir(os.path.realpath(sys.path[0]))

    with open(envFile, 'r') as envfile, open(secretsFile, 'r') as secretsfile:
        global env, secrets
        env = yaml.safe_load(Template(filename=envFile).render())
        secrets = yaml.safe_load(secretsfile)

    setOwner(secretsFile, os.getuid(), os.getgid())
    setPerms(secretsFile, 600)

    print("\nUsing socket " + env['socket'] + ".")


    action = ''
    if len(sys.argv) > 1:
        action = sys.argv[1]
    else:
        print("\nChoose action:")
        print("[1/S] Setup project")
        print("[2/U] Update project")
        print("[3/B] Backup project")

    while action == '':
        action = input("Action: ")


    projects = os.listdir("projects")
    print(f"\nProjects list: {projects}")
    if len(sys.argv) > 2:
        target_projects = sys.argv[2]
    else:
        target_projects = input("Target compose project(s), space separated, leave empty to target all: ")

    if target_projects.strip() == '':
        target_projects = projects
    else:
        target_projects = re.split(' ?, ?| ', target_projects.strip())

    print(f"Target projects: {target_projects}")


    match action.casefold():
        case '1' | 's' | 'setup':
            setNftables()

            for project in target_projects:
                try:
                    print()
                    setupProj(project)
                except Exception as e:
                    print(e, file=sys.stderr)
                    print(f"Failed to setup project {project}.", file=sys.stderr)

        case '2' | 'u' | 'update':
            for project in target_projects:
                try:
                    print()
                    updateProj(project)
                except Exception as e:
                    print(e, file=sys.stderr)
                    print(f"Failed to update project {project}.", file=sys.stderr)

        case '3' | 'b' | 'backup':
            print()
            if not os.path.exists(env['borg_repo']):
                print(f"Creating borg repository {env['borg_repo']}.")
                runBorg(["init", "--encryption", "repokey", env['borg_repo']])

            borgCreate("secrets", path=secretsFile)

            for project in target_projects:
                try:
                    print()
                    backupProj(project)
                except Exception as e:
                    print(e, file=sys.stderr)
                    print(f"Failed to backup project {project}.", file=sys.stderr)

            runBorg(["compact", env['borg_repo']])


if __name__ == "__main__":
    main()