import os
import re
import sys
import glob
import json
import shutil

import Orange.data
from Orange.data import Table, Domain, StringVariable, ContinuousVariable
from AnyQt.QtWidgets import QApplication
from Orange.widgets.settings import Setting
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
    from Orange.widgets.orangecontrib.AAIT.llm.answers_llama import load_model, load_Qwen3VL
    from Orange.widgets.orangecontrib.AAIT.utils import thread_management, base_widget, help_management
    from Orange.widgets.orangecontrib.AAIT.utils.initialize_from_ini import apply_modification_from_python_file
else:
    from orangecontrib.AAIT.llm import PageIndex_functions
    from orangecontrib.AAIT.llm.answers_llama import load_model, load_Qwen3VL
    from orangecontrib.AAIT.utils import thread_management, base_widget, 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 OWPageIndexCreate(base_widget.BaseListWidget):
    name = "PageIndex - Create"
    description = "Process all pdf files at a given location to create a PageIndex for each of them and a DatabaseIndex to store the results."
    category = "AAIT - LLM INTEGRATION"
    icon = "icons/owpageindex_create.svg"
    if "site-packages/Orange/widgets" in os.path.dirname(os.path.abspath(__file__)).replace("\\", "/"):
        icon = "icons_dev/owpageindex_create.svg"
    gui = os.path.join(os.path.dirname(os.path.abspath(__file__)), "designer/owpageindex_create.ui")
    want_control_area = False
    priority = 1060

    # Settings
    selected_column_name = Setting("path")

    class Inputs:
        data = Input("Data", Orange.data.Table)
        VLM = Input("VLM", str, auto_summary=False)
        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.data:
            self.var_selector.add_variables(self.data.domain)
            self.var_selector.select_variable_by_name(self.selected_column_name)
        if self.autorun:
            self.run()

    @Inputs.VLM
    def set_VLM(self, in_VLM_path):
        self.VLM_path = in_VLM_path
        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(470)
        self.setFixedHeight(395)
        #uic.loadUi(self.gui, self)

        # Data Management
        self.data = None
        self.VLM_path = 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.VLM_path 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.VLM_path):
            self.error(f'Your VLM path does not exist: {self.VLM_path}')
            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

        # Verification of selected column
        if not self.selected_column_name in self.data.domain:
            self.warning(f'Previously selected column "{self.selected_column_name}" does not exist in your data.')
            self.Outputs.data.send(None)
            return

        if not isinstance(self.data.domain[self.selected_column_name], StringVariable):
            self.error('You must select a text variable.')
            self.Outputs.data.send(None)
            return

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

        # 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.create_page_index, path, self.VLM_path, 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 create_page_index(self, path, VLM_path, LLM_path, progress_callback=None, argself=None):
        pdf_files = glob.glob(os.path.join(path, "*.pdf"))
        nb_files = len(pdf_files)
        self.label_1.setText(f"A total of {nb_files} documents have been detected.")
        count_errors = 0
        database_index = []
        for i, pdf_path in enumerate(pdf_files):

            pdf_name = os.path.basename(pdf_path)
            output_dir = re.sub(r"\.pdf", "_images", pdf_path, re.IGNORECASE)
            self.label_2.setText(f"Currently processing {pdf_name} - {i}/{nb_files}")

            try:
                # Step 1 - Extract PDF pages & convert them to images
                image_paths, text_per_page = PageIndex_functions.extract_pdf_pages(pdf_path=pdf_path, output_dir=output_dir)

                self.label_3.setText("Step 1/6 - Searching for the start of the document")
                # Step 2 - Find the start of the document
                VLM = load_Qwen3VL(VLM_path, verbose=False)
                start = PageIndex_functions.find_start_of_document(image_paths=image_paths, model=VLM)

                self.label_3.setText("Step 2/6 - Generating a global structure")
                # Step 3 - Go through pages to identify the structure of the document
                global_structure = PageIndex_functions.find_document_structure(image_paths=image_paths,
                                                                               start=start,
                                                                               model=VLM,
                                                                               batch_size=10,
                                                                               output_dir=output_dir,
                                                                               label=self.label_3)
                VLM.close()

                self.label_3.setText("Step 3/6 - Analysing and cleaning the table of contents")
                # Step 4 - Clear the TOC from duplicates and artifacts
                LLM = load_model(LLM_path, n_ctx=25000, use_gpu=True)
                clean_toc = PageIndex_functions.clear_toc(global_structure=global_structure, model=LLM, output_dir=output_dir)

                self.label_3.setText("Step 4/6 - Transforming the table of content into a JSON structure")
                # Step 5 - Transform the text TOC into a JSON formatted TOC
                json_toc = PageIndex_functions.toc_to_json(clean_toc=clean_toc, model=LLM, output_dir=output_dir)

                self.label_3.setText("Step 5/6 - Finalising the page index with page intervals")
                # Step 6 - Generate PageIndex
                page_index = PageIndex_functions.generate_PageIndex(json_toc=json_toc, text_per_page=text_per_page, output_dir=output_dir)

                self.label_3.setText("Step 6/6 - Generating a short summary to identify the document amongst others")
                # Step 7 - Generate a brief summary about the document, thanks to 1st page and table of content
                LLM.close()
                VLM = load_Qwen3VL(VLM_path, verbose=False)
                summary = PageIndex_functions.generate_summary_VLM(toc=clean_toc,
                                                                   pdf_name=pdf_name,
                                                                   image_path=[image_paths[0]],
                                                                   model=VLM,
                                                                   output_dir=output_dir)

                VLM.close()
            except Exception as e:
                summary = "An error happened !"
                page_index = e
                count_errors += 1


            pdf_data = {
                "id": i+1,
                "path": pdf_path,
                "name": pdf_name,
                "summary": summary,
                "page_index": f"{page_index}"
            }
            database_index.append(pdf_data)
            shutil.rmtree(re.sub(r"\.pdf", "_images", pdf_path, flags=re.IGNORECASE), ignore_errors=True)


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

        database_path = os.path.join(path, "database_index.json")
        with open(database_path, "w", encoding="utf-8") as f:
            json.dump(database_index, f, ensure_ascii=False, indent=2)

        self.label_2.setText(f"Preparation completed - {nb_files}/{nb_files}")
        self.label_3.setText(f"Your database can now be explored with the widget PageIndex - Explore")

        var_result_path = StringVariable("Result database")
        var_number_of_entries = ContinuousVariable("Number of documents")
        var_count_errors = ContinuousVariable("Number of errors")
        domain = Domain([], metas=[var_result_path, var_number_of_entries, var_count_errors])
        table = Table.from_list(domain, rows=[[database_path, len(database_index), count_errors]])
        return table





if __name__ == "__main__":
    app = QApplication(sys.argv)
    my_widget = OWPageIndexCreate()
    my_widget.show()
    if hasattr(app, "exec"):
        app.exec()
    else:
        app.exec_()
