import csv
import gzip
import os
import time
import timeit
import traceback
from datetime import timedelta
from typing import List

import click
import sqlalchemy
from sqlalchemy import Table
from sqlalchemy.exc import DatabaseError
from structlog import get_logger

from delphi_api.v3.data_models.postgres_db import (
    Artist,
    BaseModel,
    Playlist,
    PlaylistRank,
    Product,
    Session,
    Track,
    artist_product_table,
    artist_track_table,
)

LOG = get_logger(__name__)

# Default config
MAX_ROWS_PER_BATCH = 25000
BETA_DATASET = 'beta'
BETA_DATASET_SUBDIR = 'beta_only'
RANDOM_DATASET = 'random'
RANDOM_DATASET_SUBDIR = None
# Test datasets have a smaller subset of artists and tracks for faster test runs
TEST_DATASET = 'test'
TEST_DATASET_SUBDIR = 'test'
SUPPORTED_DATASETS = {BETA_DATASET, TEST_DATASET}

DEFAULT_DATASET_CONFIG = [
    # DSP and Region are now in liquibase/seeds
    {
        'model': Artist,
        'filename': 'artist.csv.gz',
    },
    {
        'model': Product,
        'filename': 'product.csv.gz',
    },
    {
        'model': Track,
        'filename': 'track.csv.gz',
    },
    {
        'model': Playlist,
        'filename': 'playlist.csv.gz',
    },
    {
        'model': PlaylistRank,
        'filename': 'playlist_rank.csv.gz',
    },
    {
        'model': artist_track_table,
        'filename': 'artist_track.csv.gz',
    },
    {
        'model': artist_product_table,
        'filename': 'artist_product.csv.gz',
    },
]

# assumes same file naming conventions
BETA_DATASET_CONFIG = DEFAULT_DATASET_CONFIG
RANDOM_DATASET_CONFIG = DEFAULT_DATASET_CONFIG
TEST_DATASET_CONFIG = DEFAULT_DATASET_CONFIG

SUBDIRS = {
    BETA_DATASET: BETA_DATASET_SUBDIR,
    TEST_DATASET: TEST_DATASET_SUBDIR,
    RANDOM_DATASET: None,
}

CONFIG = {
    BETA_DATASET: BETA_DATASET_CONFIG,
    TEST_DATASET: TEST_DATASET_CONFIG,
    RANDOM_DATASET: RANDOM_DATASET_CONFIG,
}

_divider = '--' * 40


