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

"""Base Executor"""
from contextvars import copy_context
from pathlib import Path
from typing import Generator, Optional, Type, Self, Union, Callable, Any
from types import NoneType
from contextlib import contextmanager
import time
import json
from logging import getLogger
import os
import pydantic
from pydantic_core import ValidationError
import pandera.typing as patyp
from asgi_correlation_id import correlation_id

from wurzel.exceptions import (
    ContractFailedException,
    EnvSettingsError,
    LoggedCustomException,
    StepFailed)
from wurzel.step.typed_step import (
    TypedStep,
)
from wurzel.step.history import (
    History,
    step_history,
)
from wurzel.datacontract import PanderaDataFrameModel, PydanticModel, datacontract
from wurzel.path import PathToFolderWithBaseModels
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
