#!/usr/bin/env python
# -*- coding: utf-8

import os
import sys
import glob
import argparse

import anvio.tables as t
import anvio.terminal as terminal 
import anvio.dbops as dbops
import anvio.db as db

from anvio.errors import ConfigError


run = terminal.Run()
progress = terminal.Progress()


def update_profile_db_from_v4_to_v5(profile_db_path):
    if profile_db_path is None:
        raise ConfigError("No profile database is given.")

    # make sure someone is not being funny
    dbops.is_profile_db(profile_db_path)

    # make sure the version is 4
    profile_db = db.DB(profile_db_path, None, ignore_version = True)
    if str(profile_db.get_version()) != '4':
        raise ConfigError("Version of this profile database is not 4 (hence, this script cannot really do anything).")

    progress.new("Trying to upgrade the profile database")
    progress.update('...')

    if t.states_table_name in profile_db.get_table_names():
        progress.end()
        raise ConfigError("It seems '%s' table already exists in your profile database." % t.states_table_name)

    # create the table
    profile_db.create_table(t.states_table_name, t.states_table_structure, t.states_table_types)   

    # set the version
    profile_db.remove_meta_key_value_pair('version')
    profile_db.set_version('5')

    # bye
    profile_db.disconnect()
    progress.end()
    run.info_single("An empty states table is generated.")

    # next: states
    progress.new("Importing state files")
    progress.update('...')

    profile_db_full_path = os.path.join(os.getcwd(), profile_db_path)
    dir_name = os.path.dirname(profile_db_full_path)
    
    search_pattern = os.path.join(dir_name, 'state*.json')
    state_files = glob.glob(search_pattern)
 
    for state_file_path in state_files:
        progress.update("Importing '%s' ..." % state_file_path)

        path, filename = os.path.split(state_file_path)
    
        #replace .json and use as key.
        state_id = filename.replace('.json', '')
     
        # get json file content
        with open(state_file_path) as f: content = f.read()
    
        states = dbops.TablesForStates(profile_db_path, '5')
        states.store_state(state_id, content)

    progress.end()
    run.info_single("%d state files added to the database." % len(state_files)) if len(state_files) else None

    run.info_single("Database successfully upgraded.")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='A simple script to upgrade profile database to version 5 and import all state files into database')
    parser.add_argument('profile_db', metavar = 'PROFILE_DB', help = 'Profile database')
    args = parser.parse_args()

    try:
        update_profile_db_from_v4_to_v5(args.profile_db)
    except ConfigError as e:
        print(e)
        sys.exit(-1)
