workflow_engine.py

"""
workflow_engine.py — Motor de workflow YAWL-inspired para a Trindade Pipeline.
══════════════════════════════════════════════════════════════════════════════

Conceitos YAWL implementados
─────────────────────────────
  Task        : unidade atômica de execução (um turno de LLM ou operação)
  Condition   : expressão Python avaliada contra o ExecutionContext
  XOR-split   : exatamente um arco de saída dispara (primeira condição verdadeira)
  AND-split   : todos os arcos disparam (execução paralela — futuro)
  OR-split    : um ou mais arcos disparam (subconjunto verdadeiro)
  Loop        : task repete enquanto condição for verdadeira
  SubNet      : task que encapsula um sub-workflow inteiro

Fluxo de controle
──────────────────
  WorkflowEngine.start()        → retorna a primeira TaskDef
  WorkflowEngine.advance(ctx)   → avalia transições do nó atual, retorna próximo
  WorkflowEngine.is_terminal()  → True se chegou ao nó END

O ExecutionContext é um dict simples que o pipeline preenche após cada task.
O motor lê esse dict para avaliar as condições — não conhece LLMs, HTTP, nem JSON.

Segurança do eval
──────────────────
  Condições são avaliadas com eval() em namespace restrito:
    • Apenas builtins seguros (len, int, float, bool, str, min, max, abs, round)
    • Mais todas as chaves do ExecutionContext
    • Nenhum acesso a __import__, open, os, sys, etc.
══════════════════════════════════════════════════════════════════════════════
"""
from __future__ import annotations

import copy
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional

import yaml


# ══════════════════════════════════════════════════════════════════
#  TIPOS DE SPLIT
# ══════════════════════════════════════════════════════════════════
class SplitType(str, Enum):
    XOR = "XOR"   # exatamente um arco (default)
    AND = "AND"   # todos os arcos (paralelo)
    OR  = "OR"    # um ou mais arcos


# ══════════════════════════════════════════════════════════════════
#  TRANSIÇÃO  (arco de saída de um nó)
# ══════════════════════════════════════════════════════════════════
@dataclass
class Transition:
    target:    str
    condition: Optional[str] = None   # None = default/unconditional
    label:     str           = ""

    def evaluate(self, ctx: Dict[str, Any]) -> bool:
        """Avalia a condição contra o contexto. None → sempre True."""
        if self.condition is None:
            return True
        return _safe_eval(self.condition, ctx)


# ══════════════════════════════════════════════════════════════════
#  TASK DEFINITION
# ══════════════════════════════════════════════════════════════════
@dataclass
class TaskDef:
    """
    Definição declarativa de uma task no workflow.

    Campos relevantes para o executor (pipeline_v33.py):
      id          : identificador único do nó
      task_type   : "reasoning" | "final" | "audit_reasoning" |
                    "audit_final" | "normalize" | "validate" |
                    "loop_recovery" | "end"
      phase       : chave em PAYLOAD_CONFIG (ex. "STEP1", "AUDIT_REASONING")
      provider    : slug do provider (ex. "groq", "openrouter")
      model       : model id (sobrescreve o default do PAYLOAD_CONFIG)
      stop_tokens : se True, usa STOP_TOKENS_REASONING
      is_final    : agente final (escreve JSON completo)
      loop        : TaskLoop se esta task pode ser repetida

    O executor não interpreta nenhum outro campo — passa como `task.params`.
    """
    id:          str
    task_type:   str
    phase:       str
    provider:    str                    = "groq"
    model:       Optional[str]          = None
    stop_tokens: bool                   = False
    is_final:    bool                   = False
    loop:        Optional["TaskLoop"]   = None
    transitions: List[Transition]       = field(default_factory=list)
    params:      Dict[str, Any]         = field(default_factory=dict)
    description:  str                    = ""
    prompt:       Optional[str]          = None
    instructions: Optional[str]          = None
    output_keys:  List[str]              = field(default_factory=list)

    @property
    def is_terminal(self) -> bool:
        return self.task_type == "end"

    def next_tasks(self, ctx: Dict[str, Any], split: SplitType = SplitType.XOR) -> List[str]:
        """Retorna ids dos próximos nós dado o contexto atual."""
        results: List[str] = []
        for t in self.transitions:
            if t.evaluate(ctx):
                results.append(t.target)
                if split == SplitType.XOR:
                    break   # XOR: para no primeiro verdadeiro
        return results


