from datetime import date
import glob
import json
import os.path
import subprocess
import time

import config
from refresh_neo4j import refresh_qa_backup
from lambdacommon import util
from neo4j import GraphDatabase
from owslogger import logger
import sentry_sdk


# Global logger
log = logger.setup(
    config.ENVIRONMENT,
    config.LOGGER_NAME,
    config.LOGGER_LEVEL,
    config.SERVICE_NAME,
    config.SERVICE_VERSION,
    dsn=config.LOGGER_DSN,
)


def main():
    """Execute main entrypoint."""
    call_datadog_with_metric(config.DATADOG_BACKUP_ATTEMPT_METRIC)

    if config.NEO4J_BACKUP_FROM_FOLLOWER:
        host = get_follower_node_host()
    else:
        host = config.NEO4J_HOST

    log.info(f'Attempting to perform a backup on {host}')

    prefix = f'{config.NEO4J_BACKUP_S3_PREFIX}/{date.today()}'

    backup_directory = run_backup_process(
        host=host,
        backup_dir=config.NEO4J_BACKUP_DIRECTORY,
        compress=config.NEO4J_BACKUP_COMPRESS_BACKUP,
        pagecache=config.NEO4J_BACKUP_PAGECACHE,
        timeout=config.NEO4J_BACKUP_TIMEOUT,
        database=config.NEO4J_BACKUP_DATABASE,
        include_metadata=config.NEO4J_BACKUP_INCLUDE_METADATA,
        failure_metric=config.DATADOG_BACKUP_FAILURE_METRIC,
        port=config.NEO4J_BACKUP_PORT,
    )

    # Best-effort collection of backup metadata (highest committed tx id) for Aura import.
    # A failure here must never fail an otherwise successful backup, so it returns None.
    backup_metadata = inspect_backup_metadata(
        backup_dir=config.NEO4J_BACKUP_DIRECTORY,
        timeout=config.NEO4J_BACKUP_INSPECT_TIMEOUT,
    )

    if config.NEO4J_BACKUP_UPLOAD_TO_S3:
        upload_backup_to_s3(
            backup_dir=backup_directory,
            bucket=config.NEO4J_BACKUP_S3_BUCKET,
            prefix=prefix,
            kms_key_id=config.NEO4J_BACKUP_KMS_KEY_ID,
            concurrency=config.NEO4J_BACKUP_UPLOAD_CONCURRENCY,
            timeout=config.NEO4J_BACKUP_UPLOAD_TIMEOUT,
            failure_metric=config.DATADOG_BACKUP_UPLOAD_FAILURE_METRIC,
            metadata=backup_metadata,
        )
    else:
        log.info('NEO4J_BACKUP_UPLOAD_TO_S3 set to false, skipping upload')

    if config.NEO4J_BACKUP_REFRESH_QA:
        refresh_qa_backup(prefix)
    call_datadog_with_metric(config.DATADOG_BACKUP_SUCCESS_METRIC)
    log.info('Backup process completed successfully')
    return True


def run_backup_process(
    host,
    backup_dir,
    compress,
    failure_metric,
    pagecache,
    timeout,
    port,
    database,
    include_metadata,
    additional_config='/app/additional.conf',
    backup_type='FULL',
):
    """Execute backup process."""
    backup_command_format_v4 = [
        'neo4j-admin',
        'backup',
        f'--backup-dir={backup_dir}',
        f'--from={host}:{port}',
        f'--database={database}',
        f'--pagecache={pagecache}',
        '--check-consistency=false',
        f'--include-metadata={include_metadata}',
    ]

    backup_command_format_v5 = [
        'neo4j-admin',
        'database',
        'backup',
        f'--additional-config={additional_config}',
        f'--compress={compress}',
        f'--to-path={backup_dir}',
        f'--from={host}:{port}',
        f'--pagecache={pagecache}',
        f'--include-metadata={include_metadata}',
        f'--type={backup_type}',
        f'{database}',
    ]

    if config.NEO4J_SERVER_MAJOR_VERSION == 4:
        backup_command_format = backup_command_format_v4
    elif config.NEO4J_SERVER_MAJOR_VERSION == 5:
        backup_command_format = backup_command_format_v5
    else:
        raise RuntimeError('invalid value for NEO4J_SERVER_MAJOR_VERSION')

    clean_backup_dir(backup_dir, database)
    if config.DRY_RUN:
        print(f'DRY_RUN would run: {" ".join(backup_command_format)}')
        return f'{backup_dir}{database}'
    else:
        try:
            result = subprocess.run(
                backup_command_format,
                universal_newlines=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=timeout,
                check=True,
                env={
                    'HEAP_SIZE': config.NEO4J_BACKUP_HEAP_SIZE,
                },
            )
            log.info(f'Backup process stdout is: {result.stdout}')
            result.returncode != 0 and log.error(f'Backup process stderr is {result.stderr}')
            return f'{backup_dir}{database}'
        except subprocess.CalledProcessError as error:
            sentry_sdk.capture_exception(error)
            log.error(f'Backup subprocess call experienced an error: stdout: {error.stdout}, stderr: {error.stderr}')
            call_datadog_with_metric(failure_metric)
            raise error


