import os
import sys
import json
import copy

import Orange.data
from Orange.data import Table, Domain, StringVariable
from AnyQt.QtWidgets import QApplication
from Orange.widgets import widget
from Orange.widgets.utils.signals import Input, Output
from AnyQt.QtCore import QTimer

if "site-packages/Orange/widgets" in os.path.dirname(os.path.abspath(__file__)).replace("\\", "/"):
    from Orange.widgets.orangecontrib.AAIT.llm import PageIndex_functions, prompt_management
    from Orange.widgets.orangecontrib.AAIT.llm.answers_llama import load_model, run_query, count_tokens, split_think, conversation_to_text
    from Orange.widgets.orangecontrib.AAIT.utils import thread_management, help_management
    from Orange.widgets.orangecontrib.AAIT.utils.initialize_from_ini import apply_modification_from_python_file
    from Orange.widgets.orangecontrib.AAIT.utils.import_uic import uic
else:
    from orangecontrib.AAIT.utils.import_uic import uic
    from orangecontrib.AAIT.llm import PageIndex_functions, prompt_management
    from orangecontrib.AAIT.llm.answers_llama import load_model, run_query, count_tokens, split_think, conversation_to_text
    from orangecontrib.AAIT.utils import thread_management, help_management
    from orangecontrib.AAIT.utils.initialize_from_ini import apply_modification_from_python_file


