from __future__ import annotations

import ast
import logging
import os
import re
import tempfile
from ast import literal_eval
from base64 import b64decode
from collections import OrderedDict
from importlib.util import find_spec
from io import BytesIO
from typing import TYPE_CHECKING, Any, TypeVar

import PIL.Image as Image
from typing_extensions import TypedDict

from vectice.api.http_error_handlers import VecticeException
from vectice.autolog.autolog_asset_factory import AssetFactory
from vectice.autolog.model_library import ModelLibrary
from vectice.models.dataset import Dataset
from vectice.models.table import Table
from vectice.utils.code_parser import VariableVisitor, preprocess_code

if TYPE_CHECKING:
    from tempfile import TemporaryDirectory

    from catboost.core import CatBoost
    from keras import Model as KerasModel
    from keras.layers import InputLayer
    from lightgbm.basic import Booster
    from pandas import DataFrame
    from pyspark.sql import DataFrame as SparkDF
    from sklearn.base import BaseEstimator

    # Vectice Object types
    from vectice.autolog.asset_services import TVecticeObjects
    from vectice.models import Phase
    from vectice.models.resource.metadata.db_metadata import TableType

    ModelTypes = TypeVar("ModelTypes", BaseEstimator, Booster, CatBoost, KerasModel)
    TModel = TypedDict("TModel", {"variable": str, "model": ModelTypes, "library": ModelLibrary, "metrics": dict})
    DataframeTypes = TypeVar("DataframeTypes", SparkDF, DataFrame)
    TDataset = TypedDict("TDataset", {"variable": str, "dataframe": DataframeTypes, "type": TableType})


try:
    from IPython.core.getipython import get_ipython
    from IPython.core.interactiveshell import InteractiveShell
except ModuleNotFoundError:
    raise ModuleNotFoundError(
        "To use autolog, please install extra dependencies from vectice package using '%pip install vectice[autolog]'"
    ) from None

is_plotly = False
is_matplotlib = False
if find_spec("plotly") is not None:
    is_plotly = True

if find_spec("matplotlib") is not None:
    is_matplotlib = True

_logger = logging.getLogger(__name__)

NOTEBOOK_CELLS = {}
GRAPHS = []
running = False


def _get_cell_id_and_plot_name(ipython: Any) -> tuple[str, str]:
    # Collab
    # cell_id = IP.get_parent()["metadata"]["colab"]["cell_id"]

    # Databricks
    # cell_id = IP.get_parent()["metadata"]["commandId"]
    cell_metadata = ipython.get_parent()["metadata"]  # type: ignore
    if cell_metadata.get("colab") and "cell_id" in cell_metadata.get("colab"):  # type: ignore
        cell_id = cell_metadata["colab"]["cell_id"]
        plot_name = f"vect_plot_{cell_id}"
    elif cell_metadata.get("commandId"):  # type: ignore
        cell_id = cell_metadata["commandId"]
        plot_name = f"vect_plot_{cell_id[-6:]}"
    else:
        cell_id = cell_metadata["cellId"]
        plot_name = (
            f"vect_plot_{cell_id.split('-')[-1]}"
            if "vscode-notebook-cell" not in cell_id
            else f"vect_plot_{cell_id.split('ipynb#')[-1]}"
        )
    return cell_id, plot_name


