import os
from io import StringIO
from typing import cast

import pyinfra
from pyinfra.context import host
from pyinfra.operations import apt, files, server, systemd

TZ = 'America/Los_Angeles'


def sshServer():
    systemd.service(
        service='ssh',
        running=True,
        enabled=True,
    )

    files.line(path='/etc/ssh/ssh_config', line="HashKnownHosts", replace="HashKnownHosts no")

    if 'pi' not in host.groups:
        files.line(path='/etc/ssh/sshd_config', line="^UseDNS\b", replace="UseDNS no")
        # MAYBE plus needs this fix: adding ListenAddress 0.0.0.0 to /etc/ssh/sshd_config
        systemd.service(service='sshd', reloaded=True)


def timezone():
    files.link(path='/etc/localtime', target=f'/usr/share/zoneinfo/{TZ}')
    files.replace(path='/etc/timezone', text='.*', replace=TZ)


def fstab():
    fstab_file = f'system/fstabs/{host.name}'
    if os.path.exists(fstab_file):
        files.put(src=fstab_file, dest='/etc/fstab')


def pi_tmpfs():
    if 'pi' not in host.groups:
        return

    for line in [
            'tmpfs /var/log tmpfs defaults,noatime,mode=0755 0 0',
            'tmpfs /tmp tmpfs defaults,noatime 0 0',
    ]:
        files.line(path="/etc/fstab", line=line, replace=line)

    # stop SD card corruption (along with some mounts in fstab)
    apt.packages(packages=['dphys-swapfile'], present=False)


def no_sleep():
    if host.name not in ['bang', 'pipe', 'ditto']:
        return

    server.shell(commands=['systemctl mask sleep.target suspend.target hibernate.target hybrid-sleep.target'])


def nfs_server():
    if host.name != 'ditto':
        return

    # remove when we're on longhorn
    apt.packages(packages=['nfs-kernel-server'])
    files.put(src='system/files/ditto_exports', dest='/etc/exports')


def smaller_journals():
    if host.name not in ['prime', 'ditto', 'pipe']:
        return
    files.line(name='shorter systemctl log window, for disk space',
               path='/etc/systemd/journald.conf',
               line='MaxFileSec',
               replace="MaxFileSec=7day")


def web_forward():
    if host.name != 'prime':
        return
    for port in [80, 443]:
        svc = f'web_forward_{port}'
        files.template(src="system/templates/webforward.service.j2",
                       dest=f"/etc/systemd/system/{svc}.service",
                       serv_host='bang',
                       port=port,
                       name='web',
                       fam='tcp')
        systemd.service(service=svc, enabled=True, restarted=True)


def minecraft_forward():
    if host.name != 'prime':
        return
    port = 25765
    for fam in ['tcp', 'udp']:
        svc = f'mc_smp_{fam}_forward_{port}'
        files.template(src="system/templates/webforward.service.j2",
                       dest=f"/etc/systemd/system/{svc}.service",
                       serv_host='ditto',
                       port=port,
                       name='mc_smp',
                       fam=fam)
        systemd.service(service=svc, enabled=True, restarted=True)


def pigpiod():
    if 'pi' not in host.groups:
        return
    files.put(src="system/files/pigpiod.service", dest="/etc/systemd/system/pigpiod.service")
    systemd.service(service='pigpiod', daemon_reload=True, enabled=True)


def rpi_iscsi_volumes():
    if host.name != 'ditto':
        return

    iscsi_dir = '/d2/rpi-iscsi'

    try:
        pis=pyinfra.inventory.get_group(name='pi')
    except Exception as e:
        print(f"no pi group, skipping {e!r}")
        return
    for pi_hostname in cast(list, pis):
        out = f'{iscsi_dir}/{pi_hostname}.disk'
        files.directory(path=iscsi_dir)
        server.shell(commands=f'dd if=/dev/zero of={out} count=0 bs=1 seek=10G conv=excl || true')
        files.put(dest=f"/etc/tgt/conf.d/{pi_hostname}.conf",
                  src=StringIO(f"""
<target iqn.2024-03.com.bigasterisk:{pi_hostname}.target>
    backing-store {out}
    initiator-name iqn.2024-03.com.bigasterisk:{pi_hostname}.initiator
</target>
                            """))
    # restarting is disruptive to connected pis, and they might need to be
    # visited:
    #systemd.service(service='tgt.service', running=True, restarted=True)


def hostname():
    server.hostname(hostname=host.name)



operations = [
    hostname,
    timezone,
    fstab,
    rpi_iscsi_volumes,
    pi_tmpfs,
    no_sleep,
    nfs_server,
    smaller_journals,
    web_forward,
    minecraft_forward,
    pigpiod,
]
# for space, consider:
# k3s crictl rmi --prune
# snap list --all | while read snapname ver rev trk pub notes; do if [[ $notes = *disabled* ]]; then snap remove "$snapname" --revision="$rev"; fi; done
# podman system reset