import os
from pathlib import Path
import logging
import tempfile
import concurrent.futures as confu 
import sys

import cobra

from .pruner import load_input_universe
from .pruner import prune_universe
from .pruner import load_input_eggnog
from .pruner import load_keggorg_like_eggnog
from .pruner import parse_eggnog
from .pruner import subtract_kos
from .pruner import translate_remaining_kos
from .pruner import restore_gene_annotations
from .pruner import append_keggorg_gene_annots

from .gapfillutils import include_forced

from .gapfill import gapfill_on_media

from .polishing import remove_orphans
from .polishing import remove_sinks_demands
from .polishing import remove_forced
from .polishing import remove_disconnected

from .biologcuration import biolog_on_media

from ..runsims.biosynth import biosynthesis_on_media

from ..runsims.simplegrowth import grow_on_media

from ..parsedb.cycles import verify_egc_all

from ..commons.logutils import get_logger

from ..commons.downloads import get_databases
from ..commons.downloads import eggnog_format_expansion

from ..commons.biomass import adjust_biomass_precursors
from ..commons.biomass import adjust_biomass_coefficients

from ..commons.sbmlutils import force_id_on_sbml

from ..commons.excelhub import write_excel_model
from ..commons.excelhub import comparative_table

from ..commons.metrics import log_metrics
from ..commons.metrics import log_unbalances

from ..commons.keggutils import download_keggorg

from ..commons.medium import initialize_model
from ..commons.medium import reset_growth_env 

from ..commons.memoteutils import get_memote_results_dict



