# SPDX-FileCopyrightText: 2025 Deutsche Telekom AG (opensource@telekom.de)
#
# SPDX-License-Identifier: Apache-2.0

"""Base Executor."""

import json
import os
import time
from collections.abc import Generator
from contextlib import contextmanager
from contextvars import copy_context
from logging import getLogger
from pathlib import Path
from types import NoneType
from typing import Any, Callable, Optional, Self, TypeAlias, Union

import pandas
import pandera.typing as patyp
import pydantic
from asgi_correlation_id import correlation_id
from pydantic_core import ValidationError

from wurzel.datacontract import PanderaDataFrameModel, PydanticModel, datacontract
from wurzel.exceptions import (
    ContractFailedException,
    EnvSettingsError,
    LoggedCustomException,
    StepFailed,
)
from wurzel.path import PathToFolderWithBaseModels
from wurzel.step.history import (
    History,
    step_history,
)
from wurzel.step.self_consuming_step import SelfConsumingLeafStep
from wurzel.step.typed_step import (
    TypedStep,
)
from wurzel.utils import WZ, create_model, try_get_length

log = getLogger(__name__)

StepReturnType: TypeAlias = Union[pandas.DataFrame, PydanticModel, list[PydanticModel]]


class StepReport(pydantic.BaseModel):
    """Generated by TypedStep Class."""

    model_config = pydantic.ConfigDict(frozen=True)
    inputs: int = 0
    results: int = 1
    time_to_load: float
    time_to_execute: float
    time_to_save: float
    step_name: str
    history: list[str]


def _try_sort(x: StepReturnType) -> StepReturnType:
    """Tries to sort given x.

    Returns either a sorted x or x itself
    """
    _log_extra = {"extra": {"type": type(x)}}
    if isinstance(x, PydanticModel):
        return x
    try:
        if isinstance(x, (list, set)):
            return sorted(x)
        if isinstance(x, pandas.DataFrame):
            return x.sort_values(x.columns[0])
    # pylint: disable-next=bare-except
    except:  # noqa: E722
        log.warning("Could not sort output", **_log_extra)
        return x
    log.warning("Can't sort objects of this type", **_log_extra)
    return x


@contextmanager
def step_env_encapsulation(step_cls: type[TypedStep]):
    """Create a virtual env and remove it afterwards."""
    old = os.environ.copy()
    name = step_cls.__name__.upper()
    step = WZ(step_cls)
    if BaseStepExecutor.is_allow_extra_settings():
        log.info("Allowing extra settings fields")
    if step.settings_class != NoneType:
        settings = create_model([step], allow_extra_fields=BaseStepExecutor.is_allow_extra_settings())
        try:
            inner_settings: pydantic.BaseModel = getattr(settings(), name)
        except ValidationError as err:
            e = EnvSettingsError(f"could not create {step.settings_class.__name__} from env for {step_cls.__name__}")
            e.add_note("To fix these issues setup env vars in the format <step_name>__<var>")
            raise e from err
        settings_dict = pydantic.BaseModel.model_dump(inner_settings, mode="json", exclude_none=True, exclude_unset=True)
        for field_name, value in settings_dict.items():
            if isinstance(value, (list, dict, set, tuple)):
                value = json.dumps(value)
            else:
                value = str(value)
            os.environ[field_name] = value
        log.info(
            "Running with encapsulated settings",
            extra={"settings": inner_settings.model_dump(mode="json")},
        )
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(old)


