import gradio as gr
from typing import Dict
import asyncio

from src.control.control import Controller


def run(config: Dict, controller: Controller):
    with gr.Blocks() as formatdoc:
        with gr.Row():
            with gr.Column():
                pass

            with gr.Column(scale=10):
                """
                =====================================================
                Input and style components
                ==========================
                """
                input_files_comp = gr.File(file_count="multiple", file_types=[".docx"])

                with gr.Accordion("Modifier automatiquement les styles", open=False) as style_acc:
                    templates_radio = gr.Radio(
                        label="Templates",
                        choices=config['templates'],
                        value=config['templates'][config['default_template_index']],
                    )
                    options_btn = gr.CheckboxGroup(choices=config['options'],
                                                   label="Options",
                                                   info="(Fonctions non implémentées actuellement)",
                                                   interactive=True, )

                    with gr.Accordion("Mapper les styles qui n'existent pas dans le template", open=False) \
                            as newstyles_acc:
                        with gr.Column(scale=2):
                            newstyle_comps = [gr.Dropdown(visible=False, interactive=True)
                                              for _ in range(config['max_styles'])]

                    log_comp = gr.Textbox(label="Journal des modifications", visible=False)

                    output_styles_files_comp = gr.File(file_count="multiple", file_types=[".docx"], visible=False)

                    with gr.Row():
                        run_style_btn = gr.Button("Appliquer le template et les modifications de style")
                        clear_style_btn = gr.Button("Annuler les modifications de style")

                """
                ===============================================
                Generation components
                ======================
                """
                with gr.Accordion("Générer automatiquement une première version du document", open=False) as gen_acc:

                    generate_option_btn = gr.Radio(
                        label="Automatically generate a draft based on your own database",
                        choices=["Auto generation", "No generation"],
                        value="No generation",
                        interactive=True,
                        visible=False,
                    )

                    db_list_comp = gr.CheckboxGroup(
                        label="Base de connaissance",
                        info="Ces documents constituent la source de référence. Désélectionner pour qu'ils ne soient "
                             "pas pris en compte lors de la génération automatique",
                        visible=True,
                        interactive=True,
                    )
                    db_reset_btn = gr.Button("Effacer la base de connaissance", visible=False) \
                        .style(full_width=False, size="sm")
                    with gr.Accordion("Ajouter des documents dans la base de connaissance", open=False):
                        with gr.Column(visible=True, variant="panel") as add_col:
                            with gr.Tab("Depuis Wikipedia"):
                                wiki_fetch_btn = gr.Button("Rechercher les pages Wikipedia", visible=True)
                                wiki_fetch_btn.style(full_width=False, size="sm")
                                wiki_list_comp = gr.CheckboxGroup(
                                    label="Sélectionner les pages à ajouter dans la base de connaissance",
                                    visible=False,
                                    interactive=True,
                                )
                                wiki_add_to_db_btn = \
                                    gr.Button("Ajouter les documents sélectionnés à la base de connaissance",
                                              visible=False)
                                wiki_add_to_db_btn.style(full_width=False, size="sm")

                            with gr.Tab("Depuis le disque local"):
                                my_files_list_comp = gr.Files(
                                    label="Charger ses documents",
                                    info="Les documents fournissent le contexte utilisé pour la génération de texte",
                                    visible=True,
                                )
                                my_files_add_to_db_btn = gr.Button("Add files to sources", visible=False)
                                my_files_add_to_db_btn.style(full_width=False, size="sm")

                            add_close_btn = gr.Button("Close", visible=False).style(size='sm', full_width=False)
                    with gr.Row():
                        db_add_doc_btn = gr.Button("Ajouter de nouveaux documents", visible=False)\
                            .style(full_width=False, size="sm")

                    output_files_comp = gr.Files(file_count="multiple", visible=False)

                    generate_btn = gr.Button("Générer", interactive=True)

                    clear_btn = gr.Button('Nettoyer', visible=False)
                    rerun_btn = gr.Button('Relancer', visible=False)

            with gr.Column():
                pass

        """
        ===================================================
        state variables
        ===============
        """

        wiki_source_var: [str] = gr.State([])  # list of wikipage titles of interest for the input text tasks
        wiki_db_var: [str] = gr.State([])  # list of wiki document titles in the db (as seen from the UI)
        my_files_db_var: [str] = gr.State([])  # list of titles of the files uploaded in the db (as seen from the UI)
        db_collection_var: str = gr.State("-1")  # name of the collection of documents sources in the db

        """
        ===================================================
        Input and styles functions and listeners
        ========================================
        """

        def input_files_upload_fn(input_files_):
            controller.copy_docs(input_files_)
            update_ = {
                newstyles_acc: gr.update(open=False),
                style_acc: gr.update(open=False),
                run_style_btn: gr.update(visible=True),
            }
            newstyles_update = newstyles_fn()
            update_.update(newstyles_update)
            return update_

        input_files_comp.upload(input_files_upload_fn,
                                inputs=[input_files_comp],
                                outputs=[style_acc, newstyles_acc, run_style_btn] + newstyle_comps
                                )

        def input_file_clear_fn():
            controller.clear_docs()
            update_ = {
                options_btn: gr.update(value=[]),
                log_comp: gr.update(value="", visible=False),
                output_styles_files_comp: gr.update(value=[], visible=False),
                newstyles_acc: gr.update(open=False),
                style_acc: gr.update(open=False),
                gen_acc: gr.update(open=False),
                output_files_comp: gr.update(visible=False),
            }
            newstyles_update_ = newstyles_reset()
            update_.update(newstyles_update_)
            return update_

        input_files_comp.clear(
            input_file_clear_fn,
            inputs=[],
            outputs=[options_btn, output_styles_files_comp, output_files_comp, log_comp, newstyles_acc,
                     gen_acc, style_acc] + newstyle_comps
        )

        def newstyles_fn():
            different_styles, template_styles = controller.get_difference_with_template()
            update_ = {}
            get_label = lambda i: f"document: {different_styles[i]['doc'].name} style: {different_styles[i]['style']}"
            newstyles_update_ = {
                newstyle_comps[i]: gr.update(visible=i < len(different_styles),
                                             choices=template_styles,
                                             value=None,
                                             label=get_label(i)) if i < len(different_styles) else ''
                for i in range(config['max_styles'])
            }
            update_.update(newstyles_update_)
            return update_

        def newstyles_reset():
            update_ = {
                newstyle_comps[i]: gr.update(visible=False,
                                             choices=[],
                                             label='')
                for i in range(config['max_styles'])
            }
            return update_

        def templates_fn(templates_):
            controller.set_template(templates_)
            update_ = newstyles_fn()
            return update_

        templates_radio.change(templates_fn,
                               inputs=[templates_radio],
                               outputs=newstyle_comps)

        def newstyle_fns(src_index: int):
            def newstyle_fn(newstyle_):
                controller.update_style(src_index, newstyle_)
            return newstyle_fn

        for src_index, newstyle_comp in enumerate(newstyle_comps):
            newstyle_comp.input(newstyle_fns(src_index), inputs=[newstyle_comp], outputs=[])

        def clear_style_fn(input_files_):
            controller.clear_docs()
            if input_files_:
                controller.copy_docs(input_files_)
            controller.set_template()
            update_ = {
                options_btn: gr.update(value=[]),
                log_comp: gr.update(value="", visible=False),
                output_styles_files_comp: gr.update(value=[], visible=False),
                newstyles_acc: gr.update(open=False),
                run_style_btn: gr.update(visible=True),
            }
            newstyles_update_ = newstyles_fn()
            update_.update(newstyles_update_)
            return update_

        clear_style_btn.click(clear_style_fn,
                              inputs=[input_files_comp],
                              outputs=[options_btn, output_styles_files_comp, log_comp, newstyles_acc, run_style_btn]
                                      + newstyle_comps
                              )

        def run_style_fn():
            controller.apply_template(add_front_pages=False)
            log = controller.get_log()
            output_paths = [nd.path for nd in controller.new_docs]
            update_ = {
                log_comp: gr.update(value=log, visible=True),
                output_styles_files_comp: gr.update(value=output_paths, visible=True),
                run_style_btn: gr.update(visible=False),
            }
            return update_

        run_style_btn.click(run_style_fn,
                            inputs=[],
                            outputs=[log_comp, output_styles_files_comp, run_style_btn])

        """
        =====================================================
        Generation functions
        ====================
        """

        def generate_option_fn(db_collection_):
            id_ = controller.get_or_create_collection(db_collection_)
            update_ = {
                db_collection_var: id_
            }
            return update_

        def wiki_fetch1_fn():
            """
            fetch the wikifiles interesting for solving the tasks as defined in the input doc
            """
            update_ = {
                wiki_list_comp: gr.update(visible=True),
            }
            return update_

        def wiki_fetch2_fn():
            """
            fetch the wikifiles interesting for solving the tasks as defined in the input doc
            """
            wiki_interesting_files = controller.wiki_fetch()
            wiki_files = wiki_interesting_files  # [w for w in wiki_interesting_files if w not in wiki_db_files_]
            update_ = {
                wiki_list_comp: gr.update(visible=True, value=wiki_files, choices=wiki_files),
                wiki_source_var: wiki_interesting_files,
                wiki_add_to_db_btn: gr.update(visible=True),
            }
            return update_

        async def wiki_add_to_db_fn(wiki_list_, wiki_source_, wiki_db_, db_list_, db_collection_):
            """
            adds the wikipages to the db source
            """
            wiki_to_add = [wiki for wiki in wiki_list_ if wiki not in wiki_db_]
            db_list_ += wiki_to_add
            wiki_db_ += wiki_to_add
            wiki_source_remaining = [wiki for wiki in wiki_source_ if wiki not in wiki_db_]
            tasks = [controller.wiki_upload_and_store(wiki, db_collection_) for wiki in wiki_to_add]
            await asyncio.gather(*tasks)
            db_not_empty = 0 < len(db_list_)
            wiki_to_add_not_empty = 0 < len(wiki_source_remaining)
            update_ = {
                wiki_db_var: wiki_db_,
                wiki_list_comp: gr.update(value=wiki_source_remaining, choices=wiki_source_remaining),
                wiki_add_to_db_btn: gr.update(visible=wiki_to_add_not_empty),
                db_list_comp: gr.update(
                    visible=True,
                    value=db_list_,
                    choices=db_list_,
                    label="Database content"),
                db_reset_btn: gr.update(visible=db_not_empty),
                generate_btn: gr.update(visible=True, interactive=db_not_empty),
            }
            return update_

        def generate_fn1():
            update_ = {
                output_files_comp: gr.update(visible=True)
            }
            return update_

        async def generate_fn2(db_collection_, db_list_):
            output_files = controller.generate_doc_from_db(collection_name=db_collection_,
                                                           from_files=db_list_)
            update_ = {
                output_files_comp: gr.update(value=output_files, visible=True),
            }
            return update_


        """
        =====================================================
        Generation listeners
        ====================
        """

        wiki_fetch_btn \
            .click(wiki_fetch1_fn, inputs=[], outputs=[wiki_list_comp]) \
            .then(wiki_fetch2_fn,
                  inputs=[],
                  outputs=[wiki_list_comp, wiki_source_var, wiki_add_to_db_btn])

        wiki_add_to_db_btn\
            .click(generate_option_fn,
                                  inputs=[db_collection_var],
                                  outputs=[db_collection_var])\
            .then(wiki_add_to_db_fn,
                                 inputs=[wiki_list_comp, wiki_source_var, wiki_db_var, db_list_comp, db_collection_var],
                                 outputs=[db_list_comp, wiki_list_comp, wiki_db_var,
                                          generate_btn, wiki_add_to_db_btn, db_reset_btn])

        generate_btn\
            .click(generate_fn1,
                   inputs=[],
                   outputs=[output_files_comp])\
            .then(generate_fn2,
                   inputs=[db_collection_var, db_list_comp],
                   outputs=[output_files_comp])

        """
        =====================================================
        Clear and rerun functions and listeners
        =======================================
        """

        def clear_fn():
            update_ = {
                input_files_comp: gr.update(value=None),
                output_files_comp: gr.update(value=None, visible=False),
                clear_btn: gr.update(visible=False),
                rerun_btn: gr.update(visible=False),
            }
            return update_

        clear_btn.click(clear_fn,
                        inputs=[],
                        outputs=[input_files_comp, output_files_comp, clear_btn, rerun_btn])

    return formatdoc