#!/usr/bin/env python3

import os
import sys
import glob
import shlex
import subprocess
import logging
import dateutil.parser
import pytz
import time
from datetime import datetime, timedelta
from typing import Dict, List

logging.Formatter.converter = time.gmtime
logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s")
logger = logging.getLogger(__name__)

def run(args: List[str], dry_run: bool = False) -> str:
    if dry_run:
        print("Would have run: " + " ".join(map(shlex.quote, args)))
        return ""

    p = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE)
    stdout, stderr = p.communicate()
    if p.returncode:
        logger.error("Could not invoke %s\nstdout: %r\nstderror: %r"
                     % (args[0], stdout, stderr))
        sys.exit(1)
    return stdout.decode()

recovery_val = run(['psql', '-v', 'ON_ERROR_STOP=1', '-t', '-c', 'SELECT pg_is_in_recovery()']).strip()
# Assertion to check that we're extracting the value correctly.
assert recovery_val in ['t', 'f']
if recovery_val == 't':
    # Only run if we're the master
    sys.exit(0)

is_rhel_based = os.path.exists("/etc/redhat-release")
if is_rhel_based:
    pg_data_paths = glob.glob('/var/lib/pgsql/*/data')
else:
    pg_data_paths = glob.glob('/var/lib/postgresql/*/main')
if len(pg_data_paths) != 1:
    print("Postgres installation is not unique: %s" % (pg_data_paths,))
    sys.exit(1)
pg_data_path = pg_data_paths[0]
run(['env-wal-e', 'backup-push', pg_data_path])

now = datetime.now(tz=pytz.utc)
with open('/var/lib/nagios_state/last_postgres_backup', 'w') as f:
    f.write(now.isoformat())
    f.write("\n")

backups = {}  # type: Dict[datetime, str]
lines = run(['env-wal-e', 'backup-list']).split("\n")
for line in lines[1:]:
    if line:
        backup_name, date_str, _, _ = line.split()
        backups[dateutil.parser.parse(date_str)] = backup_name

one_month_ago = now - timedelta(days=30)
for date in sorted(backups.keys(), reverse=True):
    if date < one_month_ago:
        run(['env-wal-e', 'delete', '--confirm', 'before', backups[date]])
        # Because we're going from most recent to least recent, we
        # only have to do one delete operation
        break