class BaseStepExecutor:
    """Execution Environment for Typed KP Steps.

    An Executer takes care of loading & storing data from `PathToBaseModel`.
    The Executor makes use of strong & enforced types in combination with the
    KP step itself.

    The Executor is responsible for running single KP Steps. However in the
    future this might get extendend.


    ### Usage:

    The Executor can be used as a context manager.
    ```python
    with BaseStepExecutor() as exc:
        exc(MyStep, [], Path("output"))
        # exc(...) == exc.execute_step(...)
    ```


    """

    __dont_encapsulate: bool

    @classmethod
    def is_allow_extra_settings(cls: "BaseStepExecutor") -> bool:
        """Checks os for ALLOW_EXTRA_SETTINGS.

        Args:
            cls (BaseStepExecutor)

        Returns:
            bool: ALLOW_EXTRA_SETTINGS

        """
        return os.environ.get("ALLOW_EXTRA_SETTINGS", "False").upper() == "TRUE"

    def __init__(self, dont_encapsulate: bool = False) -> None:
        self.__dont_encapsulate = dont_encapsulate

    def store(
        self,
        step: TypedStep,
        hist: History,
        obj: Union[PanderaDataFrameModel, PydanticModel, list[PydanticModel]],
        path: PathToFolderWithBaseModels,
    ) -> Path:
        """Store step result.

        Args:
            step (TypedStep): Step Object
            obj (Union[PanderaDataFrameModel, PydanticModel]): Result of step
            path (PathToBaseModel): Where to save

        """
        if not path.is_dir():
            path.mkdir()
        obj = _try_sort(obj)
        output_model_class: type[datacontract.DataModel] = step.output_model_class
        return output_model_class.save_to_path(path / f"{hist}", obj)

    def load(self, step: TypedStep, path: PathToFolderWithBaseModels):
        """Load step input.

        Args:
            step (TypedStep): step obj for which to load input
            path (PathToBaseModel): Path to input
            model_type (Union[Type, None]): Type

        Returns:
            Union[Type[None], Union[PanderaDataFrameModel, PydanticModel]]: _description_

        """
        input_model_class: type[datacontract.DataModel] = step.input_model_class
        for p in path.glob("*"):
            start = time.time()
            data = input_model_class.load_from_path(p, step.input_model_type)
            yield (
                (data, History(".".join(p.name.split(".")[:-1]), step)),
                time.time() - start,
            )

    def _load_data(
        self,
        step: TypedStep,
        inputs: set[
            Union[
                PydanticModel,
                patyp.DataFrame[PanderaDataFrameModel],
                PathToFolderWithBaseModels,
            ]
        ],
        output_path: Path,
    ) -> Generator[
        tuple[
            tuple[
                Union[
                    PydanticModel,
                    patyp.DataFrame[PanderaDataFrameModel],
                    list[PydanticModel],
                    None,
                ],
                History,
            ],
            float,
        ],
        Any,
        NoneType,
    ]:
        """Load (if necessary from disk) inputdata.

        Args:
            step (TypedStep): step to load data for
            inputs (set[Union[PydanticModel, DataFrame[PanderaDataFrameModel], PathToFolderWithBaseModels]]): inputs

        Yields:
            Loaded input: With history already extended

        """
        no_inputs_supplied = inputs in [[], set()]
        if no_inputs_supplied and isinstance(step, SelfConsumingLeafStep):
            is_run_already = False
            for (inpt, hist), took in self.load(step, output_path):
                is_run_already = True
                yield (inpt, hist), took
            if is_run_already:
                return

        if no_inputs_supplied:
            # Only yield once
            yield (None, History(step)), 0
        for inpt in inputs:
            if isinstance(inpt, (datacontract.DataModel, PydanticModel, patyp.DataFrame)):
                yield (inpt, History("[Memory]", step)), 0
            elif isinstance(inpt, (Path, PathToFolderWithBaseModels)):
                for (inpt, hist), took in self.load(step, inpt):
                    yield (inpt, hist), took
            else:
                raise NotImplementedError(f"Cannot load/convert {inpt} as input for a step")

    def _execute_step(
        self,
        step_cls: type[TypedStep],
        inputs: set[PathToFolderWithBaseModels],
        output_path: Optional[PathToFolderWithBaseModels],
    ):
        """Exceute specified step."""
        step = step_cls()
        # pylint: disable=protected-access
        if output_path:
            output_path = step._internal_output_class(output_path)
        run: Callable[[list], Any] = pydantic.validate_call(step.run, validate_return=True)
        was_called_once = False
        try:
            for (inpt, history), load_time in self._load_data(step, inputs, output_path):
                was_called_once = True
                log.info(
                    f"Start: {step_cls.__name__}.run({history[:-1]}) -> {output_path}",
                )
                history: History
                run_start = time.time()
                ctx = copy_context()
                token = step_history.set(history)
                ctx.run(step_history.set, history)
                try:
                    res = ctx.run(run, inpt)
                finally:
                    step_history.reset(token)
                run_time = time.time() - run_start
                store_time = 0
                if output_path:
                    self.store(step, history, res, output_path)
                    store_time = time.time() - run_start
                report = StepReport(
                    time_to_load=load_time,
                    time_to_execute=run_time,
                    time_to_save=store_time,
                    step_name=step_cls.__name__,
                    inputs=try_get_length(inpt),
                    results=getattr(res, "__len__", lambda: 1)(),
                    history=history.get(),
                )
                log.info(
                    f"Finished: {step_cls.__name__}.run({history[:-1]}) -> {output_path}",
                    extra=report.model_dump(),
                )
                yield res, report
            log.info(f"{step_cls.__name__}.finalize()")
            if not was_called_once:
                log.error(
                    "Run was never executed",
                    extra={
                        "step": step.__class__.__name__,
                        "inpts": inputs,
                    },
                )
            step.finalize()
        except ValidationError as err:
            raise ContractFailedException(f"{inputs} does not conform the data contract of {step.input_model_class.__name__}") from err

    def execute_step(
        self,
        step_cls: type[TypedStep],
        inputs: Optional[set[PathToFolderWithBaseModels]],
        output_dir: Optional[PathToFolderWithBaseModels],
    ) -> list[tuple[Any, StepReport]]:
        """Execute a step.

        Args:
            step_cls (Type[TypedStep]): class to execute
            inputs (Optional[set[PathToFolderWithBaseModels]]): Either objects or a folder path
            output_dir (Optional[PathToFolderWithBaseModels]): If folder is supllied, save to folder

        Returns:
            list[tuple[Any, StepReport]]: Result of step and repots in list

        """
        if not inputs:
            inputs = set()
        try:
            correlation_id.set(step_cls.__name__)
            log.info(f"{self.__class__.__name__} - start: {step_cls.__name__}")
            if self.__dont_encapsulate:
                return list(self._execute_step(step_cls, inputs, output_dir))
            with step_env_encapsulation(step_cls):
                return list(self._execute_step(step_cls, inputs, output_dir))
        except LoggedCustomException:
            raise
        except Exception as e:
            raise StepFailed(f"{self.__class__.__name__} error during {step_cls.__name__}: {str(e)}") from e
        finally:
            log.info(f"{self.__class__.__name__} - done: {step_cls.__name__}")
            correlation_id.set(None)

    __call__ = execute_step

    def __enter__(self) -> Self:
        return self

    def __exit__(self, *exc_details):
        pass
