"""
File Purpose: sqlite3 tests for registry in PlasmaCalcs.
"""

import os
import pytest

import PlasmaCalcs as pc
pc.DEFAULTS.pic_ambiguous_unit = dict(u_t=1)  # must be defined before loading an EppicCalculator.

HERE = os.path.dirname(__file__)
TESTS_DIR = os.path.dirname(HERE)
ARTIFACTS = os.path.join(TESTS_DIR, '_test_artifacts', 'registry_sqlite')
if os.path.exists(ARTIFACTS):  # remove any xarray_io artifacts before running tests.
    import shutil; shutil.rmtree(ARTIFACTS)
# <-- don't makedirs here; register() should makedirs as needed.


def get_calculator(name, **kw_init):
    '''get one of the standard "test-suite" calculators used in this file.
    These are:
        'ec_tiny': EppicCalculator from tests/test_eppic/test_eppic_tinybox
        'bc_tiny': BifrostCalculator from tests/test_bifrost/test_bifrost_tinybox
        'bm_tiny': BifrostMultifluidCalculator from tests/test_bifrost/test_bifrost_tinybox
    '''
    if name == 'ec_tiny':
        with pc.InDir(os.path.join(TESTS_DIR, 'test_eppic', 'test_eppic_tinybox')):
            return pc.EppicCalculator.from_here(**kw_init)
    elif name == 'bc_tiny':
        with pc.InDir(os.path.join(TESTS_DIR, 'test_bifrost', 'test_bifrost_tinybox')):
            return pc.BifrostCalculator(**kw_init)
    elif name == 'bm_tiny':
        with pc.InDir(os.path.join(TESTS_DIR, 'test_bifrost', 'test_bifrost_tinybox')):
            return pc.BifrostMultifluidCalculator(**kw_init)
    else:
        raise pc.InputError(f'unknown calculator name: {name}')


### --------------------- basic tests --------------------- ###

