EricRollei-AAA_Metadata_System/eric_metadata/hooks/runtime_capture.py

"""Runtime hook infrastructure for optional ComfyUI capture."""

from __future__ import annotations

import copy
import datetime
import os
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Tuple, Type

HookCallable = Callable[..., Any]


@dataclass
class HookSession:
    """Lightweight container for per-prompt runtime capture."""

    prompt_id: str
    prompt_payload: Any
    start_ts: float
    events: List[Dict[str, Any]] = field(default_factory=list)
    resolved_inputs: Dict[str, Any] = field(default_factory=dict)
    errors: List[str] = field(default_factory=list)
    # Performance tracking
    hook_call_count: int = 0
    hook_total_time: float = 0.0
    hook_max_time: float = 0.0


@dataclass
class HookState:
    """Mutable runtime hook status flags."""

    enabled: bool = False
    conflict_detected: bool = False
    auto_disabled: bool = False
    disable_reason: Optional[str] = None
    executor_class: Optional[Type[Any]] = None


_LOCK = threading.RLock()
_STATE = HookState()
_SESSIONS: Dict[str, HookSession] = {}
_EXECUTOR_PROMPTS: Dict[int, str] = {}
_BASELINES: Dict[int, Dict[str, HookCallable]] = {}
_ORIGINALS: Dict[str, HookCallable] = {}
_ERROR_SUPPRESSION: Dict[str, bool] = {}
_ENV_FLAG = "AAA_METADATA_ENABLE_HOOKS"


def _log(message: str) -> None:
    print(f"[AAA Metadata Hooks] {message}")


def _log_performance_metrics(session: HookSession) -> None:
    """Log performance metrics for a completed session."""
    try:
        from ..utils.config import get_runtime_log_performance, get_debug_enabled
        
        if not (get_runtime_log_performance() and get_debug_enabled()):
            return
        
        total_time_ms = session.hook_total_time * 1000
        max_time_ms = session.hook_max_time * 1000
        avg_time_ms = (total_time_ms / session.hook_call_count) if session.hook_call_count > 0 else 0
        
        _log(
            f"Performance: {session.prompt_id[:8]}... - "
            f"{session.hook_call_count} calls, "
            f"total={total_time_ms:.2f}ms, "
            f"avg={avg_time_ms:.3f}ms, "
            f"max={max_time_ms:.3f}ms"
        )
    except Exception:
        # Don't fail session consumption due to logging issues
        pass


def _register_baseline(executor_class: Type[Any]) -> Dict[str, HookCallable]:
    class_id = id(executor_class)
    baseline = _BASELINES.get(class_id)
    if baseline:
        return baseline

    pre_execute = getattr(executor_class, "pre_execute", None)
    pre_get_input = getattr(executor_class, "pre_get_input_data", None)

    baseline = {
        "pre_execute": pre_execute,
        "pre_get_input_data": pre_get_input,
    }
    _BASELINES[class_id] = baseline
    return baseline


def _is_conflicting(current: Optional[HookCallable], baseline: Optional[HookCallable]) -> bool:
    if current is None or baseline is None:
        return False
    if current is baseline:
        return False
    if getattr(current, "__aaa_metadata_wrapper__", False):
        return False
    return True


def _extract_prompt_id(prompt_payload: Any) -> str:
    if isinstance(prompt_payload, dict):
        for key in ("prompt_id", "id", "uuid"):
            value = prompt_payload.get(key)
            if value:
                return str(value)
    return f"prompt-{int(time.time() * 1000)}"


def _snapshot_prompt(payload: Any) -> Any:
    try:
        return copy.deepcopy(payload)
    except Exception:
        return payload


def _record_event(session: HookSession, event_type: str, data: Optional[Dict[str, Any]] = None) -> None:
    event = {
        "type": event_type,
        "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(),
    }
    if data:
        event.update(data)
    session.events.append(event)