def start_listen():
    ipython = get_ipython()  # type: ignore

    def _notebook_cell_content_update():
        global NOTEBOOK_CELLS

        cell_id, _ = _get_cell_id_and_plot_name(ipython)
        cell_content = ipython.get_parent()["content"]["code"]  # type: ignore
        NOTEBOOK_CELLS[cell_id] = cell_content

    def _check_graph_libraries(cell_content: str) -> bool:
        # check for imports and recommended aliases
        if "matplotlib" in cell_content or "plt" in cell_content:
            return True
        if "seaborn" in cell_content or "sns" in cell_content:
            return True
        if "plotly" in cell_content or "px" in cell_content:
            _logger.warning("Plotly is not supported currently.")
            return False
        return False

    def _get_graph_path(cell_content: str, graph: str) -> str | None:
        # check if the `path` is a var or a graph and get the path
        # because it's pre_run_cell we can't get the var from locals and it's possible the graph doesn't exist yet
        match = re.search(rf"{graph}\s*=\s*(.*?)\n", cell_content)
        if match:
            graph = match.group(1)
        try:
            return literal_eval(graph)
        except ValueError as e:
            _logger.debug(f"Graph path failed {e}")
            return None

    def _get_graphs(remove: bool) -> list[str]:
        graphs = []
        cell_content = ipython.get_parent()["content"]["code"]  # type: ignore
        # get commented lines and not lines with comments
        lines = cell_content.split("\n")
        matched_lines = "\n".join(
            [line for line in lines if (remove and re.search(r"^#.*", line)) or not re.search(r"^#.*", line)]
        )

        # seaborn and matplotlib support
        graphs += re.findall(r"(?<=\.savefig\()[^)\n]*[^\)\n]", matched_lines)
        # plotly support
        graphs += re.findall(r"(?<=\.write_image\()[^)\n]*[^\)\n]", matched_lines)
        # ensure we have paths
        graph_paths = []
        for graph in graphs:
            graph_path = _get_graph_path(cell_content, graph)
            if graph_path:
                graph_paths.append(graph_path)
        return graph_paths

    def _remove_commented_graphs():
        # removes a saved graph if commented out
        graphs = _get_graphs(True)

        for graph in graphs:
            if graph in GRAPHS:
                GRAPHS.remove(graph)

    def _get_saved_graphs() -> list[str]:
        graphs = _get_graphs(False)

        for graph in graphs:
            if graph not in GRAPHS:
                GRAPHS.append(graph)
        return graphs

    def _capture_graph():
        # globals allows the function to get the variable, or it outside it's scope
        global GRAPHS
        # running variable prevents recursion
        global running

        cell_content = ipython.get_parent()["content"]["code"]  # type: ignore
        # we generate plot names that are then used in the cell magic
        _, plot_name = _get_cell_id_and_plot_name(ipython)
        # check if graphs have been commented out before getting saved graphs
        _remove_commented_graphs()
        # get the saved graphs
        saved_graphs = _get_saved_graphs()
        # check if running already and that there aren't saved graphs as saved graphs are the priority
        if not running and not saved_graphs and _check_graph_libraries(cell_content):
            running = True
            ipython.run_cell_magic("capture", plot_name, cell_content)  # type: ignore
            # Check if plot name is already in our list as we don't want duplicates
            if plot_name not in GRAPHS:
                GRAPHS.append(plot_name)
        # remove event graph if save graph true, prevents double capturing
        if saved_graphs and plot_name in GRAPHS:
            GRAPHS.remove(plot_name)
        running = False

    def _load_ipython_extension():
        """Load the extension in IPython."""
        ipython.events.register("pre_execute", _capture_graph)  # type: ignore
        ipython.events.register("post_execute", _notebook_cell_content_update)  # type: ignore

    if ipython and not any([ev for ev in ipython.events.callbacks["post_execute"] if "_notebook_cell_content_update" in str(ev)]):  # type: ignore
        _load_ipython_extension()


