from typing import Optional

from shimoku.api.resources.app import App
from shimoku.api.resources.data_set import DataSet
from shimoku.api.resources.report import Report
from shimoku.api.resources.reports.modal import Modal
from shimoku.api.resources.reports.tabs_group import TabsGroup
from shimoku.api.utils import change_data_set_name_with_report

from shimoku.code_gen.file_generator import CodeGenFileHandler

import tqdm
import pandas as pd
import asyncio

import logging
from shimoku.execution_logger import ClassWithLogging

logger = logging.getLogger(__name__)


class CodeGenTree(ClassWithLogging):
    """Class for generating code generation tree."""

    _module_logger = logger

    def __init__(
        self,
        app: App,
        file_generator: CodeGenFileHandler,
        pbar: Optional[tqdm.tqdm] = None,
    ):
        self._app = app
        self._file_generator = file_generator
        self._pbar = pbar
        self.custom_data_sets_with_data = {}
        self.actual_bentobox: Optional[dict] = None
        self.all_tab_groups: dict[str, list[dict]] = {}
        self.tree: Optional[dict] = None
        self.shared_data_sets: list[str] = []
        self.needs_pandas: bool = False

    def _update_pbar(self, n: int = 1):
        if self._pbar is not None:
            self._pbar.update(n)

    async def _check_for_shared_data_sets(
        self,
        report: Report,
        seen_data_sets: set,
        individual_data_sets: dict[str, Report],
    ):
        """Check for shared data sets in a report.
        :param report: report to check
        :param seen_data_sets: set of data sets already seen
        :param individual_data_sets: list of data sets to append to
        """
        report_data_sets: list[
            Report.ReportDataSet
        ] = await report.get_report_data_sets()
        data_sets_from_rds = set([rds["dataSetId"] for rds in report_data_sets])

        for ds_id in data_sets_from_rds:
            if ds_id in self.shared_data_sets:
                continue
            if ds_id in seen_data_sets:
                self.shared_data_sets.append(ds_id)
                del individual_data_sets[ds_id]
            else:
                individual_data_sets[ds_id] = report
                seen_data_sets.add(ds_id)

    async def _check_for_create_data_set(
        self, data_set: DataSet, report: Optional[Report] = None
    ):
        """Create a file for a data set.
        :param data_set: data set to create file for
        :param report: report to create file for
        """
        if data_set["name"] is None:
            data_set["name"] = data_set["id"]

        data: list[dict] = [
            {
                k: v
                for k, v in dp.cascade_to_dict().items()
                if k not in ["id", "dataSetId"] and v is not None
            }
            for dp in await data_set.get_data_points()
        ]

        if len(data) == 0:
            return

        data_as_df = pd.DataFrame(data)

        if len(data) > 1 or "customField1" not in data[0]:
            output_name = (
                data_set["name"]
                if report is None
                else change_data_set_name_with_report(data_set, report)
            )
            reverse_columns = (
                {v: k for k, v in data_set["columns"].items()}
                if data_set["columns"]
                else None
            )
            self._file_generator.create_data_frame_file(
                output_name, data_as_df, reverse_columns
            )
            self.needs_pandas = True
        else:
            self.custom_data_sets_with_data[data_set["id"]] = data[0]["customField1"]

    async def _get_data_sets(self):
        """Create files for the data sets."""
        reports = await self._app.get_reports()

        # To store the data sets in the cache for the reports to have faster access to them
        await self._app.get_data_sets()

        individual_data_sets: dict[str, Report] = {}
        seen_data_sets = set()

        await asyncio.gather(
            *[
                self._check_for_shared_data_sets(
                    report, seen_data_sets, individual_data_sets
                )
                for report in reports
            ]
        )
        tasks = []
        for ds_id in self.shared_data_sets:
            tasks.append(
                self._check_for_create_data_set(await self._app.get_data_set(ds_id))
            )
        for ds_id, report in individual_data_sets.items():
            tasks.append(
                self._check_for_create_data_set(
                    await self._app.get_data_set(ds_id), report
                )
            )
        await asyncio.gather(*tasks)

    async def _tree_from_tabs_group(
        self,
        tree: list,
        tabs_group: TabsGroup,
        parent_tabs_index: tuple[str, str] = None,
    ):
        """Recursively build a tree of reports from a tabs group.
        :param tree: list to append to
        :param tabs_group: tabs group to build tree from
        :param parent_tabs_index: tuple of (tabs_group hash, tab name) of parent tabs group
        """
        path = tabs_group["path"] + "_" if tabs_group["path"] else ""
        name = tabs_group["properties"]["hash"][len(path) :]
        tabs_group_dict = {
            "tabs_group": tabs_group,
            "tabs": {},
            "order": tabs_group["order"],
            "parent_tabs_index": parent_tabs_index,
            "name": name,
        }

        path = tabs_group["path"]
        if path not in self.all_tab_groups:
            self.all_tab_groups[path] = []
        self.all_tab_groups[path].insert(0, tabs_group_dict)

        tree.append(tabs_group_dict)
        tabs = sorted(
            tabs_group["properties"]["tabs"].items(), key=lambda x: x[1]["order"]
        )
        for tab, tab_data in tabs:
            report_ids = tab_data["reportIds"]
            tab_dict = {"tab_groups": [], "other": []}
            tabs_group_dict["tabs"][tab] = tab_dict
            for child_id in report_ids:
                child_report = await self._app.get_report(child_id)
                if child_report["reportType"] == "TABS":
                    await self._tree_from_tabs_group(
                        tab_dict["tab_groups"],
                        child_report,
                        (tabs_group["properties"]["hash"], tab),
                    )
                else:
                    tab_dict["other"].append(child_report)
                self._update_pbar()

            tab_dict["tab_groups"] = sorted(
                tab_dict["tab_groups"], key=lambda x: x["tabs_group"]["order"]
            )
            tab_dict["other"] = sorted(tab_dict["other"], key=lambda x: x["order"])

    async def _tree_from_modal(self, tree: list, modal: Modal):
        """Recursively build a tree of reports from a modal.
        :param self: PlotApi instance
        :param tree: list to append to
        :param modal: modal to build tree from
        """
        path = modal["path"] + "_" if modal["path"] else ""
        name = modal["properties"]["hash"][len(path) :]
        modal_dict = {"modal": modal, "tab_groups": [], "other": [], "name": name}
        tree.append(modal_dict)
        for child_id in modal["properties"]["reportIds"]:
            child_report = await self._app.get_report(child_id)
            if child_report["reportType"] == "TABS":
                await self._tree_from_tabs_group(modal_dict["tab_groups"], child_report)
            else:
                modal_dict["other"].append(child_report)
            self._update_pbar()

        modal_dict["tab_groups"] = sorted(
            modal_dict["tab_groups"], key=lambda x: x["tabs_group"]["order"]
        )
        modal_dict["other"] = sorted(modal_dict["other"], key=lambda x: x["order"])

    async def generate_tree(self):
        """Generate a tree of reports from a list of reports.
        :param self: PlotApi instances
        """
        self.tree = {}
        reports = await self._app.get_reports()
        contained_reports = set()
        for report in reports:
            # TODO: make this not necessary
            if "hash" not in report["properties"]:
                report["properties"]["hash"] = "id_" + report["id"]
            if report["reportType"] == "MODAL":
                contained_reports.update(report["properties"]["reportIds"])
            elif report["reportType"] == "TABS":
                for tab in report["properties"]["tabs"]:
                    contained_reports.update(
                        report["properties"]["tabs"][tab]["reportIds"]
                    )
        reports = sorted(
            reports,
            key=lambda x: (
                x["pathOrder"] if x["pathOrder"] else 0,
                "0"
                if x["reportType"] == "MODAL"
                else "1"
                if x["reportType"] == "TABS"
                else "_" + (x["properties"]["hash"]),
            ),
        )
        for report in reports:
            if report["id"] in contained_reports:
                continue

            if report["path"] not in self.tree:
                self.tree[report["path"]] = {
                    "modals": [],
                    "tab_groups": [],
                    "other": [],
                }

            if report["reportType"] == "MODAL":
                await self._tree_from_modal(self.tree[report["path"]]["modals"], report)
            elif report["reportType"] == "TABS":
                await self._tree_from_tabs_group(
                    self.tree[report["path"]]["tab_groups"], report
                )
            else:
                self.tree[report["path"]]["other"].append(report)

            self._update_pbar()

        for path in self.tree:
            self.tree[path]["other"] = sorted(
                self.tree[path]["other"], key=lambda x: x["order"]
            )
            self.tree[path]["tab_groups"] = sorted(
                self.tree[path]["tab_groups"], key=lambda x: x["tabs_group"]["order"]
            )

        # Todo: Solve path ordering
        # self.tree = {k: v for k, v in sorted(reports.items(), key=lambda item: }
        await self._get_data_sets()