@apply_modification_from_python_file(filepath_original_widget=__file__)
class OWPageIndexExplore(widget.OWWidget):
    name = "PageIndex - Explore"
    description = "Look through your processed database with PageIndexes to answer a request"
    category = "AAIT - LLM INTEGRATION"
    icon = "icons/owpageindex_explore.svg"
    if "site-packages/Orange/widgets" in os.path.dirname(os.path.abspath(__file__)).replace("\\", "/"):
        icon = "icons_dev/owpageindex_explore.svg"
    gui = os.path.join(os.path.dirname(os.path.abspath(__file__)), "designer/owpageindex_explore.ui")
    want_control_area = False
    priority = 1060


    class Inputs:
        data = Input("Request Table", Orange.data.Table)
        path_table = Input("Path Table", Orange.data.Table)
        LLM = Input("LLM", str, auto_summary=False)

    class Outputs:
        data = Output("Data", Orange.data.Table)

    @Inputs.data
    def set_data(self, in_data):
        self.data = in_data
        if self.autorun:
            self.run()

    @Inputs.path_table
    def set_path_table(self, in_path_table):
        self.path_table = in_path_table
        if self.autorun:
            self.run()

    @Inputs.LLM
    def set_LLM(self, in_LLM_path):
        self.LLM_path = in_LLM_path
        if self.autorun:
            self.run()

    def __init__(self):
        super().__init__()
        # Qt Management
        self.setFixedWidth(598)
        self.setFixedHeight(471)
        uic.loadUi(self.gui, self)

        # Data Management
        self.data = None
        self.path_table = None
        self.LLM_path = None

        self.thread = None
        self.autorun = True
        self.result = None
        self.post_initialized()
        QTimer.singleShot(0, lambda: help_management.override_help_action(self))

    def run(self):
        self.error("")
        self.warning("")

        # if thread is running quit
        if self.thread is not None:
            self.thread.safe_quit()

        if self.data is None:
            self.Outputs.data.send(None)
            return

        if self.path_table is None:
            self.Outputs.data.send(None)
            return

        if self.LLM_path is None:
            self.Outputs.data.send(None)
            return

        if not os.path.exists(self.LLM_path):
            self.error(f'Your LLM path does not exist: {self.LLM_path}.')
            self.Outputs.data.send(None)
            return

        if not "path" in self.path_table.domain:
            self.warning(f'"path" variable required in your Path Table.')
            self.Outputs.data.send(None)
            return

        if not "request" in self.data.domain:
            self.warning(f'"request" variable required in your Request Table.')
            self.Outputs.data.send(None)
            return

        if "Answer" in self.data.domain:
            self.error(f'You must not have an "Answer" variable in your Data Table.')
            self.Outputs.data.send(None)
            return

        if "Conversation" in self.data.domain:
            self.error(f'You must not have an "Conversation" variable in your Data Table.')
            self.Outputs.data.send(None)
            return


        path = self.path_table[0]["path"].value
        if not os.path.isdir(path):
            self.error(f"Selected path is not a directory: {path}.")
            self.Outputs.data.send(None)
            return

        database_index_path = os.path.join(path, "database_index.json")
        if not os.path.exists(database_index_path):
            self.error(f"The Database Index does not exist: {database_index_path} not found.")
            self.Outputs.data.send(None)
            return

        with open(database_index_path, "r", encoding="utf-8") as f:
            database_index = json.load(f)

        # Start progress bar
        self.progressBarInit()

        # Connect and start thread : main function, progress, result and finish
        # --> progress is used in the main function to track progress (with a callback)
        # --> result is used to collect the result from main function
        # --> finish is just an empty signal to indicate that the thread is finished
        self.thread = thread_management.Thread(self.explore_page_index_for_table, self.data, database_index, self.LLM_path)
        self.thread.progress.connect(self.handle_progress)
        self.thread.result.connect(self.handle_result)
        self.thread.finish.connect(self.handle_finish)
        self.thread.start()

    def handle_progress(self, value: float) -> None:
        self.progressBarSet(value)

    def handle_result(self, result):
        try:
            self.result = result
            self.Outputs.data.send(result)
        except Exception as e:
            print("An error occurred when sending out_data:", e)
            self.Outputs.data.send(None)
            return

    def handle_finish(self):
        self.progressBarFinished()

    def post_initialized(self):
        pass



    def explore_page_index_for_table(self, table, database_index, LLM_path, progress_callback=None, argself=None):
        # Copie des données d'entrée
        data = copy.deepcopy(table)

        # Chargement modèle (llama_cpp)
        n_ctx = 32768
        LLM = load_model(model_path=LLM_path, use_gpu=True, n_ctx=n_ctx)

        # Génération sur la colonne "prompt"
        answers = []
        conversations = []
        print(data)
        for i, row in enumerate(data):
            request = row["request"].value
            try:
                answer, conversation = self.explore_page_index(question=request, database_index=database_index, LLM=LLM)
            except Exception as e:
                answer = f"An error happened: {e}"
                conversation = "Errors happened"
            answers.append(answer)
            conversations.append(conversation)


            if progress_callback is not None:
                progress_value = float(100 * (i + 1) / len(data))
                progress_callback(progress_value)
            if argself is not None:
                if argself.stop:
                    break

        # Add "Answer" column to the input data
        answer_var = StringVariable("Answer")
        conv_var = StringVariable("Conversation")
        # Add the scores as a meta attribute to the table
        data = data.add_column(answer_var, answers, to_metas=True)
        data = data.add_column(conv_var, conversations, to_metas=True)
        return data




    def explore_page_index(self, question, database_index, LLM, progress_callback=None, argself=None):
        # Parameters
        margin = 5000
        max_iter = 10

        # Default values
        final_answer = "Reached maximum number of iterations"
        path = None

        # Keep only names and summaries to show the LLM
        documents_summaries = []
        for entry in database_index:
            documents_summaries.append({"name": entry["name"], "summary": entry["summary"]})
        # Turn it into a string for the LLM
        documents_summaries = f"{documents_summaries}"

        # Verify length
        #TODO

        # Build prompt and conversation
        prompt = prompts["search_agent_3_tools"].format(QUESTION=question)
        conversation = [{"role": "user", "content": prompt}]

        # Start the search loop
        ephemeral_messages = []
        for step in range(max_iter):
            print(f"\n\n--- STEP {step} ---\n\n")

            # Apply Qwen template to the entire conversation
            prompt = prompt_management.apply_template_to_conversation("qwen", conversation)
            # Generate an answer
            answer = run_query(prompt=prompt, model=LLM, max_tokens=0, progress_callback=progress_callback, argself=argself)
            # Remove the Thinking
            answer = split_think(answer)[1]
            # Add the entry to the conversation
            conversation.append({"role": "assistant", "content": answer})

            # Detect if a tool has been called
            tool, args = PageIndex_functions.identify_tool_call(answer)
            if tool is not None:
                # Do animation for a tool !
                #TODO

                # Clear ephemeral messages to free some token memory
                PageIndex_functions.clear_ephemeral_messages(ephemeral_messages, conversation, tools=["view_documents", "get_pages"])


            if tool == "view_documents":
                self.label_1.setText("Analyse des documents disponibles...")
                tool_prompt = f"""# Appel outil
{tool}

# Résultat outil
{documents_summaries}"""

                conversation.append({"role": "user", "content": tool_prompt})
                ephemeral_message = {
                    "idx": len(conversation) - 1,
                    "tool": tool,
                    "replacement": f"# Appel outil\n{tool}\n\n# Résultat outil\n[Hidden for memory reason]"
                }
                ephemeral_messages.append(ephemeral_message)


            elif tool == "select_document":
                # Clear previously selected document in case another document has been selected
                PageIndex_functions.clear_ephemeral_messages(ephemeral_messages, conversation, tools=["select_document"])

                doc_name = args[0]
                doc_infos = PageIndex_functions.select_document(doc_name, database_index)
                self.label_1.setText(f"Sélection du document {doc_name}")
                if doc_infos == {}:
                    conversation.append({"role": "user", "content": f"# Appel outil\n{tool}\n\nNom du document incorrect"})
                    continue

                path = doc_infos["path"]
                name = doc_infos["name"]
                page_index = doc_infos["page_index"]
                tool_prompt = f"""# Appel outil
{tool}

# Résultat outil
## Document sélectionné
{name}
## Index
{page_index}"""

                conversation.append({"role": "user", "content": tool_prompt})
                ephemeral_message = {
                    "idx": len(conversation) - 1,
                    "tool": tool,
                    "replacement": f"# Appel outil\n{tool}\n\n# Résultat outil\n[Hidden for memory reason]"
                }
                ephemeral_messages.append(ephemeral_message)


            elif tool == "get_pages":
                if path is None:
                    conversation.append({"role": "user", "content": "Aucun document sélectionné."})
                    continue
                start_page = args[0]
                end_page = args[1]
                self.label_1.setText(f"Lecture des pages {start_page} - {end_page}")

                pages_text = PageIndex_functions.get_pages(path, start=start_page, end=end_page)
                tool_prompt = f"""# Appel outil
{tool}

# Résultat outil
{pages_text}"""

                conversation.append({"role": "user", "content": tool_prompt})
                ephemeral_message = {
                    "idx": len(conversation) - 1,
                    "tool": tool,
                    "replacement": f"# Appel outil\n{tool}\n\n# Résultat outil\n[Hidden for memory reason]"
                }
                ephemeral_messages.append(ephemeral_message)


            elif tool == "Unknown tool":
                tool_prompt = "The tool you tried to call has not been recognized."
                conversation.append({"role": "user", "content": tool_prompt})


            else:
                if "```final" in answer:
                    final_answer = answer.split("```final")[1].split("```")[0].strip()
                    break

                else:
                    # Log or append LLM text as fallback
                    conversation.append({"role": "user", "content": "[LLM output not recognized as tool or final answer, so continue. For final answer, use ```final\n<answer>\n```]"})
                    continue

        return final_answer, conversation_to_text(conversation)






