taskrunner/playbook/models.py

from __future__ import annotations

import yaml
import time
import functools
from datetime import datetime, timezone
from pydantic import BaseModel, ValidationInfo, field_serializer, model_validator

from playbook.action_registry import ActionRegistry, ActionFn
from playbook.logging_models import PlaybookLog, StepLogModel, ActLog, Status


CtxType = dict[str, object]


class StepModel(BaseModel):
    name: str
    action: ActionFn
    context: dict[str, object]

    @field_serializer("action")
    def serialize_action_fn(self, action_fn: ActionFn, _info) -> str:
        return getattr(action_fn, "__name__", str(action_fn))

    def run(self, ctx: dict[str, object]) -> StepLogModel:
        """ Run the step. """
        substeps: list[StepLogModel] = []

        try:
            local_ctx: CtxType = ctx | self.context
            log: StepLogModel = timed_run(
                self.action, local_ctx, self.name
            )
        except Exception as e:
            log: StepLogModel = StepLogModel.fail(
                self.name, str(e)
            )

        substeps.append(log)
        if log.status == Status.BAD:
            return StepLogModel.fail(
                self.name,
                err=f"failed on step: {log.name} with error {log.error}"
            )

        return log


class ActModel(BaseModel):
    name: str
    steps: list[StepModel]

    # NOTE: this shouldnt return a steplogmodel but an actlogmodel or something like that
    def run(self, ctx) -> ActLog:
        """
        Runs the steps in this act.
        If a step fails we return with a failed status from this act.
        """
        step_logs: list[StepLogModel] = []
        previous_step_ctx: dict[str, object] = {}
        for step in self.steps:
            stepLog: StepLogModel = timed_run(step.run, ctx | previous_step_ctx)
            step_logs.append(stepLog)

            if stepLog.failed:
                return ActLog.fail(self.name, err=stepLog.error, logs=step_logs)

            previous_step_ctx = previous_step_ctx | stepLog.pipe_ctx

        return ActLog.ok(self.name, msg="success", logs=step_logs)


class PlaybookModel(BaseModel):
    playbook_name: str
    log_dir: str
    registries: list[ActionRegistry]
    global_context: dict[str, str]
    acts: list[ActModel]

    @model_validator(mode="before")
    @classmethod
    # def parse_registries(cls, v: object) -> list[ActionRegistry]:
    def load_registries(cls, data: object, info: ValidationInfo) -> object:
        if not isinstance(data, dict):
            return data

        raw_registries = data.get("registries", [])
        loaded_registries: list[ActionRegistry] = []

        for item in raw_registries:
            if isinstance(item, str):
                loaded_registries.extend(ActionRegistry.load_registries_from_file(item))

        data["registries"] = loaded_registries

        if info.context is not None:
            info.context["registries"] = loaded_registries

        return data

    @classmethod
    def from_yaml_file(cls, fp: str):
        with open(fp, "rb") as f:
            data = yaml.safe_load(f)
        return cls.model_validate(data, context={})
        # return cls(**data)

    def _run(self) -> PlaybookLog:
        """
        Internal _run method. This one times the run.
        Runs every act registered wheter one fails or not.
        Status is set to BAD on at least one failure.
        """
        act_logs: list[ActLog] = []
        failed_on: list[str] = []
        for act in self.acts:
            log: ActLog = timed_run(act.run, self.global_context)
            act_logs.append(log)

            if log.status == Status.BAD:
                failed_on.append(log.name)
        
        if len(failed_on) > 0:
            return PlaybookLog.fail(
                self.playbook_name,
                err='failed on the act(s): ' + ','.join(failed_on),
                logs=act_logs
            )
        return PlaybookLog.ok(name=self.playbook_name, msg="success", logs=act_logs)

    def run(self) -> PlaybookLog:
        return timed_run(self._run)


class ContextChecker:
    """Safely extract values with error messages."""

    def __init__(self, ctx: dict[str, object]):
        self._ctx = ctx

    @classmethod
    def requires(cls, required: list[str]):
        def decorator(func: ActionFn):
            @functools.wraps(func)
            def wrapper(ctx: dict[str, object], name: str) -> StepLogModel:
                if all(key in ctx for key in required):
                    return func(ctx, name)
                else:
                    missing_keys = [key for key in required if key not in ctx]
                    missing_keys_msg = f"missing context keys: {missing_keys}"
                    return StepLogModel.fail(name, err=missing_keys_msg)
            return wrapper
        return decorator


def human_readable_date(time: int) -> str:
    seconds = time // 1_000_000_000
    nanos = time % 1_000_000_000

    dt = datetime.fromtimestamp(seconds, tz=timezone.utc).astimezone()

    return dt.strftime(f"%Y-%m-%dT%H:%M:%S.{nanos:09d}%:z")


# def timed_run(op: Callable[..., StepLogModel], *args: object, **kwargs: object) -> StepLogModel:
def timed_run(op, *args: object, **kwargs: object):
    start_date: int = time.time_ns()
    start_time: int = time.perf_counter_ns()

    try:
        log = op(*args, **kwargs)
    finally:
        delta = time.perf_counter_ns() - start_time

    # NOTE: we could simplyfy this by having a marshalling method
    if isinstance(log, StepLogModel | PlaybookLog | ActLog):
        end_date: int = time.time_ns()
        log.timing.duration_sec = delta / 1_000_000_000.0
        log.timing.start_date = human_readable_date(start_date)
        log.timing.start_date_timestamp = start_date
        log.timing.end_date = human_readable_date(end_date)

    return log