import os
import re
import threading

from socket import error
from uuid import uuid4

import dearpygui.dearpygui as dpg
from paramiko.ssh_exception import BadHostKeyException, AuthenticationException, UnableToAuthenticate, NoValidConnectionsError, SSHException

from utils import show_message, get_config_file_path, parse_dirs, construct_folder_structure, get_config, update_config_values, get_run_window_value, set_run_window_value, delete_item_children_and_clear_aliases
from remote_run import Connection, BuildingSession, RunningSession
from remote_run_constants import *
from style_constants import FONT_SIZE
from dynamic_ui import construct_model_output_structure_ui
from config import dump_config
from style import get_accent_color

def update_run_model_window():
    builds = get_run_window_value('builds')
    connections = get_run_window_value('connections')

    machine = LABEL_TO_MACHINE[dpg.get_value('machine')]
    model = LABEL_TO_MODEL[dpg.get_value('model')]

    if (machine, model) in builds:
        dpg.hide_item('credentials')
    elif machine in connections:
        dpg.show_item('gitlab_group')
        dpg.hide_item('lab_server_group')
        dpg.hide_item('lomonosov_server_group')
    else:
        if machine == LAB:
            dpg.show_item('gitlab_group')
            dpg.show_item('lab_server_group')
            dpg.hide_item('lomonosov_server_group')
        elif machine == LOMONOSOV:
            dpg.show_item('gitlab_group')
            dpg.hide_item('lab_server_group')
            dpg.show_item('lomonosov_server_group')

def update_run_model_window_clb(s, a, u):
    update_run_model_window()

def show_run_model_window_clb(s, a, u):
    if get_config_file_path():
        update_config_values()
        dump_config(get_config(), get_config_file_path())

    dpg.set_value('config_file_path', get_config_file_path())

    update_run_model_window()

    dpg.show_item('run_model_window')

def add_running_session_ui(running_session_id, machine, model):
    with dpg.table_row(parent='models_table', tag=running_session_id):
        dpg.add_text(MODEL_TO_LABEL[model], tag=f'{running_session_id}_model')
        dpg.add_text(MACHINE_TO_LABEL[machine], tag=f'{running_session_id}_machine')
        dpg.add_text('building', color=list(get_accent_color()[:3]) + [255], tag=f'{running_session_id}_status')
        with dpg.group(horizontal=False):
            dpg.add_spacer(height=0)
            dpg.add_progress_bar(label='progress', default_value=0.0, tag=f'{running_session_id}_progress_bar', height=FONT_SIZE, width=150)
        dpg.add_button(label='download output', enabled=False, tag=f'{running_session_id}_download_button', callback=show_download_output_window_clb, user_data=running_session_id)
    dpg.bind_item_theme(f'{running_session_id}_progress_bar', 'progress_bar_theme')

def make_execute_command():
    np = int(dpg.get_value('-np'))
    arch = dpg.get_value('-arch')
    udump = dpg.get_value('-udump')
    model_output = dpg.get_value('-model-output')

    command = f'mpirun -np {np} ./nsenx -arch {arch}'
    if dpg.get_value('use-udump'):
        command += f' -udump {dpg.get_value('-udump')}'
    if dpg.get_value('use-model-output'):
        command += f' -model-output {dpg.get_value('-model-output')}'

    return command

def update_progress(running_session_id, output):
    if not re.findall(r'([0-9]+)%', output):
        progress = dpg.get_value(f'{running_session_id}_progress_bar')
    else:
        progress = int(re.findall(r'([0-9]+)%', output)[-1]) / 100
    dpg.set_value(f'{running_session_id}_progress_bar', progress)

