from metadatastore.mds import MDS, MDSRO
import metadatastore.conf
from collections import deque
from tqdm import tqdm

def main():
    old_config = dict(metadatastore.conf.connection_config)
    new_config = old_config.copy()

    new_config['database'] = 'metadatastore_production_v1'

    old = MDSRO(version=0, config=old_config)
    new = MDS(version=1, config=new_config)

    # Drop all indexes on event collection to speed up insert.
    # They will be rebuilt the next time an MDS(RO) object connects.
    '''
    new._event_col.drop_indexes()
    total = old._runstart_col.find().count()
    for start in tqdm(old.find_run_starts(), desc='start docs', total=total):
        new.insert('start', start)
    '''
    total = old._runstop_col.find().count()
    for stop in tqdm(old.find_run_stops(), desc='stop docs', total=total):
        try:
            new.insert('stop', stop)
        except RuntimeError:
            print("error inserting run stop with uid {!r}".format(stop['uid']))
    '''
    descs = deque()
    counts = deque()
    total = old._descriptor_col.find().count()
    for desc in tqdm(old.find_descriptors(), unit='descriptors', total=total):
        d_raw = next(old._descriptor_col.find({'uid': desc['uid']}))
        num_events = old._event_col.find({'descriptor_id': d_raw['_id']}).count()
        new.insert('descriptor', desc)
        descs.append(desc)
        counts.append(num_events)

    total = sum(counts)
    with tqdm(total=total, unit='events') as pbar:
        for desc, num_events in zip(descs, counts):
            if num_events:
                # skip empty event stream of bulk insert raises
                events = old.get_events_generator(descriptor=desc, convert_arrays=False)
                try:
                    new.bulk_insert_events(descriptor=desc, events=events)
                except Exception:
                    print(desc['uid'])
                    raise
                pbar.update(num_events)
    '''
if __name__ == '__main__':
    main()