import json
import logging
import subprocess
from typing import Optional

import click
from pydantic import BaseModel
from rich.logging import RichHandler

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


class KeepStrategy(BaseModel):
    keep_last: int = 1
    keep_hourly: Optional[int] = None
    keep_daily: Optional[int] = None
    keep_weekly: Optional[int] = None
    keep_monthly: Optional[int] = None
    keep_yearly: Optional[int] = None

    def arguments(self):
        for field in self.model_fields_set:
            if val := getattr(self, field):
                yield f'--{field.replace("_", "-")}', str(val)


def forget(loki_url: str = None, dry_run: bool = False, **kwargs):
    cmd = ['restic', 'forget', '--json']

    if dry_run:
        cmd.append('--dry-run')

    keep_strat = KeepStrategy.model_validate(kwargs)
    for args in keep_strat.arguments():
        cmd.extend(args)

    logger.debug(f'Running cmd [bright_black]{" ".join(cmd)}[/]')
    with console.status('Forgetting...'):
        result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode == 0:
        try:
            line = result.stdout
            # data = json.loads(line)
            json.loads(line)
        except json.JSONDecodeError:
            logger.error('Error decoding JSON')
            logger.error(line)
        else:
            if loki_url is not None and not dry_run:
                send_to_loki(loki_url=loki_url, line=line, backup='forget')
        finally:
            logger.debug('Done')
    else:
        logger.error('Error running command')
        logger.error(result.stderr)


@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)
@click.option('-l', '--keep-last', type=int, default=None)
@click.option('-H', '--keep-hourly', type=int, default=None)
@click.option('-d', '--keep-daily', type=int, default=None)
@click.option('-w', '--keep-weekly', type=int, default=None)
@click.option('-m', '--keep-monthly', type=int, default=None)
@click.option('-y', '--keep-yearly', type=int, default=None)
def main(loki_url: str, dry_run: bool, **kwargs):
    logging.basicConfig(
        level='DEBUG', format='%(message)s', handlers=[RichHandler(markup=True, console=console)]
    )
    logging.getLogger('docker.utils.config').setLevel('WARNING')
    logging.getLogger('urllib3.connectionpool').setLevel('WARNING')

    try:
        forget(loki_url, dry_run=dry_run, **kwargs)
    except Exception:
        console.print_exception()
    else:
        prune(loki_url, dry_run)
        if not dry_run:
            snapshots.snapshot(loki_url)
            size.get_size(loki_url)
        else:
            logger.debug('Skipping getting size because of dry run')


if __name__ == '__main__':
    main()