# 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, Type, Union

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.typed_step import (
    TypedStep,
)
from wurzel.utils import WZ, create_model, try_get_length

log = getLogger(__name__)


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]


@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],
        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()
        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,
            ]
        ],
    ) -> 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:
            # Only yield once
            no_inputs_supplied = False
            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):
                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
