import datetime
from typing import Tuple, Optional

import gradio as gr

import mol_viewer
import run_utils


def run_wrapper(protein_file, ligand_file, other_args_file, *args, **kwargs) -> Tuple[str, Optional[str], str]:
    if protein_file is None:
        return "Protein file is missing! Must provide a protein file in PDB format", None, ""
    if ligand_file is None:
        return "Ligand file is missing! Must provide a ligand file in SDF format", None, ""

    if other_args_file is not None:
        kwargs["other_arg_file"] = other_args_file.name

    output_file, pdb_text, sdf_text = run_utils.run_cli_command(
        protein_file.name, ligand_file.name, *args, **kwargs
    )

    output_viz = "No visualisation created"
    if pdb_text:
        output_viz = mol_viewer.gen_3dmol_vis(pdb_text, sdf_text)

    message = f"Calculation completed at {datetime.datetime.now()}"
    return message, output_file, output_viz


def run():

    with gr.Blocks(title="DiffDock-Pocket Web") as demo:
        gr.Markdown("# DiffDock-Pocket")
        gr.Markdown("""Run [DiffDock-Pocket](https://github.com/plainerman/DiffDock-Pocket) for a single protein and ligand.
        We have provided the most important inputs as UI elements.  """)
        with gr.Row():
            protein_pdb = gr.File(label="Protein PDB (required)", file_types=[".pdb"])
            ligand_sdf = gr.File(label="Ligand SDF (required)", file_types=[".sdf"])
        with gr.Row():
            samples_per_complex = gr.Number(label="Samples Per Complex", value=1, minimum=1, maximum=100, precision=0)
            with gr.Column():
                keep_local_structures = gr.Checkbox(label="Keep Local Structures", value=True)
                save_vis = gr.Checkbox(label="Save Visualisation", value=True)
        with gr.Row():
            gr.Markdown("""Additional values can be included in "Other arguments", and will be passed 
        to [inference.py](https://github.com/plainerman/DiffDock-Pocket/blob/main/inference.py). 
        Must be a YAML file without any nesting.  """)
            other_args = gr.File(label="Other arguments (Optional, YML)", file_types=[".yml", ".yaml"], value=None)

        with gr.Row():
            run_btn = gr.Button("Run DiffDock-Pocket")

        with gr.Row():
            message = gr.Text(label="Run message", interactive=False)
        with gr.Row():
            output_file = gr.File(label="Output Files")
        with gr.Row():
            init_value = "Rank1 Reverse Processed Protein will be visualised here"
            viewer = gr.HTML(value=init_value, label="Protein Viewer", show_label=True)

        _inputs = [protein_pdb, ligand_sdf, other_args, samples_per_complex, keep_local_structures, save_vis]
        _outputs = [message, output_file, viewer]
        run_btn.click(fn=run_wrapper, inputs=_inputs, outputs=_outputs)

    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)


if __name__ == "__main__":
    run_utils.set_env_variables()
    run_utils.configure_logging()

    run()