def _find_metadata_value(record, *candidates):
    """Look up a value in a metadata record, ignoring key case and separators."""
    normalized = {key.lower().replace(' ', '').replace('_', ''): value for key, value in record.items()}
    for candidate in candidates:
        value = normalized.get(candidate.lower().replace(' ', '').replace('_', ''))
        if value is not None:
            return value
    return None


def inspect_backup_metadata(backup_dir, timeout):
    """Inspect the latest backup and return its metadata (incl. lastTxId).

    Runs `neo4j-admin backup inspect --latest-backup --show-metadata --format=JSON` over
    the backup directory, parses stdout JSON, and returns a normalized dict that always
    includes a canonical 'lastTxId' (the highestTransaction of the latest backup) for
    Aura import. `--show-metadata` is required — without it inspect only reports the file
    uri, not the transaction ids.

    This is best-effort: it returns None on the v4 branch, on DRY_RUN, or whenever the
    metadata cannot be obtained, so it never blocks the backup/upload success path.
    """
    if config.NEO4J_SERVER_MAJOR_VERSION == 4:
        log.warning('Backup inspect metadata not collected for v4; skipping')
        return None

    inspect_command_format_v5 = [
        'neo4j-admin',
        'backup',
        'inspect',
        '--latest-backup',
        '--show-metadata',
        '--format=JSON',
        backup_dir,
    ]

    if config.DRY_RUN:
        print(f'DRY_RUN would run: {" ".join(inspect_command_format_v5)}')
        return None

    try:
        result = subprocess.run(
            inspect_command_format_v5,
            universal_newlines=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout,
            check=True,
        )
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as error:
        sentry_sdk.capture_exception(error)
        log.error(f'Backup inspect subprocess call experienced an error: {error}')
        return None

    try:
        parsed = json.loads(result.stdout)
    except (json.JSONDecodeError, ValueError) as error:
        sentry_sdk.capture_exception(error)
        log.error(f'Could not parse backup inspect JSON: {error}; raw stdout: {result.stdout}')
        return None

    records = parsed if isinstance(parsed, list) else [parsed]
    if not records or not isinstance(records[0], dict):
        log.error(f'Backup inspect returned no usable records: {result.stdout}')
        return None
    record = records[0]

    highest_tx = _find_metadata_value(record, 'highestTransaction', 'highestTx', 'highestTxId', 'lastTxId')
    if highest_tx is None:
        log.error(f'highestTransaction not found in inspect record keys: {list(record.keys())}')
        return None

    metadata = dict(record)
    metadata['lastTxId'] = highest_tx
    log.info(f'Collected backup metadata: lastTxId={highest_tx}')
    return metadata