class DataSeeder:
    def __init__(self, config: List[dict],
                 data_subdir: str = '',
                 batch_size: int = MAX_ROWS_PER_BATCH,
                 update_existing: bool = False,
                 truncate: bool = False,
                 disable_constraint_checks: bool = False):
        self.config = config
        self.batch_size = batch_size
        self.truncate = truncate
        self.update_existing = update_existing if not truncate else False
        self.operation_type = 'update' if self.update_existing else 'insert'
        self.disable_constraint_checks = disable_constraint_checks
        if self.truncate:
            self.disable_constraint_checks = True
        self.curr_dir = os.path.dirname(os.path.abspath(__file__))
        self.data_dir = os.path.join(self.curr_dir, 'data', data_subdir)
        self.session = Session()
        self.start = timeit.default_timer()

    @staticmethod
    def _get_model_name(model):
        """Helper for logging which objects are currently being seeded"""
        if hasattr(model, '__name__'):
            return model.__name__
        elif hasattr(model, 'name'):
            return model.name
        else:
            return model.__class__

    @staticmethod
    def _calc_duration(start, end):
        return timedelta(seconds=(end - start)).seconds

    def _wait_for_db_ready(self, timeout=10, sleep_for=1):
        """Wait for database engine ready before attempting executions (timeout in seconds)"""
        click.echo(_divider)
        click.secho('Checking database ready...')
        remaining = timeout
        while remaining > 0:
            try:
                self.session.execute('SELECT 1')
                click.secho(' ✓ Database is ready.')
                click.echo(_divider)
                return True
            except (DatabaseError, Exception) as e:
                remaining -= sleep_for
                click.echo(f'Database not ready. Waiting {remaining} more seconds...')
                time.sleep(sleep_for)
        click.echo(f' ⓧ Database did not become ready in {timeout} seconds. Quitting', err=True)
        return False

    def _disable_fk_checks(self):
        if self.disable_constraint_checks:
            try:
                self.session.execute(sqlalchemy.text("SET session_replication_role = 'replica';"))
            except Exception as e:
                click.echo(
                    f'Failed to disable foreign key checks: \n{str(e)}'
                    '\nAttempting to continue anyway...', err=True)
                return

    def _enable_fk_checks(self):
        if self.disable_constraint_checks:
            try:
                self.session.execute(sqlalchemy.text("SET session_replication_role = 'origin';"))
            except Exception as e:
                click.echo(
                    f'Failed to re-enable foreign key checks: \n{str(e)}'
                    '\nAttempting to continue anyway...', err=True)
                return

    def _data_exists(self, model: BaseModel):
        """Check if data exists for preventing additional INSERT operations"""
        return bool(self.session.query(model).first())

    def _read_data_rows(self, filename: str) -> List[dict]:
        """Reads data from a GZIP encoded CSV file"""
        file_path = os.path.join(self.data_dir, filename)
        rows = []
        with gzip.open(file_path, 'rt') as csvfile:
            reader = csv.DictReader(csvfile)
            for row in reader:
                datum = dict(row)
                if 'ignore' in datum:
                    datum.pop('ignore')
                for k, v in datum.items():
                    # set null for empty strings
                    if v == '':
                        datum[k] = None
                    if v in {'TRUE', 'true'}:
                        datum[k] = True
                    elif v in {'FALSE', 'false'}:
                        datum[k] = False
                rows.append(datum)
        return rows

    def _write_data_rows(self, model: BaseModel, rows: List[dict]):
        """Writes a list of ``rows`` to provided ``model`` class"""
        model_name = self._get_model_name(model)
        if not rows:
            click.secho('No data rows passed. Aborting.')
            return
        args = (f' ⊙ {model_name}',
                f'– Running bulk {self.operation_type} operation for {model_name}')
        click.secho('{: <20} {: >20}'.format(*args))
        remaining = len(rows)
        offset = 0
        limit = self.batch_size
        while remaining > 0:
            limit_offset = limit + offset
            if limit_offset > len(rows):
                limit_offset = len(rows)
            partial_rows = rows[offset:limit_offset]
            args = (f' ⊙ {model_name}',
                    f'– Inserting batch slice {offset}:{limit_offset} of {remaining} remaining')
            click.secho('{: <20} {: >20}'.format(*args))
            try:
                if self.update_existing:
                    self.session.bulk_update_mappings(model, partial_rows)
                else:
                    self.session.bulk_insert_mappings(model, partial_rows,
                                                      return_defaults=False, render_nulls=True)
            except AttributeError:
                model: Table
                self.session.execute(model.insert(), rows[offset:limit_offset])
            offset += self.batch_size
            remaining -= self.batch_size
        args = (f' ✓ {model_name}', f'– Completed bulk {self.operation_type} operation.')
        click.secho('{: <20} {: >20}'.format(*args))
        duration = self._calc_duration(self.start, timeit.default_timer())
        click.secho(f'Elapsed: {duration} seconds')

    def run(self):
        """Execute the database seeding operation"""
        if not self._wait_for_db_ready():
            return
        click.secho(f'Running Data Seeder')
        self.start = timeit.default_timer()
        self._disable_fk_checks()
        success = False
        try:
            obj: dict
            for obj in self.config:
                model: BaseModel = obj['model']
                filename = obj['filename']
                model_name = self._get_model_name(model)
                try:
                    table_name = model.__tablename__
                except AttributeError:
                    table_name = model.name
                if self.truncate:
                    self.session.execute(f'TRUNCATE TABLE {table_name} CASCADE;')
                if self._data_exists(model):
                    if not self.update_existing:
                        args = (f' ⊙ {model_name}', '– Skipping seeding for model. Data exists.')
                        click.secho('{: <20} {: >20}'.format(*args))
                        continue
                    else:
                        args = (f' ⊙ {model_name}', '– Data exists in table. Updating existing.')
                        click.secho('{: <20} {: >20}'.format(*args))
                rows = self._read_data_rows(filename)

                args = (f' ⊙ {model_name}', f'– Read {len(rows)} rows from {filename}')
                click.secho('{: <20} {: >20}'.format(*args))
                self._write_data_rows(model=model, rows=rows)
            self.session.commit()
            success = True
            end = timeit.default_timer()
            duration = self._calc_duration(self.start, end)
            click.echo(_divider)
            click.secho(f'Seeding duration: {duration} seconds')
        finally:
            self._enable_fk_checks()
            self.session.close()
            if success:
                click.secho(' ✓ Seeding process completed successfully.')
            else:
                click.secho(' ⓧ Seeding process failed.', err=True)


def run_seeder(dataset, batch_size: int = MAX_ROWS_PER_BATCH, force: bool = False,
               update: bool = False, truncate: bool = False):
    click.echo(_divider)
    if dataset not in SUPPORTED_DATASETS:
        click.secho(f'Unsupported dataset "{dataset}" not in {SUPPORTED_DATASETS}', err=True)
        return
    click.secho(f'Initializing Data Seeder for dataset {dataset} with batch_size {batch_size}')

    if dataset not in SUBDIRS or dataset not in CONFIG:
        click.secho(f'Supported dataset "{dataset}" not implemented.', err=True)
        click.echo(_divider)
        return
    subdir = SUBDIRS.get(dataset)
    config = CONFIG.get(dataset)

    # main initialization and execution
    try:
        DataSeeder(config=config,
                   batch_size=int(batch_size),
                   update_existing=update,
                   truncate=truncate,
                   data_subdir=subdir,
                   disable_constraint_checks=force).run()
    except Exception as e:
        click.echo(traceback.print_exc(), err=True)
        click.secho(_divider, err=True)
        click.secho(' ⓧ Seeding process failed.', err=True)
        click.secho(_divider, err=True)
        click.secho('Reason: %s' % e, err=True)
    finally:
        click.echo(_divider)


@click.group()
def cli():
    pass


@click.command('run-seeder', short_help=f'Available datasets: {SUPPORTED_DATASETS}')
@click.argument('dataset')
@click.option('--batch-size', default=MAX_ROWS_PER_BATCH, help='Maximum number of rows per INSERT')
@click.option('--force', is_flag=True, help='Ignore foreign key constraints and write rows.')
@click.option('--update', is_flag=True, help='Update existing data instead of skipping table.')
@click.option('--truncate', is_flag=True, help='Truncate tables before seeding (--update ignored).')
def run_seeder_cmd(dataset, batch_size: int = MAX_ROWS_PER_BATCH, force: bool = False,
                   update: bool = False, truncate: bool = False):
    run_seeder(**locals())


cli.add_command(run_seeder_cmd)

if __name__ == '__main__':
    cli()