# ══════════════════════════════════════════════════════════════════
#  LOOP DEFINITION
# ══════════════════════════════════════════════════════════════════
@dataclass
class TaskLoop:
    """
    Define comportamento de loop para uma task.

    while_condition : expressão avaliada ANTES de cada iteração
    until_condition : expressão avaliada APÓS cada iteração (do-while)
    max_iterations  : teto de segurança (evita loop infinito)
    """
    while_condition: Optional[str] = None
    until_condition: Optional[str] = None
    max_iterations:  int           = 3

    def should_enter(self, ctx: Dict[str, Any]) -> bool:
        if self.while_condition:
            return _safe_eval(self.while_condition, ctx)
        return True

    def should_continue(self, ctx: Dict[str, Any], iteration: int) -> bool:
        if iteration >= self.max_iterations:
            return False
        if self.until_condition:
            return not _safe_eval(self.until_condition, ctx)
        if self.while_condition:
            return _safe_eval(self.while_condition, ctx)
        return False


# ══════════════════════════════════════════════════════════════════
#  EXECUTION CONTEXT
# ══════════════════════════════════════════════════════════════════
class ExecutionContext(dict):
    """
    Dict especializado que representa o estado de execução do workflow.

    O motor lê este dict para avaliar condições.
    O executor escreve nele após cada task.

    Chaves convencionais (o motor conhece)
    ───────────────────────────────────────
      reasoning_chars   : int   — chars de reasoning da última fase intermediária
      json_valid        : bool  — True se o JSON final passou parse
      audit_passed      : bool  — True se o último ciclo de audit retornou AUDIT_PASS
      loop_count        : int   — tentativas de recovery de loop
      atom_count        : int   — total de átomos extraídos
      correction_count  : int   — ciclos de correção do audit
      schema_errors     : int   — erros de schema após validate_and_fix
      phase_failed      : bool  — True se a fase atual falhou (content=None)
      current_phase     : str   — id do nó em execução
    """
    def update_phase(self, phase_id: str, **kwargs: Any) -> None:
        self["current_phase"] = phase_id
        self.update(kwargs)

    def increment(self, key: str, by: int = 1) -> int:
        self[key] = self.get(key, 0) + by
        return self[key]

    def snapshot(self) -> Dict[str, Any]:
        return dict(self)


