import warnings
import logging
import threading
import zipfile
from pathlib import Path
import tempfile
import json
import os
import shutil

import cobra
from escher import Builder

from .downloads import SimpleLoadingWheel



def print_json_tree(data, level=0, max_level=2):
    # explore contents of a json object
    
    if level > max_level:
        return
    indent = '  ' * level
    if isinstance(data, dict):
        for key, value in data.items():
            print(f"{indent}{key}")
            print_tree(value, level + 1, max_level)
    elif isinstance(data, list):
        for i, item in enumerate(data):
            print(f"{indent}[{i}]")
            print_tree(item, level + 1, max_level)
            
            
            
def subset_for_focus(universe, rids_in_group, outdir, focus):
            
    universe_focus = universe.copy()
    to_remove = [r for r in universe_focus.reactions if r.id not in rids_in_group]


    # trick to avoid the WARNING "cobra/core/group.py:147: UserWarning: need to pass in a list" 
    # triggered when trying to remove reactions that are included in groups. 
    with warnings.catch_warnings():  # temporarily suppress warnings for this block
        warnings.simplefilter("ignore")  # ignore all warnings
        cobra_logger = logging.getLogger("cobra.util.solver")
        old_level = cobra_logger.level
        cobra_logger.setLevel(logging.ERROR)   

        universe_focus.remove_reactions(to_remove, remove_orphans=True)

        # restore original behaviour: 
        cobra_logger.setLevel(old_level)  


    # save the subset for drawing in Escher!
    cobra.io.save_json_model(universe_focus, f'{outdir}/focus_{focus}.json')
            
            
        
def count_undrawn_rids_focus(logger, universe, lastmap, focus, outdir):
    
    
    # there could be no tracked folder / no versions for this group
    if lastmap == None:
        return
    
    
    # get modeled reads for this --focus:
    rids_in_group = set()
    try: gr = universe.groups.get_by_id(focus)
    except: 
        logger.warning(f"Group '{focus}' not found!")
        return 
    for r in gr.members:
        rids_in_group.add(r.id)


    # get rids on Escher: 
    drawn_rids = set()
    for key, value in lastmap['json'][1]['reactions'].items():
        drawn_rids.add(value['bigg_id'])

        
    # get remaining rids for this map:
    remainings = rids_in_group - drawn_rids
    remainings_krs = set()
    for rid in remainings: 
        r = universe.reactions.get_by_id(rid)
        if 'kegg.reaction' in r.annotation.keys():
            krs = r.annotation['kegg.reaction']
            for kr in krs: 
                remainings_krs.add(kr)
    
    
    if len(remainings) > 0:
        if focus != 'gr_transport':
            logger.warning(f"Current '{lastmap['filename']}' is {len(remainings)} reactions behind: {' '.join(list(remainings_krs))}.")
        else:
            logger.warning(f"Current '{lastmap['filename']}' is {len(remainings)} reactions behind.")  # usually no kegg codes for tranport reactions
        
        
        # subset the universe to ease the drawing: 
        next_version_filename = f"{focus}-v{(int(lastmap['filename'].rsplit('-v',1)[-1].replace('.json', ''))+1)}.json"
        logger.warning(f"Writing model '{outdir}/focus_{focus}.json' to ease the drawing of '{next_version_filename}'...")
        
        
        t1 = threading.Thread(target = subset_for_focus, args=(
            universe, rids_in_group, outdir, focus))
        t1.start()
        slw = SimpleLoadingWheel(msg="Please wait... ")
        while t1.is_alive():
            slw.proceed()
        slw.clear()
        
        
        logger.warning(f"'{outdir}/focus_{focus}.json' created!")
    else:
        logger.info(f"Current '{lastmap['filename']}' is 0 reactions behind. Thank you ♥")
        
        
        