def _capture_pre_execute(executor: Any, *args: Any, **kwargs: Any) -> None:
    hook_start = time.perf_counter()
    
    prompt_payload = args[0] if args else kwargs.get("prompt")
    prompt_id = _extract_prompt_id(prompt_payload)

    session = HookSession(
        prompt_id=prompt_id,
        prompt_payload=_snapshot_prompt(prompt_payload),
        start_ts=time.time(),
    )
    _record_event(session, "pre_execute", {"executor_id": id(executor)})

    with _LOCK:
        _SESSIONS[prompt_id] = session
        _EXECUTOR_PROMPTS[id(executor)] = prompt_id
    
    # Track performance
    hook_elapsed = time.perf_counter() - hook_start
    session.hook_call_count += 1
    session.hook_total_time += hook_elapsed
    session.hook_max_time = max(session.hook_max_time, hook_elapsed)


def _capture_pre_get_input_data(executor: Any, *args: Any, **kwargs: Any) -> None:
    hook_start = time.perf_counter()
    
    with _LOCK:
        prompt_id = _EXECUTOR_PROMPTS.get(id(executor))
        session = _SESSIONS.get(prompt_id) if prompt_id else None

    if not session:
        return
    
    # Track performance at end of function
    def _track_performance():
        hook_elapsed = time.perf_counter() - hook_start
        session.hook_call_count += 1
        session.hook_total_time += hook_elapsed
        session.hook_max_time = max(session.hook_max_time, hook_elapsed)

    node_id: Optional[str] = None
    if args:
        for value in args:
            if isinstance(value, dict) and "id" in value:
                node_id = str(value.get("id"))
                break
            if isinstance(value, str):
                node_id = value
                break
    if node_id is None:
        candidate = kwargs.get("node") or kwargs.get("node_id") or kwargs.get("key")
        if isinstance(candidate, (str, int)):
            node_id = str(candidate)

    payload_extract = None
    if args:
        for value in reversed(args):
            if isinstance(value, dict):
                payload_extract = value
                break
    if payload_extract is None:
        payload_extract = kwargs.get("inputs")

    safe_payload = None
    if isinstance(payload_extract, dict):
        safe_payload = {}
        for key, value in payload_extract.items():
            if isinstance(value, (str, int, float, bool)):
                safe_payload[key] = value

    _record_event(
        session,
        "pre_get_input_data",
        {
            "node_id": node_id or "unknown",
            "captured_fields": list(safe_payload.keys()) if safe_payload else [],
        },
    )

    if safe_payload:
        with _LOCK:
            session.resolved_inputs.setdefault(node_id or "unknown", {}).update(safe_payload)
    
    # Track performance
    _track_performance()


def _wrap_callback(name: str, original: HookCallable, capture_fn: Callable[..., None]) -> HookCallable:
    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        try:
            capture_fn(self, *args, **kwargs)
        except Exception as exc:  # pragma: no cover - safety path
            key = f"{name}_error"
            if not _ERROR_SUPPRESSION.get(key):
                _ERROR_SUPPRESSION[key] = True
                _log(f"Error inside {name} hook: {exc}. Auto-disabling hooks.")
            with _LOCK:
                _STATE.auto_disabled = True
            disable_hooks(auto=True)
        return original(self, *args, **kwargs)

    wrapper.__aaa_metadata_wrapper__ = True
    wrapper.__aaa_metadata_original__ = original
    return wrapper


def enable_hooks(*, executor_class: Optional[Type[Any]] = None) -> bool:
    with _LOCK:
        if _STATE.enabled:
            return True

        target_cls = executor_class or _import_prompt_executor()
        if target_cls is None:
            _STATE.disable_reason = "PromptExecutor not available"
            return False

        baseline = _register_baseline(target_cls)
        current_pre_execute = getattr(target_cls, "pre_execute", None)
        current_pre_get = getattr(target_cls, "pre_get_input_data", None)

        if _is_conflicting(current_pre_execute, baseline["pre_execute"]):
            _STATE.conflict_detected = True
            _STATE.disable_reason = "pre_execute already wrapped"
            return False
        if _is_conflicting(current_pre_get, baseline["pre_get_input_data"]):
            _STATE.conflict_detected = True
            _STATE.disable_reason = "pre_get_input_data already wrapped"
            return False

        if baseline["pre_execute"] is None:
            _STATE.disable_reason = "pre_execute missing"
            return False

        wrapped_pre_execute = _wrap_callback("pre_execute", baseline["pre_execute"], _capture_pre_execute)
        setattr(target_cls, "pre_execute", wrapped_pre_execute)
        _ORIGINALS["pre_execute"] = baseline["pre_execute"]

        if baseline["pre_get_input_data"] is not None:
            wrapped_pre_get = _wrap_callback("pre_get_input_data", baseline["pre_get_input_data"], _capture_pre_get_input_data)
            setattr(target_cls, "pre_get_input_data", wrapped_pre_get)
            _ORIGINALS["pre_get_input_data"] = baseline["pre_get_input_data"]

        _STATE.enabled = True
        _STATE.executor_class = target_cls
        _STATE.auto_disabled = False
        _STATE.disable_reason = None
        _STATE.conflict_detected = False
        _log("Runtime hooks enabled")
        return True