def create_model_incore(params):
    
    
    # load params
    annotation_source, universe, pruned_universe, eggpath, dbuni, dbexp, args, multistrain = params
    
    
    # get the logger:
    logger = get_logger('gsrap_queued', args.verbose)  # loggers can't be pickled!
    
    
    # disable logging if strains are more than 1:
    # only errors will be recorded if multistrain mode
    if multistrain:
        logger.setLevel(logging.ERROR)  



    ###### STRAIN-SPECIFIC PRUNING
    
    # create a copy of the universe and define the model ID
    logger.info("Starting reconstruction from a copy of whole universe...")
    model = universe.copy()  # 'universe' and not 'pruned_universe' by design
    
    
    # assign model ID
    if annotation_source == 'keggorg':
        model.id = args.keggorg
    elif annotation_source == 'eggnog':
        model.id = Path(eggpath).stem 

        
    # remove all sinks and demands, excluding sn_meanfa (removed later).
    n = remove_sinks_demands(logger, model, keep_meanfa=True)  
    logger.info(f"Removed {n} universal sinks or demand reactions.")
    

    # remove all orphans 
    n = remove_orphans(logger, model, keep_transport=False) 
    logger.info(f"Removed {n} universal orphan reactions.")

    
    # get the annotation dictionaries
    if annotation_source == 'keggorg':
        eggnog_style_table = load_keggorg_like_eggnog(logger, args.keggorg, args.outdir)
        logger.info(f"Reading annotation for organism code '{args.keggorg}'...")
    elif annotation_source == 'eggnog':
        eggnog_style_table = load_input_eggnog(logger, eggpath)
        logger.info("Reading provided eggnog-mapper annotation...")
    # get important dictionaries: 'eggnog_ko_to_gids' and 'eggonog_gid_to_kos'
    eggnog_ko_to_gids, eggonog_gid_to_kos = parse_eggnog(eggnog_style_table)  
      

    # prune reactions
    subtract_kos(logger, model, eggnog_ko_to_gids)

    
    # translate KOs to the actual genes
    translate_remaining_kos(logger, model, eggnog_ko_to_gids)
    restore_gene_annotations(logger, model, universe, eggonog_gid_to_kos)
    # insert gene annotation if starting from kegg organisms:
    if annotation_source == 'keggorg': 
        append_keggorg_gene_annots(logger, model, args.keggorg, args.outdir)

    

    ###### GAPFILLING
    
    # force inclusion of reactions prior to gap-filling:  
    include_forced(logger, model, universe, args.include)

    
    # remove missing conditional precursors + get the 'cond_col_dict' dict.
    # 'cond_col_dict' is str-to-str: {'pheme_c': 'M00868: 1/8; M00121: 2/12;', 'hemeO_c': 'gr_HemeO: 0/1'}
    cond_col_dict = adjust_biomass_precursors(logger, model, universe, args.conditional, dbuni) 

    
    # adjust biomass coefficients
    response = adjust_biomass_coefficients(logger, model, universe, dbuni, dbexp, args.biomass)
    if response == 1: return 1


    # gap-fill based on media:
    df_B = gapfill_on_media(logger, model, pruned_universe, universe, dbuni, dbexp, args.gapfill, cond_col_dict)
    if type(df_B)==int: return 1


    # force removal of reactions
    # experimental feature, not public: it's main purpose was to test gap-filling in biolog_on_media().
    setattr(args, 'remove', '-')  
    remove_forced(logger, model, universe, args.remove)

    
    # perform Biolog(R) curation based on media
    df_P = biolog_on_media(logger, model, pruned_universe, universe, dbexp, args.gapfill, args.biolog, args.cnps)
    if type(df_P)==int: return 1

    
    # remove disconnected metabolites
    n = remove_disconnected(logger, model)
    logger.info(f"Removed {n} disconnected metabolites.")

    
    # log metrics
    log_metrics(logger, model)
    log_unbalances(logger, model)


    
    ###### POST-RECONSTCTION
    
    # check erroneous EGCs
    verify_egc_all(logger, model, args.outdir)
    
    
    # check blocked metabolites / dead-ends
    df_S = biosynthesis_on_media(logger, model, dbexp, args.gapfill, args.biosynth)
    if type(df_S)==int: return 1


    # reset growth environment befor saving the model (changed during growth sims)
    reset_growth_env(model)

    
    # initialize model
    response = initialize_model(logger, model, dbexp, args.initialize, args.gapfill)
    if response==1: return 1


    # compute Memote metrics
    memote_results_dict = get_memote_results_dict(logger, model)


    # write the model on disk:
    logger.info("Writing strain-specific model...")
    cobra.io.save_json_model(model, f'{args.outdir}/{model.id}.json')        # JSON
    logger.info(f"'{args.outdir}/{model.id}.json' created!")
    cobra.io.write_sbml_model(model, f'{args.outdir}/{model.id}.xml')        # SBML   # groups are saved only to SBML
    logger.info(f"'{args.outdir}/{model.id}.xml' created!")
    force_id_on_sbml(f'{args.outdir}/{model.id}.xml', model.id)   # force introduction of the 'id=""' field
    sheets_dict = write_excel_model(model, f'{args.outdir}/{model.id}.mkmodel.xlsx', args.nofigs, memote_results_dict, None, df_B, df_P, df_S)  
    logger.info(f"'{args.outdir}/{model.id}.mkmodel.xlsx' created!")
        
        
    return sheets_dict    



def multicore_generation_handling(universe, pruned_universe, args, annotation_source, dbuni, dbexp, multistrain):
    
    
    # base objects
    error_raised = False
    sheets_dicts = []
    executor =  confu.ProcessPoolExecutor(max_workers=args.cores)

    
    # run executor according to the annotation source
    if annotation_source == 'keggorg':
        futures = [executor.submit(create_model_incore, (
            annotation_source, universe, pruned_universe, None, dbuni, dbexp, args, multistrain
        ))]
    elif annotation_source == 'eggnog':
        futures = [executor.submit(create_model_incore, (
            annotation_source, universe, pruned_universe, eggpath, dbuni, dbexp, args, multistrain
        )) for eggpath in args.eggnog]
    
    
    # collect results as the runs complete
    for f in confu.as_completed(futures):
        sheets_dict = f.result() 
        
        
        # error raised in one of the cores:
        if type(sheets_dict) == int: 
            error_raised = True
            break   # a single error from a child process is sufficient to shutdown all the executor
        
        
        # this strain reached the end:
        else:
            if multistrain:
                sheets_dicts.append(sheets_dict)
                print(f"{len(sheets_dicts)}/{len(args.eggnog)} ({int(len(sheets_dicts)/len(args.eggnog)*100)}%) completed!", end='\r', file=sys.stderr)
    
    
    # hide last progress trace ('sheets_dicts' unused if not in multi-strain mode):
    if multistrain and sheets_dicts != []:
        last_trace = f"{len(sheets_dicts)}/{len(args.eggnog)} ({int(len(sheets_dicts)/len(args.eggnog)*100)}%) completed!"
        whitewash = ''.join([' ' for i in range(len(last_trace))])
        print(whitewash, end='\r', file=sys.stderr)   
    
    
    # multiproces part terminated: safely shut down the executor
    executor.shutdown(wait=True)
    
    
    # propagate errors if any:
    if error_raised:
        return 1
    
    
    return sheets_dicts