def parse_zipped_escher(logger, universe, escherzip, outdir): 
    # used to parse the zipped 'escher' folder downloaded from Google Drive.
    # 'escherzip' is the path to zipped folder. 
    
    
    logger.info("Processing collection of hand-drawn Escher maps...")
    
    
    # prepare empty logs folder
    shutil.rmtree(f'{outdir}/escherapps', ignore_errors=True)
    os.makedirs(f'{outdir}/escherapps', exist_ok=True)
    
    
    drawngid_to_maps = dict()
    drawnrid_to_maps = dict()
    drawnpuremid_to_maps = dict()
    with tempfile.TemporaryDirectory() as tmp_dir:
        
        # check path existance: 
        zip_path = Path(escherzip)  # convert to Path for convenience
        if not zip_path.exists():
            logger.error(f"Zipped escher does not exist at provided path '{zip_path}'.")
            return 1
        
        # extract in a auto-deleting temporary directory 
        with zipfile.ZipFile(zip_path, "r") as z:
            z.extractall(tmp_dir)

        # extract the inner child dirs:
        tmp_dir = Path(tmp_dir)
        child_dirs = [p for p in tmp_dir.iterdir() if p.is_dir()]
        escher_dir = child_dirs[0]  # the top-level folder in the ZIP
        inner_child_dirs = [p for p in escher_dir.iterdir() if p.is_dir()]
        
        # iterate the inner child dirs:
        for map_dir in inner_child_dirs:
            map_id = map_dir.name.rsplit('/',1)[-1]
            map_files = [f for f in map_dir.iterdir() if f.is_file()]
            map_files_name = set([f.name.rsplit('/',1)[-1].split('-v',1)[0] for f in map_files])
            
            # check presence of maps:
            if len(map_files_name) == 0:
                # still no maps in this folder
                continue
            
            # check name consistency:
            if len(map_files_name) != 1:
                logger.error(f"Name inconsistency in folder '{map_dir_name}'!")
                return 1
            
            # check consistency with the folder 
            if list(map_files_name)[0] != map_id:
                logger.error(f"Versions in '{map_id}' refer to a different map!")
            
            # get latest version
            map_versions = [int(f.name.rsplit('/',1)[-1].split('-v',1)[-1].replace('.json','')) for f in map_files]
            latest_version = max(map_versions)
            
            
            # read json (last version) to fill the dict
            latest_filepath = str(map_dir / f"{map_id}-v{latest_version}.json")
            with open(latest_filepath, 'r') as file:
                json_data = json.load(file)
        
                # get elements on Escher:
                map_drawngids = set()
                map_drawnrids = set()
                map_drawnpuremids = set()
                for key, value in json_data[1]['reactions'].items():
                    map_drawnrids.add(value['bigg_id'])
                    for i in value['genes']:
                        map_drawngids.add(i['bigg_id'])
                    #for i in value['metabolites']:
                    #    puremeid = i['bigg_id'].rsplit('_',1)[0]
                    #    map_drawnpuremids.add(puremeid)
                for key, value in json_data[1]['nodes'].items():
                    if value['node_type'] == 'metabolite':
                        puremeid = value['bigg_id'].rsplit('_',1)[0]
                        map_drawnpuremids.add(puremeid)
                    
                # populat dicts
                for drawngid in map_drawngids:
                    if drawngid not in drawngid_to_maps.keys():
                        drawngid_to_maps[drawngid] = set()
                    drawngid_to_maps[drawngid].add(map_id)
                for drawnrid in map_drawnrids:
                    if drawnrid not in drawnrid_to_maps.keys():
                        drawnrid_to_maps[drawnrid] = set()
                    drawnrid_to_maps[drawnrid].add(map_id)
                for drawnpuremid in map_drawnpuremids:
                    if drawnpuremid not in drawnpuremid_to_maps.keys():
                        drawnpuremid_to_maps[drawnpuremid] = set()
                    drawnpuremid_to_maps[drawnpuremid].add(map_id)
                    
                    
            # read json (last version) to create the escher app
            builder = Builder(
                map_json = latest_filepath,   
                model_json = None, 
            )
            builder.never_ask_before_quit = True
            builder.scroll_behavior = 'zoom'

            builder.menu = 'all'
            builder.enable_editing = False

            builder.enable_keys = False  # switch not working
            builder.enable_search = True  # switch not working
            builder.full_screen_button = True  # switch not working

            builder.highlight_missing = True 
            builder.show_gene_reaction_rules = True
            builder.enable_tooltips = 'label'
    
            builder.save_html(f"{outdir}/escherapps/{map_id}.html")
            
        
    # apply annotation to universe
    for g in universe.genes: 
        if g.id in drawngid_to_maps.keys():
            g.annotation['drawn_in_maps'] = list(drawngid_to_maps[g.id])
        else:
            g.annotation['drawn_in_maps'] = []
    for r in universe.reactions: 
        if r.id in drawnrid_to_maps.keys():
            r.annotation['drawn_in_maps'] = list(drawnrid_to_maps[r.id])
        else:
            r.annotation['drawn_in_maps'] = []
    for m in universe.metabolites: 
        puremid = m.id.rsplit('_',1)[0]
        if puremid in drawnpuremid_to_maps.keys():
            m.annotation['drawn_in_maps'] = list(drawnpuremid_to_maps[puremid])
        else:
            m.annotation['drawn_in_maps'] = []
            
    
    return universe