backend/core/generation/expectation_loop.py

"""
Expectation quality loop: evaluate → targets → re-evaluate until ready or max-iter.

Pure control plane (no LLM). Callers inject evaluate/fix callables so unit tests
drive the real loop without network.
"""
from __future__ import annotations

import logging
import os
from dataclasses import asdict, dataclass, field
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)

DEFAULT_MAX_ITER = int(os.environ.get("ADVISOR_EXPECTATION_MAX_ITER", "3"))
DEFAULT_MIN_SCORE = int(os.environ.get("ADVISOR_EXPECTATION_MIN_SCORE", "70"))


@dataclass
class ExpectationLoopResult:
    ready: bool
    iterations: int
    max_iterations: int
    final_score: int
    regulation_grounded_pass: bool
    stop_reason: str  # ready | max_iter | blocked_grounding | evaluate_error
    history: List[Dict[str, Any]] = field(default_factory=list)
    remaining_blockers: List[str] = field(default_factory=list)
    final_report: Optional[Dict[str, Any]] = None
    state: Optional[Dict[str, Any]] = None

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


def is_regulation_ready(
    report: Any,
    *,
    min_score: int = DEFAULT_MIN_SCORE,
    require_regulation_grounded: bool = True,
) -> bool:
    """Stop condition: advisor report meets readiness criteria."""
    if report is None:
        return False
    if isinstance(report, dict):
        score = int(report.get("score") or 0)
        grounded = bool(report.get("regulation_grounded_pass"))
        passed = bool(report.get("passed"))
        blockers = list(report.get("blockers") or [])
        mode = str(report.get("grounding_mode") or "regulation")
    else:
        score = int(getattr(report, "score", 0) or 0)
        grounded = bool(getattr(report, "regulation_grounded_pass", False))
        passed = bool(getattr(report, "passed", False))
        blockers = list(getattr(report, "blockers", None) or [])
        mode = str(getattr(report, "grounding_mode", "regulation") or "regulation")

    if mode in ("structure_only", "blocked", "blind"):
        # Never regulation-ready in ungrounded modes
        return False
    if require_regulation_grounded and not grounded:
        return False
    if blockers:
        return False
    if score < min_score:
        return False
    return passed or grounded


def extract_blockers(report: Any) -> List[str]:
    if report is None:
        return ["brak raportu doradcy"]
    if isinstance(report, dict):
        blockers = list(report.get("blockers") or [])
        if blockers:
            return [str(b) for b in blockers]
        findings = report.get("findings") or []
        out = []
        for f in findings:
            if isinstance(f, dict) and (f.get("blocking") or f.get("severity") == "critical"):
                out.append(str(f.get("message") or f.get("code") or "finding"))
        return out
    blockers = list(getattr(report, "blockers", None) or [])
    if blockers:
        return [str(b) for b in blockers]
    out = []
    for f in getattr(report, "findings", None) or []:
        if getattr(f, "blocking", False) or getattr(f, "severity", "") == "critical":
            out.append(str(getattr(f, "message", "") or getattr(f, "code", "finding")))
    return out


def report_to_dict(report: Any) -> Dict[str, Any]:
    if report is None:
        return {}
    if isinstance(report, dict):
        return report
    if hasattr(report, "to_dict"):
        return report.to_dict()
    return {
        "passed": getattr(report, "passed", False),
        "score": getattr(report, "score", 0),
        "regulation_grounded_pass": getattr(report, "regulation_grounded_pass", False),
        "blockers": list(getattr(report, "blockers", None) or []),
        "grounding_mode": getattr(report, "grounding_mode", ""),
        "summary": getattr(report, "summary", ""),
    }


