import json
import logging
import re
import subprocess
import sys

import click
from rich.logging import RichHandler

from restic import size, snapshots
from restic.console import console, logger
from restic.loki import send_to_loki

field_regex = re.compile(
    r'^(?P<name>[\w ]+):\s+(?P<blobs>\d+) blobs \/ (?P<size>\d+(\.\d+)? (?:Mi|Gi)?B)', re.MULTILINE
)


def convert_size(size_str: str) -> int:
    base_size = float(size_str.split(' ')[0])

    if size_str.endswith('GiB'):
        scale = 8589934592
    elif size_str.endswith('MiB'):
        scale = 8388608
    else:
        scale = 1
    return int(round(base_size * scale))


def field_gen(result_str: str):
    for m in field_regex.finditer(result_str):
        d = m.groupdict()
        d['blobs'] = int(d['blobs'])
        d['size'] = convert_size(d['size'])
        yield d


def prune(loki_url: str = None, dry_run: bool = False):
    cmd = ['restic', 'prune']
    if dry_run:
        cmd.append('--dry-run')
    logger.debug(f'Running cmd [bright_black]{" ".join(cmd)}[/]')

    with console.status('Pruning...'):
        result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        logger.error(result.stderr)
        sys.exit(1)

    d = {f['name']: {'blobs': f['blobs'], 'size': f['size']} for f in field_gen(result.stdout)}
    logger.debug(json.dumps(d, indent=4))

    if loki_url is not None and not dry_run:
        send_to_loki(loki_url, line=json.dumps(d), backup='prune')


@click.command()
@click.option(
    '--loki-url',
    type=str,
    help='Loki URL for logging. Defaults to the LOKI_URL env variable',
    envvar='LOKI_URL',
    default=None,
)
@click.option('-n', '--dry-run', type=bool, default=False, is_flag=True)
def main(loki_url: str, dry_run: bool):
    logging.basicConfig(
        level='DEBUG', format='%(message)s', handlers=[RichHandler(markup=True, console=console)]
    )

    try:
        prune(loki_url=loki_url, dry_run=dry_run)
    except Exception as e:
        raise e
    else:
        snapshots.snapshot(loki_url)
        size.get_size(loki_url)


if __name__ == '__main__':
    main()