def test_registry_basic():
    '''test basic registry functionality, e.g. create a registry, register run'''
    ec = get_calculator('ec_tiny')
    bc = get_calculator('bc_tiny')
    bm = get_calculator('bm_tiny')

    ## in-memory registry lo register a single run
    registry00 = pc.Registry(':memory:')
    with registry00.connect() as con:
        # check column names match what we expect:
        runs_cols = [row[1] for row in con.execute('PRAGMA table_info(runs)').fetchall()]
        assert set(runs_cols) >= {'run_id', 'pc_uuid', 'calculator_type', 'registered_at_first', 'registered_at'}
        params_cols = [row[1] for row in con.execute('PRAGMA table_info(params)').fetchall()]
        assert set(params_cols) >= {'run_id', 'param_id', 'name', 'value', 'value_text'}
        # ... could repeat for other tables too, but that's a lot of stuff to hard-code,
            # which might get annoying to update here if the schema ever changes.
        # check that all tables we expect to exist actually exist:
        tables = [row[0] for row in con.execute('SELECT name FROM sqlite_master WHERE type="table"').fetchall()]
        assert set(tables) >= {'runs', 'params', 'metrics', 'register_history'}
        # register a run to this registry:
        (reg, _r00) = ec.register(registry00)
        assert _r00 is registry00
        assert len(reg) == 2   # (run_id, history_id)
        assert reg.history_id is not None   # i.e., actually added a new run.
        run_id = reg.run_id
    # check that the run we just registered is actually in the registry
    # (this is mostly a check that the ":memory:" connection logic works as expected;
    #  i.e. the registry00 connection lives if the registry00 object lives.)
    with registry00.connect() as con:
        assert (run_id,) in con.execute('SELECT run_id FROM runs').fetchall()
        assert 1 == con.execute('SELECT COUNT(*) FROM runs').fetchone()[0]  # 1 entry
        assert 1 == registry00.sizes['runs']  # check that sizes works properly
    # check that if we make a new registry in ":memory:" it is not the same registry
    # (this is another check that the special ":memory:" processing works properly
    registry01 = pc.Registry(':memory:')
    with registry01.connect() as con:
        assert (run_id,) not in con.execute('SELECT run_id FROM runs').fetchall()
        assert 0 == con.execute('SELECT COUNT(*) FROM runs').fetchone()[0]  # 0 entries
    # check that we can make a (different) registry in ":memory:" from a calculator
    reg, registry02 = bc.register(None)
    assert registry02.location == ':memory:'   # None <--> ":memory:"
    with registry02.connect() as con:
        assert 1 == con.execute('SELECT COUNT(*) FROM runs').fetchone()[0]  # 1 entry
        assert 'BifrostCalculator' == con.execute('SELECT calculator_type FROM runs WHERE run_id=?',
                                                  (reg.run_id,)).fetchone()[0]

    ## registry in a file:
    REG10_PATH = os.path.join(ARTIFACTS, 'registry10.db')
    # check can register to file:
    (ec_run_id, ec_history_id), registry10 = ec.register(REG10_PATH)
    with registry10.connect() as con:
        assert (ec_run_id,) in con.execute('SELECT run_id FROM runs').fetchall()
        assert 1 == con.execute('SELECT COUNT(*) FROM runs').fetchone()[0]  # 1 entry
    # check can register another run:
    (bc_run_id, _), _ = bc.register(registry10)
    with registry10.connect() as con:
        assert 2 == con.execute('SELECT COUNT(*) FROM runs').fetchone()[0]  # 2 entries

    # check can register a run to same location by provided path:
    # (start with clean registry)
    REG11_PATH = os.path.join(ARTIFACTS, 'registry11.db')
    _, registry11 = ec.register(REG11_PATH)
    _, registry11b = bc.register(REG11_PATH)
    assert registry11b is not registry11  # different objects
    assert registry11b.location == registry11.location  # same file location

    ## misc check: ensure pc_uuid works as expected, i.e. reads from file when it exists:
    EC_TINYBOX_UUID = "c51a5543c96d4a41b456228bec4f9f22"   # value in tests/test_eppic/test_eppic_tinybox/_pc_uuid.txt
    with registry10.connect() as con:
        ec_uuid = con.execute('SELECT pc_uuid FROM runs WHERE run_id=?', (ec_run_id,)).fetchone()[0]
        assert ec_uuid == EC_TINYBOX_UUID   # same value as in file
        bc_uuid = con.execute('SELECT pc_uuid FROM runs WHERE run_id=?', (bc_run_id,)).fetchone()[0]
        assert bc_uuid != EC_TINYBOX_UUID   # new uuids should (practically) always be unique.

    ## re-registering a run, updating metrics (re-uses some variables from above.)
    ec.register(REG10_PATH)
    with registry10.connect() as con:
        # no new runs (this run already exists)
        assert 2 == con.execute('SELECT COUNT(*) FROM runs').fetchone()[0]
        # no new history entries (no new metrics / info entered)
        assert 2 == con.execute('SELECT COUNT(*) FROM register_history').fetchone()[0]
    # repeat those tests but for bc_tiny (which doesn't have _pc_uuid.txt file natively on repo)
    #  (this is mostly a test that _pc_uuid.txt file is being properly generated.)
    bc.register(REG10_PATH)
    with registry10.connect() as con:
        assert 2 == con.execute('SELECT COUNT(*) FROM runs').fetchone()[0]
        assert 2 == con.execute('SELECT COUNT(*) FROM register_history').fetchone()[0]
    # okay, now add metrics:
    bc.register(REG10_PATH, metrics=['max_mod_B', {'custom_metric0': 7}])
    with registry10.connect() as con:
        assert 2 == con.execute('SELECT COUNT(*) FROM runs').fetchone()[0]  # no new runs
        assert 3 == con.execute('SELECT COUNT(*) FROM register_history').fetchone()[0]  # 1 new history entry
        assert 2 == registry10.sizes['runs']  # check that sizes works
        assert 3 == registry10.sizes['register_history']  # check that sizes works
        # check that the new metrics are actually in the metrics table:
        modB = con.execute('''
            SELECT run_id, value, is_pc_quantity
            FROM metrics WHERE name = "max_mod_B"
        ''').fetchone()
        assert modB[0] == bc_run_id
        assert modB[1] == bc('max_mod_B')
        assert modB[2]   # came from bc('max_mod_B') so is_pc_quantity should be True
        custom0 = con.execute('''
            SELECT value, is_pc_quantity
            FROM metrics WHERE name="custom_metric0"
        ''').fetchone()
        assert custom0[0] == 7
        assert not custom0[1]  # not from pc quantity, so is_pc_quantity should be False
    # adding more of the same metrics with same values changes nothing:
    bc.register(REG10_PATH, metrics=['max_mod_B'])
    assert 2 == registry10.sizes['runs']
    assert 3 == registry10.sizes['register_history']
    # updating a metric value changes something:
    (_, history_id_4), _ = bc.register(REG10_PATH, metrics=[{'custom_metric0': 42}])
    assert 4 == registry10.sizes['register_history']
    with registry10.connect() as con:
        custom0 = con.execute('''
            SELECT value, is_pc_quantity
            FROM metrics WHERE name="custom_metric0"
        ''').fetchone()
        assert custom0[0] == 42      # updated value stored in value column,
        assert not custom0[1]  # not from pc quantity, so is_pc_quantity should be False
        history_entry = con.execute('''
            SELECT updated_metadata, updated_metrics
            FROM register_history WHERE history_id=?
        ''', (history_id_4,)).fetchone()
        assert history_entry is not None
        assert history_entry[0] == ''  # [TODO] maybe should be None instead? comes from ', '.join(empty list)
        assert history_entry[1] == 'custom_metric0'
    
    ## check that registering a run with a different calculator type but same _pc_uuid.txt
    # will be treated as the same run, and update the entry appropriately.
    (bm_run_id, _), _ = bm.register(REG10_PATH)
    assert bm_run_id == bc_run_id
    assert 2 == registry10.sizes['runs']
    assert 5 == registry10.sizes['register_history']
    # check that the calculator_type was updated to the new one
    # (also checks that registry.execute() works)
    calc_type = registry10.execute('SELECT calculator_type FROM runs WHERE run_id=?', (bm_run_id,)).fetchone()[0]
    assert calc_type == 'BifrostMultifluidCalculator'