prompts = {
    # {PAGE_INDEX}
    # {QUESTION}
    "search_agent": """# Contexte
Tu es un assistant avec accès à deux outils :
1. Un outil pour sélectionner un document
2. Un outil pour parcourir les pages d'un document sélectionné

Tu disposes d’un ensemble de documents, chacun décrit par un résumé.


# Règles
- Tu n'as PAS le droit d'inventer ou de deviner.
- Tu n'as PAS le droit de conclure que la réponse n'existe pas **sans** avoir lu de pages pertinentes avec l’outil.
- Tu peux utiliser les outils autant que tu veux jusqu'à avoir suffisamment d'informations.

---

# Comportement attendu
## Étape 1 — Sélection du document
- Si la question concerne les documents → tu DOIS d'abord sélectionner un document
- Choisis le document le plus pertinent en te basant sur les résumés fournis
- Tu ne peux pas appeler `get_pages` sans avoir sélectionné un document

## Étape 2 — Lecture des pages
- Une fois le document sélectionné → tu peux appeler `get_pages`

---

# Utilisation d'un outil
## Outil 1 - Sélection d'un document
Réponds avec

```tool
select_document(<pdf_name>)
```

puis attend le résultat avant de continuer tout réflexion ou rédaction de réponse. 
Tu vas recevoir un plan du document, te permettant de cibler les pages qui t'intéressent.


## Outil 2 - Lecture de pages
Réponds avec

```tool
get_pages(<page_debut>, <page_fin>)
```

puis attend le résultat avant de continuer toute réflexion ou rédaction de réponse.

---

# Arrêt
Lorsque tu as suffisamment d'informations pour répondre :
- Tu dois répondre directement à la question
- Tu ne dois PAS appeler l'outil
- Ta réponse doit être complète et précise

Format de réponse final :
```final
<ta réponse ici>
```

---

# Documents disponibles
{DOCUMENTS_SUMMARIES}


# Question
{QUESTION}""",


    # {QUESTION}
    "search_agent_3_tools": """# Contexte
Tu es un assistant avec accès à trois outils :
1. Un outil pour ouvrir voir l'ensemble des documents disponibles
2. Un outil pour sélectionner un document
3. Un outil pour parcourir les pages d'un document sélectionné


# Règles
- Tu n'as PAS le droit d'inventer ou de deviner.
- Tu peux utiliser les outils autant que tu veux jusqu'à avoir suffisamment d'informations.

---


# Comportement attendu
## Étape 1 - Visualisation des documents disponibles

## Étape 2 — Sélection d'un document
- Si la question concerne les documents → tu DOIS d'abord sélectionner un document
- Choisis le document le plus pertinent en te basant sur les résumés fournis
- Tu ne peux pas appeler `get_pages` sans avoir sélectionné un document

## Étape 3 — Lecture des pages
- Une fois le document sélectionné → tu peux appeler `get_pages`

---

# Utilisation d'un outil
Lorsque tu utilises un outil, respecte les 2 consignes suivantes :
- attend le résultat avant de continuer toute réflexion ou rédaction de réponse
- après avoir reçu le résultat, récapitule toute information pertinente que tu aurais reçue. S'il n'y en a pas, indique brièvement que ça ne t'a pas aidé.


## Outil 1 - Visualisation des documents disponibles
```tool
view_documents()
```

Tu vas recevoir un ensemble de documents avec leurs résumés.


## Outil 2 - Sélection d'un document
```tool
select_document(<pdf_name>)
```

Tu vas recevoir un plan du document, te permettant de cibler les pages qui t'intéressent.


## Outil 3 - Lecture de pages
```tool
get_pages(<page_debut>, <page_fin>)
```

Tu vas recevoir le contenu des pages demandées.


---


# Arrêt
Lorsque tu as suffisamment d'informations pour répondre :
- Tu dois répondre directement à la question
- Tu ne dois PAS appeler d'outil
- Ta réponse doit être complète et précise

Format de réponse final :
```final
<ta réponse ici>
```

---


# Question
{QUESTION}"""
}



if __name__ == "__main__":
    app = QApplication(sys.argv)
    my_widget = OWPageIndexExplore()

    my_widget.show()
    # if hasattr(app, "exec"):
    #     app.exec()
    # else:
    #     app.exec_()


    table = Table.from_list(Domain([], metas=[StringVariable("request")]), rows=[["Comment changer l'huile de turbine ?"]])
    LLM_path = r"C:\Users\lucas\AppData\Local\Programs\aait_store\Models\NLP\Qwen3.5-9B-Q6_K.gguf"
    database_path = r"C:\Users\lucas\Documents\Projets\IFIA\RAG\Dataset\Test_PageIndex"
    my_widget.data = table
    my_widget.LLM_path = LLM_path
    my_widget.path_table = Table.from_list(Domain([], metas=[StringVariable("path")]), rows=[[database_path]])

    my_widget.run()
