#!/usr/bin/env python3

""" This script takes a metadata file generated by the script 
``create_metadata.py`` and stores all the images listed in that file
stores them in a MongoDB using GridFS.
"""

import os
from typing import Any, List, Dict, Union
from pathlib import Path
import json
import multiprocessing
from multiprocessing import Pool, cpu_count, current_process
import sys
import inspect
import argparse
from unicodedata import name

from numpy import insert
from forensicfit.core import Tape
from forensicfit.database import Database
from forensicfit.utils.image_tools import IMAGE_EXTENSIONS


def get_chunks(files: List[Dict], n_processors: int) -> List[List[Dict]]:
    """divides list of files for multiprocessing

    Parameters
    ----------
    files : List[Dict]
        list of files in a dictionary format
    n_processors : int
        number of processors for Pool

    Returns
    -------
    List[List[Dict]]


    """
    
    ret = [ [] for x in range(n_processors)]
    n_files = len(files)
    for i, ifile in enumerate(files):
        ret[i % n_processors].append(ifile)
    return ret

def worker(args: Dict):
    """inserts images provided in the list to the MongoDB

    connects to a MongoDB, loads the images provided in the args, and 
    inserts them in the DB.

    Parameters
    ----------
    args : Dict
        dictionary containing the files, the database settings, and
        the inser options (skip, overwrite, etc)

    """
    
    files = args['files'] 
    db_settings = args['db_settings']
    insert_options = args['insert_options']
    db = Database(**db_settings)
    for _, entry in enumerate(files):
        file_path = Path(entry['source'])
        if not file_path.exists():
            print(f"{file_path.as_posix()} does not exist")
            return 
        if file_path.suffix not in IMAGE_EXTENSIONS:
            continue
        tape = Tape.from_file(file_path)
        # print(f"{file_path.stem}")
        for key in entry:
            tape.metadata[key] = entry[key]
        db.insert(tape, **insert_options)
        
def store_on_db(
        metadata_file: Union[str, Path],
        db_name: str,
        host:str ='localhost',
        port: int=27017,
        username: str="",
        password:str ="",
        n_processors: int=1,
        overwrite: bool=False,
        skip: bool=True):
    """given a metadata file about images, stores them on a database.

    Given a metadata file generated by create_metadata.py, it divides
    them into chunks, passes them to multiple processors, and inserts
    them on a MongoDB

    Parameters
    ----------
    metadata_file : Union[str, Path]
        Metadata file generated by ``create_metadata.py``
    db_name : str
        Name of the MongoDB to be stored on.
    host : str
        Address to the host of the MongoDB server.
    port : int
        Port of the MongoDB host
    username : str
        User name for the MongoDB host
    password : str
        Password for the MongoDB host
    n_processors : int
        Number of processors to be used.
    overwrite : bool
        If the file already exists, overwrite or create a new one
    skip : bool
        If the file already exists on the DB, skip or not.

    """

    if n_processors > cpu_count():
        n_processors = cpu_count()
        
    with open(metadata_file, 'r') as rf:
        metadata = json.load(rf)
    chunks = get_chunks(metadata, n_processors)
    
    db_settings = dict(name=db_name,
                       host=host,
                       port=port,
                       username=username,
                       password=password)
    insert_options = dict(overwrite=overwrite,
                          skip=skip)
    
    args = [{
            'files': x,
            'db_settings': db_settings,
            'insert_options': insert_options,
        } for x in chunks]
    with Pool(n_processors) as p:
        p.map(worker, iterable=args)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-dbn', '--database-name', 
                           dest='db_name',
                           help='The name of the database.')
    parser.add_argument('-dbh', '--database-host', 
                           dest='db_host',
                           default='localhost',
                           help='The name of the host for the database server.')
    parser.add_argument('-dbp', '--database-port', 
                           dest='db_port',
                           type=int,
                           default=27017,
                           help='The port of the database server.')

    parser.add_argument('-i', '--input',
                                dest='input',
                                help=('Path to the metadata.json file that' 
                                      'contains path all the files. This file ' 
                                      'can be generated by create_metadata.py '
                                      'script')
                                )
    parser.add_argument('-np', '--n-processors', 
                        dest='n_processors',
                        type=int,
                        help='Number of processors available for this task.',
                        default=1)
    parser.add_argument('--overwrite',
                        dest='overwrite',
                        help=('To overwrite the existing entries if they already exist'),
                        action=argparse.BooleanOptionalAction,
                        type=bool,
                        default=False)
    parser.add_argument('--skip',
                        dest='skip',
                        help=('To overwrite the existing entries if they already exist'),
                        action=argparse.BooleanOptionalAction,
                        type=bool,
                        default=False)
    args = parser.parse_args()
    store_on_db(
        metadata_file=args.input,
        db_name=args.db_name,
        host=args.db_host,
        port=args.db_port,
        n_processors=args.n_processors,
        overwrite=args.overwrite,
        skip=args.skip)