def upload_backup_to_s3(
    backup_dir,
    bucket,
    prefix,
    kms_key_id,
    concurrency,
    timeout,
    failure_metric,
    metadata=None,
):
    """Upload backup to s3."""
    upload_command_format_v4 = [
        's5cmd',
        'sync',
        '-c',
        concurrency,
        '--sse',
        'aws:kms',
        '--sse-kms-key-id',
        kms_key_id,
        backup_dir,
        f's3://{bucket}/{prefix}/',
    ]

    upload_command_format_v5 = [
        's5cmd',
        'sync',
        '-c',
        concurrency,
        '--sse',
        'aws:kms',
        '--sse-kms-key-id',
        kms_key_id,
        f'{backup_dir}*.backup',
        f's3://{bucket}/{prefix}/',
    ]

    if config.NEO4J_SERVER_MAJOR_VERSION == 4:
        upload_command_format = upload_command_format_v4
    elif config.NEO4J_SERVER_MAJOR_VERSION == 5:
        upload_command_format = upload_command_format_v5
    else:
        raise RuntimeError('invalid value for NEO4J_SERVER_MAJOR_VERSION')

    if config.DRY_RUN:
        print(f'DRY_RUN would run: {" ".join(upload_command_format)}')
    else:
        try:
            result = subprocess.run(
                upload_command_format,
                universal_newlines=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=timeout,
                check=True,
            )
            log.info(f'Upload process stdout is: {result.stdout}')

            result.returncode != 0 and log.error(f'Upload process stderr is {result.stderr}')
        except subprocess.CalledProcessError as error:
            sentry_sdk.capture_exception(error)
            log.error(f'Upload subprocess call experienced an error: stdout: {error.stdout}, stderr: {error.stderr}')
            call_datadog_with_metric(failure_metric)
            raise error

    with open('/tmp/backup-complete.txt', 'w') as backup_complete_file:
        backup_complete_file.write(f'prefix: {prefix} upload complete')

    upload_command_format_complete = [
        's5cmd',
        'sync',
        '/tmp/backup-complete.txt',
        f's3://{bucket}/{prefix}/',
    ]

    if config.DRY_RUN:
        print(f'DRY_RUN would run: {" ".join(upload_command_format_complete)}')
    else:
        try:
            result = subprocess.run(
                upload_command_format_complete,
                universal_newlines=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=timeout,
                check=True,
            )
            log.info(f'Upload backup-complete.txt stdout is: {result.stdout}')

        except subprocess.CalledProcessError as error:
            sentry_sdk.capture_exception(error)
            log.error(
                'Upload backup-complete.txt subprocess call experienced an error: '
                f'stdout: {error.stdout}, stderr: {error.stderr}'
            )
            call_datadog_with_metric(failure_metric)
            raise error

    if metadata is not None:
        metadata_path = '/tmp/backup-metadata.json'
        # Best-effort: the durable .backup artifacts are already uploaded, so a failure to
        # write the sidecar must not fail the run (do not raise) — log and skip the upload.
        try:
            with open(metadata_path, 'w') as metadata_file:
                json.dump(metadata, metadata_file, indent=2, default=str)
        except (OSError, TypeError, ValueError) as error:
            sentry_sdk.capture_exception(error)
            log.error(f'Failed to write backup-metadata.json to {metadata_path}: {error}')
            return

        upload_command_format_metadata = [
            's5cmd',
            'sync',
            metadata_path,
            f's3://{bucket}/{prefix}/',
        ]

        if config.DRY_RUN:
            print(f'DRY_RUN would run: {" ".join(upload_command_format_metadata)}')
        else:
            # Best-effort: the durable .backup artifacts are already uploaded, so a
            # metadata sidecar failure must not fail the run (do not raise).
            try:
                result = subprocess.run(
                    upload_command_format_metadata,
                    universal_newlines=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    timeout=timeout,
                    check=True,
                )
                log.info(f'Upload backup-metadata.json stdout is: {result.stdout}')
            except subprocess.CalledProcessError as error:
                sentry_sdk.capture_exception(error)
                log.error(
                    'Upload backup-metadata.json subprocess call experienced an error: '
                    f'stdout: {error.stdout}, stderr: {error.stderr}'
                )


def clean_backup_dir(backup_dir, database):
    """Clean backup directory."""
    if not backup_dir:
        raise Exception(
            f'Will not clean backup dir where either backup_dir({backup_dir}) or database({database}) is empty'
        )
    if not config.NEO4J_BACKUP_DIRECTORY_CLEAN:
        log.info('Not cleaning backup dir as NEO4J_BACKUP_DIRECTORY_CLEAN is false')
        return True

    # Find all files that start with the database name and end with .backup
    pattern = os.path.join(backup_dir, f'{database}*.backup')
    backup_files = glob.glob(pattern)

    if backup_files:
        log.info(f'Found {len(backup_files)} backup files matching pattern: {pattern}')
        for backup_file in backup_files:
            try:
                log.info(f'Attempting to delete backup file: {backup_file}')
                os.remove(backup_file)
                log.info(f'Successfully deleted: {backup_file}')
            except Exception as e:
                log.error(f'Unable to delete backup file {backup_file}. Exception: {e}')
                raise e
    else:
        log.info(f'No backup files found matching pattern: {pattern}')


def call_datadog_with_metric(metric):
    """
    Wrap the DataDog metric call.

    Args:
        metric (str): custom metric name to send to DataDog.
    """
    with util.datadog_connection(api_key=config.DATADOG_API_KEY, app_key=config.DATADOG_APP_KEY) as datadog:
        now = time.time()
        datadog.Metric.send(
            [
                {
                    'metric': '{}.{}'.format(config.SERVICE_NAME, metric),
                    'type': 'count',
                    'interval': 60,
                    'points': (now, 1),
                    'tags': [
                        'environment:{}'.format(config.ENVIRONMENT),
                        'service_name:{}'.format(config.SERVICE_NAME),
                    ],
                }
            ]
        )


def get_follower_node_host():
    """Get the host of the follower node."""
    log.info('Trying to get follower node host...')
    driver = GraphDatabase.driver(
        config.NEO4J_URL,
        auth=(config.NEO4j_USERNAME, config.NEO4j_PASSWORD),
    )
    database_info = "show database graph.db where role = 'follower';"
    with driver.session() as source:
        follower_nodes = source.run(database_info).data()
        host, port = follower_nodes[0]['address'].split(':')
        return host


if __name__ == '__main__':
    main()