####### Autolog logic
class Autolog:
    def __init__(
        self,
        phase: Phase | None,
        ipy: InteractiveShell,
        is_notebook: bool,
        create: bool = True,
        note: str | None = None,
        capture_schema_only: bool = True,
    ):
        from vectice.services.phase_service import PhaseService

        if phase is None:
            raise ValueError("Login")

        if create is True:
            iteration = phase.create_iteration()
        else:
            iteration = PhaseService(
                phase._client  # pyright: ignore[reportPrivateUsage]
            ).get_active_iteration_or_create(phase)

        self._iteration = iteration
        self._iteration_service = iteration._service  # pyright: ignore[reportPrivateUsage]
        self._capture_schema_only = capture_schema_only
        self._ip = ipy
        self._local_vars = self._ip.user_global_ns
        self._cell_content = self._get_notebook_cell_content() if is_notebook is True else self._get_cell_content()
        self._vectice_data = self._get_variable_matches(self._local_vars)
        self._failed_assets = []

        # Get back objects to log
        self._assets = self._get_assets()
        graphs = self._get_graphs(is_notebook)

        if note:
            self._iteration_service.log_comment(note)
            _logger.info(f"Note logged in iteration {self._iteration.index!r}.")

        # Log objects
        self._log_assets(self._assets)

        if graphs:
            for graph in graphs:
                try:
                    self._iteration_service.log_image_or_file(graph)
                    graph = graph if isinstance(graph, str) else graph.filename
                    _logger.info(f"Graph {graph!r} logged in iteration {self._iteration.index!r}.")
                except Exception as e:
                    self._failed_assets.append({"reason": e, "asset": graph, "type": "Graph"})

        if len(self._failed_assets):
            _logger.warning("The following assets failed to log:")
        for failed_asset in self._failed_assets:
            _logger.warning(f"{failed_asset['type']} {failed_asset['asset']!r}, reason: {failed_asset['reason']}")

    def _get_variable_matches(self, local_vars: dict[str, Any]) -> list[dict[Any, Any]]:
        vectice_data = []
        all_vectice_calls = set()
        for cell_not_processed in self._cell_content:
            cell = preprocess_code(cell_not_processed)
            vectice_match = {}
            all_vars: OrderedDict[str, Any] = OrderedDict()

            vectice_match["cell"] = cell

            # Parse the cell content using the VariableVisitor
            try:
                tree = ast.parse(cell)
            except SyntaxError as err:
                raise SyntaxError(
                    "Autolog is unable to parse the code. Make sure all non-Python syntax, such as bash commands, are properly indented and preceded by '%' or '!'. If the error persist, please contact your sales representative."
                ) from err

            visitor = VariableVisitor()
            visitor.visit(tree)
            # Keep set of all vectice call vars
            all_vectice_calls = all_vectice_calls.union(visitor.vectice_call_vars)

            # Update all_vars with variable names and values in the order they appear
            for var in visitor.variables:
                if var in local_vars and var not in all_vars:
                    all_vars[var] = local_vars[var]

            vectice_match["variables"] = all_vars
            vectice_data.append(vectice_match)
        clean_vectice_data = []
        for data in vectice_data:
            for var in all_vectice_calls:
                # the order of cells misses vars
                if var in data["variables"]:
                    del data["variables"][var]
            # check that there are vars for the cell and then append
            if len(data["variables"]) >= 1:
                clean_vectice_data.append(data)
        return clean_vectice_data

    def _get_notebook_cell_content(self) -> list[Any]:
        return list(NOTEBOOK_CELLS.values())

    def _get_cell_content(self) -> list[Any]:
        """Used by autolog cell to get the content of the cell. This is used to parse for variables."""
        cell_content = self._ip.get_parent()["content"]["code"]  # pyright: ignore[reportAttributeAccessIssue]
        if cell_content is None:
            raise ValueError("Failed to get cell content.")
        return [cell_content]

    def _get_assets(self) -> list[TModel | TDataset | TVecticeObjects]:
        assets = []
        for data in self._vectice_data:
            for key in data["variables"].keys():
                # skip cell inputs/outputs
                if key.startswith("_"):
                    continue
                asset = data["variables"][key]
                try:
                    asset_information = AssetFactory.get_asset_service(key, asset, data).get_asset()
                    if asset_information is not None:
                        assets.append(asset_information)
                except VecticeException:
                    pass
        return self._deduplicate_assets(assets)

    def _deduplicate_assets(self, assets: list[dict[str, Any]]) -> list[TModel | TDataset | TVecticeObjects]:
        unique_dict = {}
        result = []

        for item in assets:
            variable = item["variable"]

            if variable not in unique_dict.keys():
                # If variable is not seen before, add it to the dictionary
                unique_dict[variable] = item
                result.append(item)
            else:
                existing_item = unique_dict[variable]
                if "vectice_object" in item:
                    existing_item["vectice_object"] = item["vectice_object"]
                elif "dataframe" in item:
                    existing_item["type"] = item["type"]
                else:
                    existing_item["library"] = item["library"]

                    # Merge dictionaries for the "metrics" attribute
                    existing_item_metrics = existing_item["metrics"]
                    item_metrics = item["metrics"]

                    for key, value in item_metrics.items():
                        existing_item_metrics[key] = value

        return result

    def _get_lightgbm_info(self, model: Booster) -> tuple[str, dict[str, Any]] | tuple[None, None]:
        try:
            params = {
                key: value
                for key, value in model._get_loaded_param().items()  # pyright: ignore[reportPrivateUsage]
                if value is not None
            }
            return "lightgbm", params
        except AttributeError:
            return None, None

    def _get_model_library(self, model: ModelTypes):
        if "xgboost" in str(model.__class__):
            return "xgboost"
        if "lightgbm" in str(model.__class__):
            return "lightgbm"
        return "sklearn"

    def _get_sklearn_or_xgboost_or_lgbm_info(
        self, model: BaseEstimator
    ) -> tuple[str, dict[str, Any]] | tuple[None, None]:
        try:
            library = self._get_model_library(model)
            params = {str(key): value for key, value in model.get_params().items() if value is not None}
            return library, params
        except AttributeError:
            return None, None

    def _format_keras_params(self, model: KerasModel) -> dict[str, Any]:
        params: dict[str, Any] = {}

        def _get_output_shape(layer: InputLayer) -> tuple:
            try:
                # Keras 2.15.1 & below
                return layer.output_shape
            except AttributeError:
                return layer.output.shape  # pyright: ignore[reportAttributeAccessIssue]

        for i, layer in enumerate(model.layers):
            output_shape = _get_output_shape(layer)
            params[f"Layer-{i}"] = {
                "name": layer.name,
                "param": layer.count_params(),
                "output shape": output_shape,
            }
        params["Total # of weights"] = model.count_params()

        return params

    def _get_keras_info(self, model: KerasModel) -> tuple[str, dict[str, Any]] | tuple[None, None]:
        try:
            return "keras", self._format_keras_params(model)
        except Exception:
            return None, None

    def _log_assets(
        self,
        assets: list[TModel | TDataset | TVecticeObjects],
    ):
        for asset in assets:
            try:
                if "vectice_object" in asset:
                    self._log_vectice_asset(asset)
                elif "model" in asset:
                    self._log_model(asset)
                elif "dataframe" in asset:
                    self._log_dataset(asset)
            except Exception as e:
                self._failed_assets.append(
                    {"reason": e, "asset": asset["variable"], "type": "Model" if "model" in asset else "Dataset"}
                )

    def _get_catboost_info(self, model: CatBoost) -> tuple[str, dict[str, Any]] | tuple[None, None]:
        try:
            params = {str(key): value for key, value in model.get_all_params().items() if value is not None}
            return "catboost", params
        except AttributeError:
            return None, None

    def _get_model_params(self, asset: TModel) -> tuple[str, dict[str, Any]] | tuple[None, None]:
        library = asset["library"]
        model = asset["model"]

        if library is ModelLibrary.SKLEARN:
            return self._get_sklearn_or_xgboost_or_lgbm_info(model)  # pyright: ignore[reportArgumentType]
        if library is ModelLibrary.LIGHTGBM:
            return self._get_lightgbm_info(model)  # pyright: ignore[reportArgumentType]
        if library is ModelLibrary.CATBOOST:
            return self._get_catboost_info(model)  # pyright: ignore[reportArgumentType]
        if library is ModelLibrary.KERAS:
            return self._get_keras_info(model)  # pyright: ignore[reportArgumentType]
        return None, None

    def _check_for_single_model(self):
        model_list = [asset for asset in self._assets if "model" in asset]
        if len(model_list) > 1:
            return False
        return True

    def _get_single_model_metrics(self) -> dict:
        from vectice.autolog.asset_services.metric_service import MetricService

        # get a single models metrics for notebook function
        metrics = {}
        if not self._check_for_single_model():
            return {}
        for data in self._vectice_data:
            cell_metrics = MetricService(data["cell"])._get_model_metrics(data)  # type: ignore[reportPrivateUsage]
            if cell_metrics:
                metrics.update(cell_metrics)
        return metrics

    def _log_model(self, model: TModel):
        from vectice import Model

        temp_dir: TemporaryDirectory | None = None
        temp_file_path: str | None = None
        if model["library"] is ModelLibrary.KERAS:
            graph = None

            try:
                from keras.utils import plot_model

                temp_dir = tempfile.TemporaryDirectory()
                file_name = f"{model['variable']!s}_plot.png"
                temp_file_path = rf"{temp_dir.name}\{file_name}"
                graph = plot_model(model["model"], to_file=temp_file_path, show_shapes=True, show_layer_names=False)

                if graph is None:
                    temp_file_path = None
                    _logger.info(
                        "Unable to generate the model plot. Please check the 'plot_model' function from the Keras library is working correctly. Make sure that the graphviz and pydot packages are installed and configured properly."
                    )

            except Exception as e:
                _logger.info(
                    f"Unable to generate the model plot. Ensure that the Keras library is correctly installed and up-to-date. Error details: {e}"
                )

        library, params = self._get_model_params(model)

        algorithm = str(model["model"].__class__).split(".")[-1]
        algorithm = algorithm.replace("'>", "")
        model_name = f"{self._iteration.phase.id}-{model['variable']}"
        # safety check, if no metrics and single model, get metrics
        model_metrics = model["metrics"] if model["metrics"] else self._get_single_model_metrics()
        self._iteration_service.log_model(
            Model(
                library=library,
                technique=algorithm,
                metrics=model_metrics,
                properties=params,
                name=model_name,
                predictor=model["model"],
                attachments=temp_file_path,
            )
        )
        _logger.info(f"Model {model_name!r} logged in iteration {self._iteration.index!r}.")

        if temp_dir is not None:
            temp_dir.cleanup()

    def _log_vectice_asset(self, asset: TVecticeObjects) -> None:
        asset_to_log = asset["vectice_object"]
        asset_name = asset["variable"]
        if isinstance(asset_to_log, Dataset):
            dataset = asset_to_log
            dataset_name = dataset.name if dataset.name else asset_name
            self._iteration_service.log_dataset(asset_to_log)
            _logger.info(f"Dataset {dataset_name!r} logged in iteration {self._iteration.index!r}.")
        elif isinstance(asset_to_log, Table):
            table = asset_to_log
            table_name = table.name if table.name else asset_name
            self._iteration_service.log_table(asset_to_log)
            _logger.info(f"Table {table_name!r} logged in iteration {self._iteration.index!r}.")
        else:
            model = asset_to_log
            model_name = model.name if model.name else asset_name
            self._iteration_service.log_model(asset_to_log)
            _logger.info(f"Model {model_name!r} logged in iteration {self._iteration.index!r}.")

    def _log_dataset(self, dataset: TDataset) -> None:
        from vectice import Dataset, DatasetType, FileResource, NoResource

        resource = self._get_dataset_resource(dataset)
        dataset_name = f"{self._iteration.phase.id}-{dataset['variable']}"
        no_resource_dataset = Dataset(
            type=DatasetType.UNKNOWN,
            resource=NoResource(
                dataframes=dataset["dataframe"],
                origin="DATAFRAME",
                type=dataset["type"],
                capture_schema_only=self._capture_schema_only,
            ),
            name=dataset_name,
        )
        if resource:
            # TODO Dataset type ?
            vec_dataset = Dataset(
                type=DatasetType.UNKNOWN,
                resource=FileResource(
                    paths=resource,
                    dataframes=dataset["dataframe"],
                    capture_schema_only=self._capture_schema_only,
                ),
                name=dataset_name,
            )
        else:
            vec_dataset = no_resource_dataset
        try:
            self._iteration_service.log_dataset(vec_dataset)
        except FileNotFoundError:
            self._iteration_service.log_dataset(no_resource_dataset)
        _logger.info(f"Dataset {dataset_name!r} logged in iteration {self._iteration.index!r}.")

    def _get_dataset_resource(self, dataset: TDataset) -> str | None:
        import re

        if not self._cell_content:
            return None
        try:
            # Avoid getting stray dataset with autolog.notebook()
            match = []
            for cell in self._cell_content:
                match += re.findall(rf"(?<={dataset['variable']} = pd.read_csv\().*[^,)\n]", cell)
            if len(match) < 1:
                return None
            # TODO update regex
            # check if read csv has comma dominated arguments
            return literal_eval(match[0].split(",")[0]) if match[0].find(",") else literal_eval(match[0])
        except TypeError:
            return None

    def _get_image_data(self, captured: Any) -> str | None:
        try:
            return captured.outputs[0].data["image/png"]
        except IndexError:
            return None

    def _get_graphs(self, is_notebook: bool) -> list[Any]:
        global GRAPHS

        _, plot_name = _get_cell_id_and_plot_name(self._ip)

        graphs = []
        # Create png data and don't convert saved graphs
        graphs_data = [
            (graph, self._get_image_data(self._local_vars[graph]))
            for graph in GRAPHS
            if not os.path.isfile(graph) and self._get_image_data(self._local_vars[graph])
        ]
        # Create the Images from the png_data from the cell magic
        for graph_name, png_bytes in graphs_data:
            if isinstance(png_bytes, str):
                png_bytes = b64decode(png_bytes)
                bytes_io = BytesIO(png_bytes)
                image = Image.open(bytes_io)
                image.filename = f"{graph_name}.png"  # type: ignore
                # Check if it is cell event graph and make sure graph belongs to cell with the plot_name /  if notebook get all graph png data. Else statement causes duplicates
                if (not is_notebook and plot_name == graph_name) or is_notebook:
                    graphs.append(image)

        for graph in GRAPHS:
            is_str_graph = isinstance(graph, str)
            is_saved_graph_in_cell = "savefig" in str(self._cell_content) or "write_image" in str(self._cell_content)
            is_cell_condition = not is_notebook and graph in str(self._cell_content) and is_saved_graph_in_cell
            is_notebook_condition = is_notebook and os.path.isfile(graph)
            # check for cell saved graphs, check if graph is in cell as multiple cells can be run / check for notebook saved graphs, filepath check makes sure it's a saved graph and not a plot_name
            if is_str_graph and (is_cell_condition or is_notebook_condition):
                graphs.append(graph)
        return graphs
