from datetime import date
from time import sleep
from typing import Optional, List

import boto3
import typer
import config
from owslogger import logger
from neo4j import GraphDatabase
import sentry_sdk


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

app = typer.Typer(pretty_exceptions_show_locals=False)


def refresh_qa_backup(prefix: str) -> Optional[str]:
    """Compatibility wrapper to fetch latest backup URL for a prefix."""
    backup_url = get_latest_backup(prefix)
    if backup_url:
        refresh_neo4j_backup(backup_url)
    else:
        log.error('No matching backup found.')
        raise typer.Exit(code=1)


def refresh_neo4j_backup(
    backup_url: str,
    neo4j_url: Optional[str] = None,
    username: Optional[str] = None,
    password: Optional[str] = None,
    new_database_name: Optional[str] = None,
    apply: bool = True,
) -> Optional[str]:
    """Recreate the graph database from backup on a Neo4j 5.26.x server.

    - Connects to server
    - Lists databases starting with 'graph.'
    - Finds which database alias 'graph.db' points to
    - Plans which 'graph.*' databases to drop (never drops alias or a database named 'graph.db')
    - If apply=True, drops planned databases
    - Determines the new database name from backup URL (or provided)
    - If apply=True, creates a new database with topology 3 PRIMARIES from the provided S3 backup URL

    Returns the name of the created database.
    """
    effective_url = neo4j_url or config.NEO4J_REFRESH_URL
    effective_user = username or config.NEO4J_REFRESH_CONNECTION_USER
    effective_pass = password or config.NEO4J_REFRESH_CONNECTION_PASSWORD

    try:
        driver = GraphDatabase.driver(effective_url, auth=(effective_user, effective_pass))
        with driver.session(database='system') as session:
            alias_target = _get_graph_db_alias_target(session)

            graph_dbs = _list_graph_databases(session)
            log.info(f'graph.* databases discovered: {graph_dbs}')

            names_to_drop: List[str] = []
            for db_name in graph_dbs:
                if db_name == 'graph.db':
                    log.info("Skipping drop of protected database name 'graph.db'")
                    continue
                if alias_target and db_name == alias_target:
                    log.info(f"Skipping drop of '{db_name}' because alias 'graph.db' points to it")
                    continue
                names_to_drop.append(db_name)

            if names_to_drop:
                log.info(f'Databases planned to drop: {names_to_drop}')
            else:
                log.info('No databases to drop')

            if apply:
                for db_name in names_to_drop:
                    _drop_database(session, db_name)
                    log.info(f"Dropped database '{db_name}', sleeping for 10 seconds")
                    sleep(10)
            else:
                log.info('apply=False, skipping DROP operations')

            created_db = new_database_name or _derive_database_name_from_backup_url(backup_url)
            if created_db == 'graph.db':
                created_db = f'graph.db-{date.today()}'

            if apply:
                _create_database_from_backup(session, created_db, backup_url)
                log.info(f"Created database '{created_db}' from {backup_url}")
            else:
                log.info(f"apply=False, would create database '{created_db}' from {backup_url}")
            return created_db
    except Exception as error:
        sentry_sdk.capture_exception(error)
        log.error(f'Failed to refresh Neo4j from backup: {error}')
        raise


def get_latest_backup(
    prefix: str,
    bucket: Optional[str] = None,
    region: Optional[str] = None,
    suffix: str = '.backup',
) -> Optional[str]:
    """Return full s3:// URL of the most recently modified backup file in prefix.

    Args:
        prefix: S3 key prefix to search under.
        bucket: Optional S3 bucket override. Defaults to configured bucket.
        region: Optional AWS region override. Defaults to configured region.
        suffix: File suffix to filter on. Defaults to ".backup".

    Returns:
        Full s3:// URL string to the latest backup, or None if none found.
    """
    used_bucket = bucket or config.NEO4J_BACKUP_S3_BUCKET
    used_region = region or config.AWS_REGION

    log.info(f"Looking for latest backup in s3://{used_bucket}/{prefix} with suffix '{suffix}'")

    s3_client = boto3.client('s3', region_name=used_region)
    paginator = s3_client.get_paginator('list_objects_v2')

    latest_backup_obj = None

    for page in paginator.paginate(Bucket=used_bucket, Prefix=prefix):
        for obj in page.get('Contents', []):
            key = obj.get('Key')
            if not key or not key.endswith(suffix):
                continue
            if latest_backup_obj is None or obj['LastModified'] > latest_backup_obj['LastModified']:
                latest_backup_obj = obj

    if not latest_backup_obj:
        log.warning(f'No {suffix} files found under s3://{used_bucket}/{prefix}')
        return None

    latest_key = latest_backup_obj['Key']
    full_url = f's3://{used_bucket}/{latest_key}'
    log.info(
        'Selected latest backup: '
        f'{full_url} '
        f'(LastModified={latest_backup_obj["LastModified"]}, Size={latest_backup_obj.get("Size")})'
    )
    return full_url