def run_model():
    dpg.hide_item('run_model_window')

    config_file_path = dpg.get_value('config_file_path')

    if not os.path.exists(config_file_path):
        show_message(f'no {config_file_path} file :(')
        return

    builds = get_run_window_value('builds')
    connections = get_run_window_value('connections')

    machine = LABEL_TO_MACHINE[dpg.get_value('machine')]
    model = LABEL_TO_MODEL[dpg.get_value('model')]

    if machine not in connections:
        if machine == LAB:
            username = dpg.get_value('server_username')
            password = dpg.get_value('server_password')
        elif machine == LOMONOSOV:
            username = dpg.get_value('lomonosov_username')
            if 'ID_RSA_PATH' in os.environ:
                id_rsa_path = os.environ['ID_RSA_PATH']
            else:
                id_rsa_path = None
        try:
            if machine == LAB:
                conn = Connection(MACHINE_TO_HOST[machine], username, password)
            elif machine == LOMONOSOV:
                if id_rsa_path:
                    conn = Connection(MACHINE_TO_HOST[machine], username, id_rsa_path=id_rsa_path)
                else:
                    conn = Connection(MACHINE_TO_HOST[machine], username)
        except BadHostKeyException:
            show_message('server’s host key could not be verified\ntry again')
            return
        except AuthenticationException:
            show_message('wrong username or password')
            return
        except error:
            show_message('check your internet connection :)')
            return
        except NoValidConnectionsError:
            show_message('something wrong with the server ¯\\_(ツ)_/¯')
            return
        except SSHException:
            show_message('i have no idea what happend\nplease contact @maryshca')
            return
        connections[machine] = conn
        set_run_window_value('connections', connections)
    else:
        conn = connections[machine]

    running_session_id = str(uuid4())

    if machine == LOMONOSOV:
        running_session_id = '~/_scratch/' + running_session_id

    add_running_session_ui(running_session_id, machine, model)

    if (machine, model) not in builds:
        dpg.hide_item(f'{running_session_id}_progress_bar')

        gitlab_username = dpg.get_value('gitlab_username')
        gitlab_password = dpg.get_value('gitlab_password')
        build_commands = BUILD_COMMANDS[(machine, model)]
        exec_file_path = EXECUTABLE_FILE_PATH[(machine, model)]
        building_session = BuildingSession(conn, build_commands, exec_file_path, gitlab_username, gitlab_password)

        while building_session.building():
            output = building_session.output()
            update_progress(running_session_id, output)

        executable_file_path = building_session.executable_file_path

        builds[(machine, model)] = executable_file_path
        set_run_window_value('builds', builds)

        dpg.show_item(f'{running_session_id}_progress_bar')
    else:
        executable_file_path = builds[(machine, model)]

    dpg.set_value(f'{running_session_id}_status', 'running')

    if machine == LOMONOSOV:
        run_commands = RUN_COMMANDS[(machine, model)] # + make_run_sh_for_slurm() + ['sbatch ./run.sh']
    else:
        run_commands = RUN_COMMANDS[(machine, model)] + [make_execute_command()]

    running_session = RunningSession(conn, config_file_path, run_commands, executable_file_path, running_session_id)

    if machine == LOMONOSOV:
        output = running_session.output()
        sbatch_id = re.findall(r'([0-1]+)', output)[0]
        return

    while running_session.running():
        output = running_session.output()
        update_progress(running_session_id, output)

    dpg.set_value(f'{running_session_id}_status', 'successful run')
    dpg.set_value(f'{running_session_id}_progress_bar', 1.0)
    dpg.enable_item(f'{running_session_id}_download_button')

def run_model_clb():
    threading.Thread(target=run_model).start()

def show_download_output_window_clb(s, a, running_session_id):
    delete_item_children_and_clear_aliases('output_header')

    machine = LABEL_TO_MACHINE[dpg.get_value(f'{running_session_id}_machine')]

    conn = get_run_window_value('connections')[machine]
    stdout = conn.exec_command(f'cd ./{running_session_id}/output; find . -type f')
    output = stdout.read().decode()
    dirs, filepaths = parse_dirs(output)
    construct_model_output_structure_ui(filepaths, dirs)

    dpg.set_item_user_data('model_output_window', (dirs, filepaths, running_session_id))

    dpg.show_item('model_output_window')

def open_download_to_folder_dialog_clb():
    dpg.show_item('download_to_folder_dialog')

def open_config_file_path_dialog_clb():
    dpg.show_item('config_file_path_dialog')

def set_download_output_to_folder_clb(s, a, u):
    dpg.set_value('download_output_to_folder', a['file_path_name'])

def set_config_file_path_clb(s, a, u):
    dpg.set_value('config_file_path', a['file_path_name'])

def download_output():
    dpg.hide_item('model_output_window')
    dirs, filepaths, running_session_id = dpg.get_item_user_data('model_output_window')

    download = []
    for filepath in filepaths:
        if dpg.get_value(filepath):
            download.append(filepath)

    machine = LABEL_TO_MACHINE[dpg.get_value(f'{running_session_id}_machine')]

    download_to_folder = dpg.get_value('download_output_to_folder')

    dpg.set_value(f'{running_session_id}_status', 'downloading')

    prefix = os.path.join(download_to_folder)
        
    construct_folder_structure(dirs, prefix=prefix)

    conn = get_run_window_value('connections')[machine]
    
    for filepath in download:
        conn.download(f'{running_session_id}/output/{filepath}', os.path.join(prefix, filepath))

    dpg.set_value(f'{running_session_id}_status', 'downloaded')

def download_output_clb(s, a, u):
    threading.Thread(target=download_output).start()