"""Reflexion — episodic failure memory injected between workflow retries.

Between a failed (below-threshold) critique and the next revision dispatch,
the supervisor fires a small LLM call ("why did this fail, what will I do
differently next time"), stores the verbal self-correction in a bounded
deque keyed by ``(agent_key, step_id)``, and prepends the accumulated
reflections into the next revision's task.

Inspired by Shinn et al. 2023 (arxiv:2303.11366). This module is a pure
re-implementation of the pattern against our WorkflowSupervisor primitives
— no third-party code is imported.

Design:
  - **Bounded buffer**: max 3 reflections per ``(agent_key, step_id)``
    tuple. Oldest dropped when full. Cleared when the enclosing
    ``WorkflowExecution`` completes so reflections don't leak across
    unrelated workflow runs.
  - **Fail-open**: when the reflection LLM call fails (timeout, provider
    down, parse error), the revision loop continues without reflection.
    Reflexion is a quality boost, never a blocker.
  - **Skill-driven**: the prompt, token budget, timeout, max-per-step,
    and enabled flag all live in ``_meta/reflexion-protocol/config/reflexion.yaml``.
    Legacy defaults in :data:`_FALLBACK` preserve behavior when the skill
    isn't loaded or the YAML is malformed.
"""

from __future__ import annotations

import asyncio
import logging
from collections import deque
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from agntspace.llm.base import LLMProvider
    from agntspace.skills.loader import SkillLoader

log = logging.getLogger(__name__)


# Load-bearing safety net when the skill YAML is missing or malformed.
# Keep in sync with the shipped `reflexion.yaml`.
_FALLBACK: dict[str, Any] = {  # arch:deterministic — safety net for missing skill YAML
    "enabled": True,
    "max_reflections_per_step": 3,
    "model_tier": "fast",
    "max_tokens": 220,
    "timeout_seconds": 30,
    "prompt": (
        "You are a reflection agent helping another agent learn from a failed attempt.\n\n"
        "Task: {task}\n\n"
        "Prior output (scored {score:.2f} by a reviewer):\n{output}\n\n"
        "Reviewer feedback:\n{critique}\n\n"
        "In 1-2 short sentences, state EXACTLY what you would do differently next time. "
        "Focus on a concrete, actionable change — not generic advice. Start with 'Next time: '."
    ),
}


@dataclass
class Reflection:
    """One self-correction entry. Short, dated, scoped."""

    text: str
    at: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
    score_before: float = 0.0


class ReflectionBuffer:
    """Bounded per-(agent, step) deque of Reflection entries.

    Designed for in-memory use within a single WorkflowExecution. The
    supervisor owns one instance per ``execute()`` call and drops it when
    the workflow finishes — no cross-run leak.
    """

    def __init__(self, max_per_step: int) -> None:
        self._max = max(1, int(max_per_step))
        self._store: dict[tuple[str, str], deque[Reflection]] = {}

    def append(self, agent_key: str, step_id: str, reflection: Reflection) -> None:
        key = (agent_key, step_id)
        buf = self._store.get(key)
        if buf is None:
            buf = deque(maxlen=self._max)
            self._store[key] = buf
        buf.append(reflection)

    def get(self, agent_key: str, step_id: str) -> list[Reflection]:
        buf = self._store.get((agent_key, step_id))
        return list(buf) if buf else []

    def clear(self) -> None:
        self._store.clear()

    def size(self) -> int:
        return sum(len(b) for b in self._store.values())


def format_reflections_block(reflections: list[Reflection]) -> str:
    """Human-readable block to prepend to the revised task prompt.

    Returns empty string when the list is empty so callers can
    unconditionally concatenate.
    """
    if not reflections:
        return ""
    lines = ["Past reflections (self-corrections from prior attempts on this step):"]
    for ref in reflections:
        body = ref.text.strip()
        if not body:
            continue
        if not body.lower().startswith("next time"):
            body = "Next time: " + body
        lines.append(f"- {body}")
    if len(lines) == 1:
        return ""
    return "\n".join(lines)


def load_reflexion_config(skill_loader: SkillLoader | None) -> dict[str, Any]:
    """Load the reflexion YAML; always returns a valid dict.

    Missing skill, malformed YAML, or missing keys → fall through to
    :data:`_FALLBACK`. Merges config over fallback so partial YAMLs
    still get the safety-net defaults for keys they didn't set.
    """
    cfg = dict(_FALLBACK)
    if skill_loader is None:
        return cfg
    try:
        loaded = skill_loader.get_skill_config("reflexion-protocol", "reflexion")
    except Exception:
        return cfg
    if not isinstance(loaded, dict):
        return cfg
    # Accept either the whole skill dict (with a 'reflexion:' block) or
    # the reflexion block directly.
    block = loaded.get("reflexion", loaded) if isinstance(loaded.get("reflexion"), dict) else loaded
    if isinstance(block, dict):
        cfg.update({k: v for k, v in block.items() if v is not None})
    return cfg


async def generate_reflection(
    *,
    llm: LLMProvider,
    config: dict[str, Any],
    task: str,
    output: str,
    critique: str,
    score: float,
    model_router: Any | None = None,
) -> Reflection | None:
    """Fire one reflection LLM call. Returns None on any failure.

    Never raises — the supervisor's revision loop must survive a dead LLM.
    Timeout is enforced via ``asyncio.wait_for``. The parsed text is
    trimmed to a single paragraph (collapse newlines) and capped at 400
    chars so a verbose model can't flood the next revision's prompt.
    """
    from agntspace.llm.base import Message as LLMMessage

    if not config.get("enabled", True):
        return None

    template = config.get("prompt") or _FALLBACK["prompt"]
    try:
        prompt_text = template.format(
            task=(task or "").strip()[:1500],
            output=(output or "").strip()[:1500],
            critique=(critique or "").strip()[:1000],
            score=float(score),
        )
    except Exception:
        log.debug("reflexion: prompt template formatting failed", exc_info=True)
        return None

    tier = str(config.get("model_tier", "fast"))
    timeout_s = float(config.get("timeout_seconds", 30))
    max_tokens = int(config.get("max_tokens", 220))

    # Resolve the tier → model when a router is wired; otherwise let the
    # provider pick its default. Matches the pattern daily_cycle uses.
    model_override = None
    if model_router is not None:
        try:
            model_override = model_router.resolve(tier=tier)
        except Exception:
            model_override = None

    try:
        call = llm.complete(
            messages=[LLMMessage(role="user", content=prompt_text)],
            system="Reply with ONE short self-correction starting with 'Next time:'. No preamble.",
            max_tokens=max_tokens,
            temperature=0.3,
            **({"model": model_override} if model_override else {}),
        )
        response = await asyncio.wait_for(call, timeout=timeout_s)
    except TimeoutError:
        log.debug("reflexion: LLM call timed out after %.1fs", timeout_s)
        return None
    except Exception as exc:
        log.debug("reflexion: LLM call failed: %s", exc, exc_info=True)
        return None

    text = getattr(response, "content", "") or ""
    # Collapse to a single paragraph — no newlines in stored reflections.
    cleaned = " ".join(text.split()).strip()
    if not cleaned:
        return None
    if len(cleaned) > 400:
        cleaned = cleaned[:400].rsplit(" ", 1)[0] + "…"
    return Reflection(text=cleaned, score_before=float(score))