def _get_graph_db_alias_target(session) -> Optional[str]:
    """Return the database name targeted by alias 'graph.db', if it exists."""
    records = session.run('SHOW ALIASES FOR DATABASE').data()
    for record in records:
        alias_name = record.get('name') or record.get('alias')
        target = record.get('database') or record.get('target')
        if alias_name == 'graph.db' and isinstance(target, str) and target:
            log.info(f"Graph DB alias 'graph.db' points to '{target}'")
            return target
    log.info("No graph DB alias 'graph.db' found")
    return None


def _list_graph_databases(session, show_starting: bool = False) -> List[str]:
    """Return names of databases starting with 'graph.'

    When show_starting is False (default), databases whose currentStatus is
    'starting' are excluded. Results are de-duplicated while preserving order.
    """
    records = session.run('SHOW DATABASES').data()
    names: List[str] = []
    seen: set[str] = set()
    for record in records:
        name = record.get('name')
        status = record.get('currentStatus')
        if not (isinstance(name, str) and name.startswith('graph')):
            continue
        if not show_starting and status == 'starting':
            continue
        if name not in seen:
            names.append(name)
            seen.add(name)
    log.info(f'Graph DB databases discovered: {names}')
    return names


def _drop_database(session, database_name: str) -> None:
    """Drop a database by name in the system database (safe quoting)."""
    if database_name in {'graph.db', 'system'}:
        log.warning(f"Refusing to drop protected database '{database_name}'")
        return
    escaped = database_name.replace('`', '``')
    cypher = f'DROP DATABASE `{escaped}` IF EXISTS'
    log.info(f"Dropping database '{database_name}'")
    session.run(cypher)


def _create_database_from_backup(session, database_name: str, backup_url: str) -> None:
    """Create a database with 3 primaries from the given backup URL."""
    escaped_db = database_name.replace('`', '``')
    escaped_url = backup_url.replace("'", "''")
    cypher = (
        f'CREATE DATABASE `{escaped_db}` TOPOLOGY 3 PRIMARIES OPTIONS '
        f"{{ existingData: 'use', seedURI:'{escaped_url}' }}"
    )
    log.info(f"Creating database '{database_name}' from backup URL")
    session.run(cypher)


def _derive_database_name_from_backup_url(backup_url: str) -> str:
    """Derive a database name from the backup file name, prefixed with 'graph.' if needed."""
    # Strip s3://bucket/ and take final path segment
    path_part = backup_url.split('//', 1)[-1]
    path_part = path_part.split('/', 1)[-1] if '/' in path_part else path_part
    filename = path_part.split('/')[-1]
    base = filename.rsplit('.', 1)[0] if '.' in filename else filename
    # Sanitize and ensure starts with graph.
    safe = ''.join(ch if ch.isalnum() or ch in ('-', '_', '.') else '-' for ch in base)
    if not safe.startswith('graph.'):
        safe = f'graph.{safe}'
    return safe


@app.command()
def main(
    prefix: str,
    bucket: Optional[str] = typer.Option(config.NEO4J_BACKUP_S3_BUCKET, help='S3 bucket name override'),
    region: Optional[str] = typer.Option(config.AWS_REGION, help='AWS region override'),
    suffix: str = typer.Option('.backup', help='File suffix filter (default: .backup)'),
    neo4j_url: Optional[str] = typer.Option(config.NEO4J_REFRESH_URL, help='Neo4j bolt+s URL override'),
    username: Optional[str] = typer.Option(config.NEO4J_REFRESH_CONNECTION_USER, help='Neo4j username override'),
    password: Optional[str] = typer.Option(config.NEO4J_REFRESH_CONNECTION_PASSWORD, help='Neo4j password override'),
    new_database_name: Optional[str] = typer.Option(None, help='Name for the new database'),
    apply: bool = typer.Option(False, help='If set, perform the refresh on Neo4j'),
):
    url = get_latest_backup(prefix=prefix, bucket=bucket, region=region, suffix=suffix)
    if not url:
        log.error('No matching backup found.')
        raise typer.Exit(code=1)
    created = refresh_neo4j_backup(
        backup_url=url,
        neo4j_url=neo4j_url,
        username=username,
        password=password,
        new_database_name=new_database_name,
        apply=apply,
    )
    typer.echo(created or url)


if __name__ == '__main__':
    app()