def disable_hooks(*, auto: bool = False) -> bool:
    with _LOCK:
        target_cls = _STATE.executor_class
        if not target_cls:
            _STATE.enabled = False
            if auto:
                _STATE.auto_disabled = True
            _SESSIONS.clear()
            _EXECUTOR_PROMPTS.clear()
            return False

        baseline = _register_baseline(target_cls)
        if "pre_execute" in _ORIGINALS and baseline["pre_execute"] is not None:
            setattr(target_cls, "pre_execute", baseline["pre_execute"])
        if "pre_get_input_data" in _ORIGINALS and baseline["pre_get_input_data"] is not None:
            setattr(target_cls, "pre_get_input_data", baseline["pre_get_input_data"])

        _STATE.enabled = False
        _STATE.executor_class = None
        _STATE.auto_disabled = _STATE.auto_disabled or auto
        _STATE.disable_reason = None
        _STATE.conflict_detected = False
        _ORIGINALS.clear()
        _ERROR_SUPPRESSION.clear()
        _SESSIONS.clear()
        _EXECUTOR_PROMPTS.clear()
        _log("Runtime hooks disabled")
        return True


def hooks_enabled() -> bool:
    with _LOCK:
        return _STATE.enabled and not _STATE.auto_disabled


def get_session(prompt_id: str) -> Optional[HookSession]:
    with _LOCK:
        session = _SESSIONS.get(prompt_id)
        return copy.deepcopy(session) if session else None


def consume_session(prompt_id: str) -> Optional[HookSession]:
    with _LOCK:
        session = _SESSIONS.pop(prompt_id, None)
        if session:
            keys_to_remove = [key for key, value in _EXECUTOR_PROMPTS.items() if value == prompt_id]
            for key in keys_to_remove:
                _EXECUTOR_PROMPTS.pop(key, None)
            
            # Log performance metrics if enabled
            _log_performance_metrics(session)
        
        return session


def active_prompt_ids() -> List[str]:
    with _LOCK:
        return list(_SESSIONS.keys())


def auto_enable_from_env(*, executor_class: Optional[Type[Any]] = None) -> None:
    """
    Auto-enable hooks based on configuration.
    
    Priority order:
    1. Environment variable AAA_METADATA_ENABLE_HOOKS
    2. config.json runtime_hooks.enabled setting
    3. Default: disabled
    """
    try:
        from ..utils.config import get_runtime_hooks_enabled
        if get_runtime_hooks_enabled():
            enable_hooks(executor_class=executor_class)
    except Exception as exc:
        # Fallback to environment variable only
        flag = os.environ.get(_ENV_FLAG, "").strip().lower()
        if flag in {"1", "true", "yes", "on"}:
            enable_hooks(executor_class=executor_class)


def _import_prompt_executor() -> Optional[Type[Any]]:
    try:
        from comfy.executor import PromptExecutor  # type: ignore
    except Exception:  # pragma: no cover - not available during tests
        return None
    return PromptExecutor


def reset_state_for_tests() -> None:
    with _LOCK:
        if _STATE.executor_class:
            disable_hooks()
        _SESSIONS.clear()
        _EXECUTOR_PROMPTS.clear()
        _ERROR_SUPPRESSION.clear()
        _STATE.enabled = False
        _STATE.conflict_detected = False
        _STATE.auto_disabled = False
        _STATE.disable_reason = None
        _STATE.executor_class = None


__all__ = [
    "HookSession",
    "enable_hooks",
    "disable_hooks",
    "hooks_enabled",
    "get_session",
    "consume_session",
    "active_prompt_ids",
    "auto_enable_from_env",
    "reset_state_for_tests",
]