"""Post-task reflection: run Critic and write lessons/heuristics to reflective memory."""

from typing import Any, Callable, Protocol

from fusionagi.schemas.messages import AgentMessage, AgentMessageEnvelope
from fusionagi._logger import logger


class CriticLike(Protocol):
    """Protocol for critic agent: must have identity and handle_message."""

    identity: str

    def handle_message(self, envelope: AgentMessageEnvelope) -> AgentMessageEnvelope | None:
        ...


class ReflectiveMemoryLike(Protocol):
    """Protocol for reflective memory: must have add_lesson and set_heuristic."""

    def add_lesson(self, lesson: dict[str, Any]) -> None:
        ...

    def set_heuristic(self, key: str, value: Any) -> None:
        ...


ReflectionCallback = Callable[[str, dict[str, Any]], None]
"""Callback (event_type, payload) -> None. Emits 'reflection_done' with task_id and evaluation."""


def run_reflection(
    critic_agent: CriticLike,
    task_id: str,
    outcome: str,
    trace: list[dict[str, Any]],
    plan: dict[str, Any] | None,
    reflective_memory: ReflectiveMemoryLike | None,
    orchestrator_callback: ReflectionCallback | None = None,
) -> dict[str, Any] | None:
    """
    Trigger reflection: send evaluate_request to Critic, then write evaluation
    to reflective memory (lessons, heuristics). Optionally notify orchestrator
    via orchestrator_callback(event_type, payload); e.g. "reflection_done" with task_id and evaluation.
    Returns evaluation dict or None.
    """
    envelope = AgentMessageEnvelope(
        message=AgentMessage(
            sender="orchestrator",
            recipient=critic_agent.identity,
            intent="evaluate_request",
            payload={
                "outcome": outcome,
                "trace": trace,
                "plan": plan,
            },
        ),
        task_id=task_id,
    )
    response = critic_agent.handle_message(envelope)
    if not response or response.message.intent != "evaluation_ready":
        return None
    evaluation = response.message.payload.get("evaluation", {})
    if reflective_memory:
        reflective_memory.add_lesson({
            "task_id": task_id,
            "outcome": outcome,
            "evaluation": evaluation,
        })
        suggestions = evaluation.get("suggestions", [])
        for i, s in enumerate(suggestions[:5]):
            reflective_memory.set_heuristic(f"suggestion_{task_id}_{i}", s)
    if orchestrator_callback:
        try:
            orchestrator_callback("reflection_done", {"task_id": task_id, "evaluation": evaluation})
        except Exception:
            logger.exception(
                "Orchestrator callback failed (reflection_done); callback is best-effort",
                extra={"task_id": task_id},
            )
    return evaluation