def run_expectation_loop(
    *,
    evaluate: Callable[[Dict[str, Any]], Any],
    apply_fixes: Optional[Callable[[Dict[str, Any], Any], Dict[str, Any]]] = None,
    initial_state: Optional[Dict[str, Any]] = None,
    max_iterations: int = DEFAULT_MAX_ITER,
    min_score: int = DEFAULT_MIN_SCORE,
    require_regulation_grounded: bool = True,
) -> ExpectationLoopResult:
    """
    Loop:
      evaluate(state) → if ready: stop
      else apply_fixes(state, report) → state' → re-evaluate
      until ready or max_iterations.

    structure_only / blind never stop as regulation-ready (is_regulation_ready=False).
    """
    state: Dict[str, Any] = dict(initial_state or {})
    max_iterations = max(1, int(max_iterations))
    history: List[Dict[str, Any]] = []
    last_report: Any = None

    for i in range(1, max_iterations + 1):
        try:
            last_report = evaluate(state)
        except Exception as e:
            logger.warning("[ExpectationLoop] evaluate failed at iter %s: %s", i, e)
            return ExpectationLoopResult(
                ready=False,
                iterations=i,
                max_iterations=max_iterations,
                final_score=0,
                regulation_grounded_pass=False,
                stop_reason="evaluate_error",
                history=history,
                remaining_blockers=[str(e)],
                final_report={"error": str(e)},
                state=state,
            )

        rd = report_to_dict(last_report)
        history.append({"iteration": i, "phase": "evaluate", "report": rd})

        mode = str(rd.get("grounding_mode") or state.get("grounding_mode") or "").lower()
        ext = state.get("external_context") if isinstance(state.get("external_context"), dict) else {}
        if not mode:
            mode = str(ext.get("grounding_mode") or "regulation").lower()
            rd["grounding_mode"] = mode

        if mode in ("structure_only", "blocked", "blind"):
            # Explicit: never mark regulation pass; stop early with blockers if blocked
            if mode == "blocked":
                return ExpectationLoopResult(
                    ready=False,
                    iterations=i,
                    max_iterations=max_iterations,
                    final_score=int(rd.get("score") or 0),
                    regulation_grounded_pass=False,
                    stop_reason="blocked_grounding",
                    history=history,
                    remaining_blockers=extract_blockers(last_report)
                    or ["Generacja zablokowana — brak regulaminu i zgody."],
                    final_report=rd,
                    state=state,
                )
            # structure_only: continue loop for structural quality but never ready=True for regulation
            # Fall through to fix attempts, but is_regulation_ready stays False

        if is_regulation_ready(
            rd,
            min_score=min_score,
            require_regulation_grounded=require_regulation_grounded,
        ):
            return ExpectationLoopResult(
                ready=True,
                iterations=i,
                max_iterations=max_iterations,
                final_score=int(rd.get("score") or 0),
                regulation_grounded_pass=bool(rd.get("regulation_grounded_pass")),
                stop_reason="ready",
                history=history,
                remaining_blockers=[],
                final_report=rd,
                state=state,
            )

        if i >= max_iterations:
            break

        if apply_fixes is None:
            # No fixer → cannot improve; stop next iteration boundary
            continue

        try:
            new_state = apply_fixes(state, last_report)
            if isinstance(new_state, dict):
                state = new_state
            history.append(
                {
                    "iteration": i,
                    "phase": "fix",
                    "fixed_sections": list(
                        (new_state or {}).get("fixed_sections")
                        or state.get("last_fixed_sections")
                        or []
                    )
                    if isinstance(new_state, dict)
                    else [],
                }
            )
        except Exception as e:
            logger.warning("[ExpectationLoop] apply_fixes failed at iter %s: %s", i, e)
            history.append({"iteration": i, "phase": "fix_error", "error": str(e)})

    rd = report_to_dict(last_report)
    return ExpectationLoopResult(
        ready=False,
        iterations=max_iterations,
        max_iterations=max_iterations,
        final_score=int(rd.get("score") or 0),
        regulation_grounded_pass=bool(rd.get("regulation_grounded_pass")),
        stop_reason="max_iter",
        history=history,
        remaining_blockers=extract_blockers(last_report),
        final_report=rd,
        state=state,
    )


def run_advisor_expectation_loop(
    state: Dict[str, Any],
    *,
    max_iterations: int = DEFAULT_MAX_ITER,
    min_score: int = DEFAULT_MIN_SCORE,
    apply_fixes: Optional[Callable[[Dict[str, Any], Any], Dict[str, Any]]] = None,
) -> ExpectationLoopResult:
    """
    Default loop using world_class_advisor.evaluate_from_generator_state.
    Optional apply_fixes for rewrite integration (LLM or pure test stub).
    """
    from agents.world_class_advisor import evaluate_from_generator_state

    def _eval(st: Dict[str, Any]):
        return evaluate_from_generator_state(st)

    return run_expectation_loop(
        evaluate=_eval,
        apply_fixes=apply_fixes,
        initial_state=state,
        max_iterations=max_iterations,
        min_score=min_score,
        require_regulation_grounded=True,
    )