# ══════════════════════════════════════════════════════════════════
#  WORKFLOW ENGINE
# ══════════════════════════════════════════════════════════════════
class WorkflowEngine:
    """
    Máquina de estados finitos YAWL-inspired.

    Uso típico no pipeline_v33.py
    ───────────────────────────────
        engine = WorkflowEngine.from_yaml("trindade_workflow.yaml")
        ctx    = ExecutionContext(defaults)
        task   = engine.start()

        while not task.is_terminal:
            result = executor.run(task, memory)
            ctx.update(result)
            task = engine.advance(ctx)
    """

    def __init__(self, tasks: Dict[str, TaskDef], start_id: str) -> None:
        self._tasks       = tasks
        self._start_id    = start_id
        self._current     = start_id
        self._wrapper_key = "output"   # sobrescrito por from_dict
        self._id_key      = "id"       # sobrescrito por from_dict

    @property
    def wrapper_key(self) -> str:
        """Chave de wrapper do resultado final (ex: 'manifestacao_juridica')."""
        return self._wrapper_key

    @property
    def id_key(self) -> str:
        """Chave de id dentro do wrapper (ex: 'id_manifestacao')."""
        return self._id_key

    # ── Navegação ─────────────────────────────────────────────────
    def start(self) -> TaskDef:
        self._current = self._start_id
        return self._tasks[self._current]

    def current(self) -> TaskDef:
        return self._tasks[self._current]

    def advance(self, ctx: ExecutionContext) -> TaskDef:
        """
        Avalia as transições do nó atual e move para o próximo.
        Retorna a TaskDef do próximo nó (pode ser END).
        """
        current_task = self._tasks[self._current]
        split_type   = SplitType(current_task.params.get("split", SplitType.XOR.value))
        next_ids     = current_task.next_tasks(ctx, split=split_type)

        if not next_ids:
            # Sem transição → END implícito
            self._current = "__END__"
            return TaskDef(id="__END__", task_type="end", phase="END")

        # Para XOR e OR, executa o primeiro target (AND é futuro)
        next_id = next_ids[0]
        if next_id not in self._tasks:
            raise KeyError(f"Nó '{next_id}' referenciado mas não definido no workflow")

        self._current = next_id
        ctx["current_phase"] = next_id
        return self._tasks[next_id]

    def peek_next(self, ctx: ExecutionContext) -> Optional[str]:
        """Retorna o id do próximo nó sem avançar o estado."""
        task = self._tasks[self._current]
        ids  = task.next_tasks(ctx)
        return ids[0] if ids else None

    def reset(self) -> None:
        self._current = self._start_id

    # ── Factory ───────────────────────────────────────────────────
    @classmethod
    def from_yaml(cls, path: str) -> "WorkflowEngine":
        """Carrega um workflow de um arquivo YAML."""
        with open(path, encoding="utf-8") as f:
            spec = yaml.safe_load(f)
        return cls.from_dict(spec)

    @classmethod
    def from_dict(cls, spec: Dict[str, Any]) -> "WorkflowEngine":
        """Constrói o engine a partir de um dict (já parseado)."""
        wf        = spec["workflow"]
        start_id  = wf["start"]
        tasks: Dict[str, TaskDef] = {}

        # ── Defaults do workflow (base para todos os tasks) ────────
        wf_defaults = wf.get("defaults", {})
        # Campos top-level da TaskDef — não vão para params
        _top_level = {"id", "type", "phase", "provider", "model", "stop_tokens",
                      "is_final", "loop", "transitions", "description",
                      "prompt", "instructions", "output_keys", "params"}

        # ── Configuração de saída agnóstica ────────────────────────
        out_cfg     = wf.get("output", {})
        wrapper_key = out_cfg.get("wrapper_key", "output")
        id_key      = out_cfg.get("id_key", "id")

        for raw in wf["tasks"]:
            task_id = raw["id"]

            # ── Loop ──────────────────────────────────────────────
            loop = None
            if "loop" in raw:
                lraw = raw["loop"]
                loop = TaskLoop(
                    while_condition = lraw.get("while"),
                    until_condition = lraw.get("until"),
                    max_iterations  = lraw.get("max_iterations", 3),
                )

            # ── Transitions ───────────────────────────────────────
            transitions = []
            for t in raw.get("transitions", []):
                transitions.append(Transition(
                    target    = t["target"],
                    condition = t.get("condition"),
                    label     = t.get("label", ""),
                ))

            # ── Params: defaults do YAML ← sobrescritos por task ──
            params = {k: v for k, v in wf_defaults.items() if k not in _top_level}
            params.update({k: v for k, v in raw.items() if k not in _top_level})

            tasks[task_id] = TaskDef(
                id           = task_id,
                task_type    = raw.get("type", "task"),
                phase        = raw.get("phase", task_id),
                provider     = raw.get("provider", wf_defaults.get("provider", "groq")),
                model        = raw.get("model", wf_defaults.get("model")),
                stop_tokens  = raw.get("stop_tokens", False),
                is_final     = raw.get("is_final", False),
                loop         = loop,
                transitions  = transitions,
                params       = params,
                description  = raw.get("description", ""),
                prompt       = raw.get("prompt"),
                instructions = raw.get("instructions"),
                output_keys  = raw.get("output_keys", []),
            )

        engine = cls(tasks=tasks, start_id=start_id)
        engine._wrapper_key = wrapper_key
        engine._id_key      = id_key
        return engine

    # ── Inspeção ──────────────────────────────────────────────────
    def task_ids(self) -> List[str]:
        return list(self._tasks.keys())

    def describe(self) -> str:
        """Retorna representação textual do grafo para debug."""
        lines = [f"WorkflowEngine  start={self._start_id}  nodes={len(self._tasks)}"]
        for tid, t in self._tasks.items():
            arrow = " → ".join(
                f"{tr.target}[{tr.condition or 'default'}]"
                for tr in t.transitions
            ) or "(terminal)"
            lines.append(f"  {tid:30s} ({t.task_type:20s})  {arrow}")
        return "\n".join(lines)


# ══════════════════════════════════════════════════════════════════
#  SAFE EVAL  —  avalia condições em namespace restrito
# ══════════════════════════════════════════════════════════════════
_SAFE_BUILTINS = {
    "len":   len,
    "int":   int,
    "float": float,
    "bool":  bool,
    "str":   str,
    "min":   min,
    "max":   max,
    "abs":   abs,
    "round": round,
    "True":  True,
    "False": False,
    "None":  None,
    "all":   all,
    "any":   any,
}


def _safe_eval(expr: str, ctx: Dict[str, Any]) -> bool:
    """
    Avalia `expr` (string Python) contra `ctx`.

    Namespace: builtins seguros + ctx.
    Qualquer exceção → False (condição não satisfeita).

    Exemplos válidos de expressões:
      "reasoning_chars > 100"
      "audit_passed == True"
      "loop_count < 2 and not phase_failed"
      "atom_count >= 3"
      "schema_errors == 0"
    """
    namespace = {**_SAFE_BUILTINS, **ctx}
    try:
        result = eval(expr, {"__builtins__": {}}, namespace)  # noqa: S307
        return bool(result)
    except Exception:
        return False