def main(args, logger):
    
    
    # adjust out folder path                      
    while args.outdir.endswith('/'):              
        args.outdir = args.outdir[:-1]
    os.makedirs(f'{args.outdir}/', exist_ok=True)
    
    
    # check compatibility of input parameters:
    if args.cnps == '-' and args.biolog != '-':
        logger.error("Missing starting C/N/P/S sources: --biolog must be used in conjunction with --cnps.")
        return 1
    if args.transrule not in ['rescue', 'ignore']:
        logger.error("Unrecognized value for --transrule: it must be either 'rescue' or 'ignore'.")
        return 1
    
    
    # determine available cores
    cpu_count = os.cpu_count()
    if cpu_count == None: 
        logger.debug("Failed to count cpus. Reset to 1.")
        cpu_count = 1  
    if args.cores == 0:
        args.cores = cpu_count
    elif args.cores > cpu_count:
        logger.debug(f"Parameter --cores {args.cores} is greater than the maximum available ({cpu_count}): reset to {cpu_count}.")
        args.cores = cpu_count
    
    
    # format the --eggnog param
    args.eggnog = eggnog_format_expansion(logger, args.eggnog)  # now 'args.eggnog' could still be '-'
    
    
    # get the kegg organism if requested
    if args.keggorg != '-':
        response = download_keggorg(logger, args.keggorg, args.outdir)
        if response == 1: return 1
        
    
    # determine the source of functional annotation:
    annotation_source = None
    if args.keggorg != '-':  # keggorg has precedence
        annotation_source = 'keggorg'
    elif args.eggnog != '-':
        annotation_source = 'eggnog'
        if args.cores > len(args.eggnog):
            logger.debug(f"Parameter --cores {args.cores} is greater than the number of strains ({len(args.eggnog)}): reset to {len(args.eggnog)}.")
            args.cores = len(args.eggnog)
    else:
        logger.error("No valid functional annotations provided: please use '--keggorg' or '--eggnog'.")
        return 1
    
    
    # get dbexp (dbuni not really used):
    logger.info("Downloading updated gsrap assets...")
    response = get_databases(logger)
    if type(response)==int: return 1
    else: dbuni, dbexp = response 
    
    
    # load universe (fresh or built-in):
    universe = load_input_universe(logger, args.universe)
    if type(universe)==int:
        return 1
    
    
    # check and use the --taxon param to prune the universe
    pruned_universe = prune_universe(logger, universe, args.taxon, args.transrule)
    if type(pruned_universe) == int: return 1
    
    
    # determine if mkmodel will be run in multistrain mode:
    if annotation_source == 'keggorg':
        multistrain = False
    elif annotation_source == 'eggnog':
        multistrain = len(args.eggnog) > 1
        if multistrain:
            logger.info(f"Number of provided strains is >1: logging will be disabled.")
            logger.info(f"Performing {len(args.eggnog)} reconstructions relying on {args.cores} cores... ")
            
            
    # create strain-specific GSMMs using multi-core
    sheets_dict = multicore_generation_handling(universe, pruned_universe, args, annotation_source, dbuni, dbexp, multistrain)
    if type(sheets_dict) == int: 
        return 1  # some error was raised 
    
        
    # everything went fine, and there are > 1 strains: time to make the comparison table:
    if multistrain:
        logger.info("Collecting comparative tables...")
        comparative_table(logger, args.outdir, sheets_dicts)
            
        
    return 0