def test_registry_merge():
    '''test registry merging.
    [TODO] more comprehensive tests...
        e.g., "if some runs in both original registries, keep values from first, only."
    '''
    ec = get_calculator('ec_tiny')
    bc = get_calculator('bc_tiny')

    _, r0 = ec.register(':memory:')
    _, r1 = bc.register(':memory:')
    r2 = r1.merge(r0)

    assert r2.sizes['sources'] == 3  # source r2, and sources from r0 and r1.
    assert r2.sizes['runs'] == 2

    rid0, uuid0 = r0.execute('SELECT run_id, pc_uuid FROM runs').fetchone()
    rid1, uuid1 = r1.execute('SELECT run_id, pc_uuid FROM runs').fetchone()

    # check that the run from r0 is in r2 with same pc_uuid, and source_run_id assigned appropriately.
    assert (rid0, uuid0) in r2.execute('SELECT source_run_id, pc_uuid FROM runs').fetchall()
    # check that the run from r1 is in r2 with same pc_uuid, and source_run_id assigned appropriately.
    assert (rid1, uuid1) in r2.execute('SELECT source_run_id, pc_uuid FROM runs').fetchall()

    # ensure that source names were transferred appropriately
    oname0 = r0.execute('SELECT name FROM sources WHERE source_id = 0').fetchone()[0]
    oname1 = r1.execute('SELECT name FROM sources WHERE source_id = 0').fetchone()[0]
    oname2 = r2.execute('SELECT name FROM sources WHERE source_id = 0').fetchone()[0]

    onames = [row[0] for row in r2.execute('SELECT name FROM sources').fetchall()]

    assert set([oname0, oname1, oname2]) == set(onames)
