backend/core/generation/quality_loop.py

"""
Document quality loop: holistic/panel findings → targeted multi-section rewrite.

Replaces the old nuclear strategy (wipe all sections and regenerate from zero)
with keep-good / rewrite-weak passes for higher pass rates before PDF/DOCX export.
"""
from __future__ import annotations

import difflib
import logging
import os
import re
from typing import Any, Dict, List, Optional, Sequence, Tuple

logger = logging.getLogger(__name__)

DEFAULT_MIN_KEEP_SCORE = int(os.environ.get("QUALITY_LOOP_MIN_SECTION_CHARS", "120"))
MAX_SECTIONS_PER_PASS = int(os.environ.get("QUALITY_LOOP_MAX_SECTIONS_PER_PASS", "8"))


def _norm(s: str) -> str:
    return re.sub(r"\s+", " ", (s or "").lower().strip())


def section_title_match(label: str, candidates: Sequence[str], threshold: float = 0.45) -> Optional[str]:
    """Fuzzy-map free-text section label onto plan titles."""
    if not label or not candidates:
        return None
    nl = _norm(label)
    if nl in ("ogólne", "ogolne", "całość", "calosc", "all", "general", "wniosek", "dokument"):
        return None
    best: Tuple[float, Optional[str]] = (0.0, None)
    for c in candidates:
        nc = _norm(c)
        if not nc:
            continue
        if nl in nc or nc in nl:
            return c
        r = difflib.SequenceMatcher(None, nl, nc).ratio()
        if r > best[0]:
            best = (r, c)
    return best[1] if best[0] >= threshold else None


def extract_section_targets_from_holistic(
    report: Any,
    plan_titles: Sequence[str],
) -> Dict[str, List[str]]:
    """
    Build map title -> list of fix instructions from holistic report.
    """
    targets: Dict[str, List[str]] = {t: [] for t in plan_titles}
    global_notes: List[str] = []

    recs = []
    if hasattr(report, "key_recommendations"):
        recs = list(report.key_recommendations or [])
    elif isinstance(report, dict):
        recs = list(report.get("key_recommendations") or [])

    for rec in recs:
        rec_s = str(rec).strip()
        if not rec_s:
            continue
        matched = section_title_match(rec_s, plan_titles)
        # Try to find section name mentioned in recommendation
        if not matched:
            for t in plan_titles:
                if _norm(t) in _norm(rec_s):
                    matched = t
                    break
        if matched:
            targets.setdefault(matched, []).append(rec_s)
        else:
            global_notes.append(rec_s)

    # Category feedback → map to typical sections
    cat_map = {
        "budget_consistency": ("budżet", "harmonogram", "koszt", "finans"),
        "logical_flow": ("streszczenie", "opis", "cel", "logika"),
        "program_alignment": ("innowacyj", "uzasadn", "dopasow", "program"),
        "dnsh_assessment": ("środowisk", "dnsh", "zrównoważ", "klimat"),
    }
    for cat_name, keywords in cat_map.items():
        cat = getattr(report, cat_name, None) if not isinstance(report, dict) else report.get(cat_name)
        feedback = ""
        flags: List[str] = []
        score = 100
        if cat is not None:
            if hasattr(cat, "feedback"):
                feedback = cat.feedback or ""
                flags = list(getattr(cat, "inconsistencies_flagged", None) or [])
                score = int(getattr(cat, "score", 100) or 100)
            elif isinstance(cat, dict):
                feedback = cat.get("feedback") or ""
                flags = list(cat.get("inconsistencies_flagged") or [])
                score = int(cat.get("score", 100) or 100)
        if score >= 70 and not flags:
            continue
        note = feedback
        if flags:
            note = (note + " " if note else "") + "; ".join(flags[:5])
        if not note:
            continue
        hit = False
        for t in plan_titles:
            nt = _norm(t)
            if any(k in nt for k in keywords):
                targets.setdefault(t, []).append(f"[{cat_name}] {note}")
                hit = True
        if not hit:
            global_notes.append(f"[{cat_name}] {note}")

    # Distribute global notes to weakest sections (shortest content heuristic later)
    if global_notes:
        for t in plan_titles:
            targets.setdefault(t, []).extend(global_notes[:3])

    # Drop empty
    return {k: v for k, v in targets.items() if v}


def extract_section_targets_from_panel_issues(
    issues: Sequence[Any],
    plan_titles: Sequence[str],
) -> Dict[str, List[str]]:
    targets: Dict[str, List[str]] = {}
    for issue in issues or []:
        if isinstance(issue, dict):
            affected = str(issue.get("affected_section") or "Ogólne")
            msg = str(issue.get("message") or "")
            rec = str(issue.get("recommendation") or "")
            sev = str(issue.get("severity") or "?")
            line = f"[{sev}] {msg}" + (f" | Rek: {rec}" if rec else "")
        else:
            affected = str(getattr(issue, "affected_section", "Ogólne") or "Ogólne")
            msg = str(getattr(issue, "message", "") or "")
            rec = str(getattr(issue, "recommendation", "") or "")
            sev = str(getattr(issue, "severity", "?") or "?")
            line = f"[{sev}] {msg}" + (f" | Rek: {rec}" if rec else "")
        matched = section_title_match(affected, plan_titles)
        if not matched:
            # general → all titles get a light note (limited later)
            for t in plan_titles:
                targets.setdefault(t, []).append(line)
        else:
            targets.setdefault(matched, []).append(line)
    return targets


def extract_section_targets_from_compliance_checklist(
    checklist: Any,
    plan_titles: Sequence[str],
) -> Dict[str, List[str]]:
    """Map compliance_checklist.missing_sections onto plan titles (P0 priority)."""
    targets: Dict[str, List[str]] = {}
    if not checklist:
        return targets
    if isinstance(checklist, dict):
        missing = list(checklist.get("missing_sections") or [])
        coverage = checklist.get("coverage_score")
    else:
        missing = list(getattr(checklist, "missing_sections", None) or [])
        coverage = getattr(checklist, "coverage_score", None)
    for raw in missing:
        label = str(raw or "").strip()
        if not label:
            continue
        # Strip " (pusta treść)" suffix used by regulation_checklist
        clean = re.sub(r"\s*\(pusta tre[sś][cć]\)\s*$", "", label, flags=re.I).strip()
        matched = section_title_match(clean, plan_titles) or section_title_match(label, plan_titles)
        note = (
            f"[compliance_checklist] Brakuje wymaganej sekcji regulaminu: {label}. "
            f"Uzupełnij merytorykę i jawne odniesienie do regulaminu programu."
        )
        if coverage is not None:
            note += f" (coverage={coverage}%)"
        if matched:
            targets.setdefault(matched, []).append(note)
        else:
            # Unmapped required section → push to first plan title + global note on all short titles
            for t in plan_titles:
                targets.setdefault(t, []).append(note)
                break
    return targets


def extract_section_targets_from_citation_failures(
    citation_data: Any,
    plan_titles: Sequence[str],
    *,
    min_score: float = 0.72,
) -> Dict[str, List[str]]:
    """Build rewrite targets from citation / faithfulness failures (P0 priority)."""
    targets: Dict[str, List[str]] = {}
    if not citation_data:
        return targets

    # Forms: list of {section, overall_score, issues, recommendation}
    # or dict section->payload, or flat {overall_score, issues, section}
    items: List[Dict[str, Any]] = []
    if isinstance(citation_data, list):
        items = [c for c in citation_data if isinstance(c, dict)]
    elif isinstance(citation_data, dict):
        if any(k in citation_data for k in ("overall_score", "overall_citation_score", "issues", "section")):
            items = [citation_data]
        else:
            for sec, payload in citation_data.items():
                if isinstance(payload, dict):
                    items.append({**payload, "section": payload.get("section") or sec})
                elif isinstance(payload, (int, float)):
                    items.append({"section": sec, "overall_score": float(payload)})

    for item in items:
        score = item.get("overall_score")
        if score is None:
            score = item.get("overall_citation_score")
        try:
            score_f = float(score) if score is not None else None
        except (TypeError, ValueError):
            score_f = None
        issues = item.get("issues") or []
        quality = str(item.get("quality") or item.get("citation_quality") or "")
        weak = (
            (score_f is not None and score_f < min_score)
            or quality.lower() in ("poor", "low", "weak", "fail", "failed")
            or bool(issues)
        )
        if not weak:
            continue
        section = str(item.get("section") or item.get("affected_section") or "Ogólne")
        rec = str(item.get("recommendation") or "").strip()
        issue_bits: List[str] = []
        for iss in (issues if isinstance(issues, list) else [issues])[:4]:
            if isinstance(iss, list):
                issue_bits.extend(str(x) for x in iss[:2] if x)
            elif iss:
                issue_bits.append(str(iss))
        score_txt = f"{score_f:.2f}" if score_f is not None else "?"
        line = (
            f"[citation_failure] Ugruntowanie cytowaniami zbyt niskie (score={score_txt}"
            f"{', quality=' + quality if quality else ''}). "
            f"Dodaj twarde odniesienia do regulaminu/snapshotu i usuń nieugruntowane twierdzenia."
        )
        if issue_bits:
            line += " Problemy: " + "; ".join(issue_bits[:3])
        if rec:
            line += f" | Rek: {rec}"
        matched = section_title_match(section, plan_titles)
        if matched:
            targets.setdefault(matched, []).append(line)
        else:
            for t in plan_titles:
                targets.setdefault(t, []).append(line)
    return targets


def extract_section_targets_from_trap_issues(
    trap_data: Any,
    plan_titles: Sequence[str],
) -> Dict[str, List[str]]:
    """Build rewrite targets from Kruczkowski trap detections (P0 priority)."""
    targets: Dict[str, List[str]] = {}
    if not trap_data:
        return targets

    traps: List[Any] = []
    if isinstance(trap_data, list):
        traps = list(trap_data)
    elif isinstance(trap_data, dict):
        if trap_data.get("detected") is not None:
            traps = list(trap_data.get("detected") or [])
            # Also honor high/critical risk as a doc-level signal when no per-trap list
            risk = str(trap_data.get("risk_level") or trap_data.get("trap_risk") or "").lower()
            if not traps and risk in ("high", "critical"):
                traps = [{"type": "trap_risk", "message": f"trap_risk={risk}", "severity": risk}]
        elif trap_data.get("traps") is not None:
            traps = list(trap_data.get("traps") or [])
        else:
            # section -> payload map
            for sec, payload in trap_data.items():
                if isinstance(payload, dict):
                    detected = payload.get("detected") or payload.get("traps") or []
                    if isinstance(detected, list):
                        for d in detected:
                            if isinstance(d, dict):
                                traps.append({**d, "section": d.get("section") or sec})
                            else:
                                traps.append({"message": str(d), "section": sec})
                    elif payload.get("risk_level") in ("high", "critical", "medium"):
                        traps.append({
                            "section": sec,
                            "severity": payload.get("risk_level"),
                            "message": f"trap_risk={payload.get('risk_level')}",
                        })
                elif isinstance(payload, list):
                    for d in payload:
                        traps.append(d if isinstance(d, dict) else {"message": str(d), "section": sec})

    for trap in traps:
        if isinstance(trap, dict):
            section = str(trap.get("section") or trap.get("affected_section") or "Ogólne")
            ttype = str(trap.get("type") or trap.get("trap_type") or trap.get("category") or "trap")
            msg = str(trap.get("message") or trap.get("description") or trap.get("detail") or ttype)
            sev = str(trap.get("severity") or trap.get("risk_level") or "medium")
            rec = str(trap.get("recommendation") or "").strip()
        else:
            section = "Ogólne"
            ttype = "trap"
            msg = str(trap)
            sev = "medium"
            rec = ""
        line = (
            f"[trap:{ttype}|{sev}] {msg}. "
            f"Usuń niekwalifikowalne koszty/pułapki regulaminowe i jawnie wyklucz ryzyko."
        )
        if rec:
            line += f" | Rek: {rec}"
        matched = section_title_match(section, plan_titles)
        if matched:
            targets.setdefault(matched, []).append(line)
        else:
            for t in plan_titles:
                targets.setdefault(t, []).append(line)
    return targets


def _merge_target_maps(*maps: Dict[str, List[str]]) -> Dict[str, List[str]]:
    """Merge title->notes maps preserving order (first maps win priority position)."""
    out: Dict[str, List[str]] = {}
    for m in maps:
        if not m:
            continue
        for title, notes in m.items():
            bucket = out.setdefault(title, [])
            for n in notes or []:
                if n and n not in bucket:
                    bucket.append(n)
    return out


def build_priority_rewrite_targets(
    plan_titles: Sequence[str],
    *,
    compliance_checklist: Any = None,
    citation_failures: Any = None,
    trap_issues: Any = None,
    report: Any = None,
    panel_issues: Optional[Sequence[Any]] = None,
    source: str = "holistic",
    advisor_report: Any = None,
) -> Dict[str, List[str]]:
    """
    P0 target order: advisor brief gaps → compliance → citation → traps → holistic/panel.
    Compliance/citation/trap notes are placed first so pick_sections_to_fix prioritizes them.
    """
    advisor_targets: Dict[str, List[str]] = {}
    if advisor_report is not None:
        try:
            from agents.world_class_advisor import advisor_findings_to_rewrite_targets

            advisor_targets = advisor_findings_to_rewrite_targets(advisor_report, plan_titles)
        except Exception as e:
            logger.debug("[QualityLoop] advisor targets skipped: %s", e)

    compliance_targets = extract_section_targets_from_compliance_checklist(
        compliance_checklist, plan_titles
    )
    citation_targets = extract_section_targets_from_citation_failures(
        citation_failures, plan_titles
    )
    trap_targets = extract_section_targets_from_trap_issues(trap_issues, plan_titles)

    secondary: Dict[str, List[str]] = {}
    if source == "holistic" and report is not None:
        secondary = extract_section_targets_from_holistic(report, plan_titles)
    elif panel_issues is not None:
        secondary = extract_section_targets_from_panel_issues(panel_issues, plan_titles)
    elif source == "panel" and report is not None:
        # tolerate report-as-issues misuse
        secondary = extract_section_targets_from_panel_issues(
            getattr(report, "issues", None) or [], plan_titles
        )

    return _merge_target_maps(
        advisor_targets,
        compliance_targets,
        citation_targets,
        trap_targets,
        secondary,
    )


def collect_grounding_signals_from_state(state: Dict[str, Any]) -> Dict[str, Any]:
    """Harvest checklist / citation / trap signals from generator state + external_context."""
    ext = state.get("external_context") if isinstance(state.get("external_context"), dict) else {}
    checklist = (
        state.get("compliance_checklist")
        or ext.get("compliance_checklist")
        or {}
    )
    # Citations: prefer explicit state keys, else aggregate v5_verification from traceability
    citations = state.get("citation_failures") or state.get("citation_scores") or ext.get("citation_failures")
    traps = state.get("trap_issues") or ext.get("trap_issues") or ext.get("v5_grounding_certificate")

    if citations is None or traps is None:
        trace = state.get("traceability_data") or {}
        cit_list: List[Dict[str, Any]] = []
        trap_map: Dict[str, Any] = {}
        for section_key, events in (trace.items() if isinstance(trace, dict) else []):
            if not isinstance(events, list):
                continue
            for ev in events:
                if not isinstance(ev, dict):
                    continue
                if ev.get("type") != "v5_verification":
                    continue
                data = ev.get("data") or {}
                if not isinstance(data, dict):
                    continue
                sec_name = data.get("section") or section_key
                cit = data.get("citation") or {}
                if citations is None and isinstance(cit, dict):
                    cit_list.append({
                        "section": sec_name,
                        "overall_score": cit.get("overall_score"),
                        "quality": cit.get("quality"),
                        "issues": cit.get("issues") or [],
                        "recommendation": cit.get("recommendation") or "",
                    })
                tr = data.get("traps") or {}
                if traps is None and isinstance(tr, dict):
                    trap_map[str(sec_name)] = tr
        if citations is None and cit_list:
            citations = cit_list
        if traps is None and trap_map:
            traps = trap_map

    # Certificate-level trap signal
    if traps is None:
        v5c = ext.get("v5_grounding_certificate") or {}
        if isinstance(v5c, dict) and (v5c.get("trap_risk") or v5c.get("detected")):
            traps = v5c

    return {
        "compliance_checklist": checklist,
        "citation_failures": citations,
        "trap_issues": traps,
    }


def pick_sections_to_fix(
    plan: Sequence[Any],
    generated: Dict[str, str],
    targets: Dict[str, List[str]],
    *,
    max_sections: int = MAX_SECTIONS_PER_PASS,
) -> List[str]:
    """Prefer targeted weak sections; if none, pick shortest / incomplete ones."""
    titles = []
    for s in plan:
        t = s.get("title") if isinstance(s, dict) else str(s)
        if t:
            titles.append(t)

    ranked: List[Tuple[int, str]] = []
    for t in titles:
        notes = targets.get(t) or []
        content = (generated or {}).get(t) or ""
        incomplete = 1 if ("[UZUPEŁNIĆ" in content or "[DO WERYFIKACJI" in content or len(content) < DEFAULT_MIN_KEEP_SCORE) else 0
        priority = len(notes) * 10 + incomplete * 5 + max(0, 500 - len(content)) // 50
        if notes or incomplete:
            ranked.append((priority, t))
    ranked.sort(key=lambda x: x[0], reverse=True)
    chosen = [t for _, t in ranked[:max_sections]]
    if not chosen and titles:
        # Always fix at least top-N shortest when critic failed without mapping
        by_len = sorted(titles, key=lambda t: len((generated or {}).get(t) or ""))
        chosen = by_len[: min(3, max_sections)]
    return chosen


def rewrite_section_with_feedback(
    *,
    title: str,
    section_type: str,
    current_content: str,
    instructions: List[str],
    program_name: str,
    project_description: str = "",
    company_context: str = "",
    external_context: Optional[dict] = None,
    regulation_boost: str = "",
) -> str:
    """LLM rewrite of one section with critic/audit instructions + regulation grounding."""
    from agents.helpers import generate_section_light
    from core.llm_router import get_llm
    from langchain_core.messages import HumanMessage

    instr = "\n".join(f"- {i}" for i in (instructions or [])[:12])
    if not instr:
        instr = "- Podnieś merytorykę, spójność z resztą wniosku i zgodność z programem."

    reg_block = (regulation_boost or "").strip()
    if reg_block:
        reg_block = reg_block[:6000]
        reg_section = f"""
Kontekst regulaminowy (ŹRÓDŁO PRAWDY — cytuj reguły, nie wymyślaj):
--------------------
{reg_block}
--------------------
Wzmocnij ugruntowanie: jawne odniesienia do regulaminu/snapshotu, zero niekwalifikowalnych kosztów.
"""
    else:
        reg_section = (
            "\nBrak snapshotu regulaminu w kontekście — unikaj kategorycznych twierdzeń o "
            "kwalifikowalności; oznacz niepewne miejsca [DO WERYFIKACJI: regulamin].\n"
        )

    # Prefer focused rewrite when content exists
    if current_content and len(current_content.strip()) > 80:
        llm = get_llm(task_type="writing")
        prompt = f"""Jesteś redaktorem wniosków unijnych. PRZEPISZ sekcję „{title}" tak, aby usunąć wskazane wady.
Zachowaj język polski, styl urzędowy, Markdown. NIE wymyślaj faktów spoza kontekstu.
Dla braków danych użyj [DO WERYFIKACJI: …] — nie zostawiaj pustych miejsc.
Priorytet: checklist regulaminu, cytowania/ugruntowanie, pułapki Kruczkowskiego.

Program: {program_name}
Opis projektu (skrót):
{(project_description or '')[:2500]}

Dane firmy / kontekst:
{(company_context or '')[:2000]}
{reg_section}
Wady / instrukcje poprawy:
{instr}

Obecna treść sekcji:
--------------------
{current_content[:12000]}
--------------------

Zwróć WYŁĄCZNIE poprawioną treść sekcji (bez preambuły)."""
        try:
            resp = llm.invoke([HumanMessage(content=prompt)])
            content = resp.content if hasattr(resp, "content") else str(resp)
            if content and len(content.strip()) > 40:
                return content.strip()
        except Exception as e:
            logger.warning("[QualityLoop] rewrite failed for %s: %s", title, e)

    # Fallback: generate_section_light with feedback + regulation in context
    ctx_parts = [project_description or ""]
    if reg_block:
        ctx_parts.append(reg_block)
    ctx_parts.append(f"INSTRUKCJE POPRAWY SEKCJI:\n{instr}")
    ctx_parts.append(f"Poprzednia treść (do ulepszenia):\n{(current_content or '')[:4000]}")
    ctx = "\n\n".join(p for p in ctx_parts if p)
    return generate_section_light(
        section_type=section_type or title,
        context=ctx,
        external_context=external_context or {},
        program_name=program_name,
        light_mode=False,
    )


def run_quality_expectation_step(
    state: Dict[str, Any],
    *,
    source: str = "holistic",
    report: Any = None,
    issues: Optional[Sequence[Any]] = None,
    regulation_boost: str = "",
    min_score: int = 70,
) -> Dict[str, Any]:
    """
    One production expectation step: advisor evaluate → if not regulation-ready,
    apply targeted fixes → re-evaluate. Used by generator quality path.

    Returns dict with advisor_before/after, expectation ready flags, fixed state.
    """
    from agents.world_class_advisor import evaluate_from_generator_state
    from core.generation.expectation_loop import (
        extract_blockers,
        is_regulation_ready,
        report_to_dict,
    )

    before = evaluate_from_generator_state(state)
    ready = is_regulation_ready(before, min_score=min_score)
    out: Dict[str, Any] = {
        "advisor_before": report_to_dict(before),
        "regulation_ready": ready,
        "remaining_blockers": extract_blockers(before),
        "fixed_sections": [],
        "generated_sections": dict(state.get("generated_sections") or {}),
    }
    if ready:
        out["stop_reason"] = "ready"
        out["advisor_after"] = out["advisor_before"]
        return out

    fixed = apply_targeted_section_fixes(
        state,
        source=source,
        report=report,
        issues=issues,
        regulation_boost=regulation_boost,
        advisor_report=before,
    )
    new_state = {**state, "generated_sections": fixed.get("generated_sections") or state.get("generated_sections")}
    after = evaluate_from_generator_state(new_state)
    out["advisor_after"] = report_to_dict(after)
    out["regulation_ready"] = is_regulation_ready(after, min_score=min_score)
    out["remaining_blockers"] = extract_blockers(after)
    out["fixed_sections"] = list(fixed.get("fixed_sections") or [])
    out["generated_sections"] = fixed.get("generated_sections") or out["generated_sections"]
    out["targets"] = fixed.get("targets") or {}
    out["stop_reason"] = "ready" if out["regulation_ready"] else "needs_retry"
    return out


def apply_targeted_section_fixes(
    state: Dict[str, Any],
    *,
    source: str,
    report: Any = None,
    issues: Optional[Sequence[Any]] = None,
    regulation_boost: str = "",
    advisor_report: Any = None,
) -> Dict[str, Any]:
    """
    Returns updated generated_sections (partial rewrite) + metadata for telemetry.

    Rewrite targets are built in P0 order:
      world-class advisor → compliance → citation → traps → holistic/panel.
    regulation_boost (if provided or buildable) is injected into each section rewrite.
    """
    plan = state.get("sections_plan") or []
    generated = dict(state.get("generated_sections") or {})
    titles = []
    type_by_title: Dict[str, str] = {}
    for s in plan:
        if isinstance(s, dict):
            t = s.get("title") or s.get("type") or ""
            titles.append(t)
            type_by_title[t] = s.get("type") or t
        else:
            titles.append(str(s))
            type_by_title[str(s)] = str(s)

    signals = collect_grounding_signals_from_state(state)
    if advisor_report is None:
        # Only auto-run world-class advisor when regulation brief signals exist —
        # avoid rewriting all sections solely for empty-brief noise.
        ext0 = state.get("external_context") if isinstance(state.get("external_context"), dict) else {}
        has_brief_signals = bool(
            ext0.get("advisor_brief")
            or ext0.get("required_sections")
            or ext0.get("regulation_key_rules")
            or ext0.get("key_rules")
            or ext0.get("attention_points")
        )
        if has_brief_signals:
            try:
                from agents.world_class_advisor import evaluate_from_generator_state

                advisor_report = evaluate_from_generator_state(state)
            except Exception as e:
                logger.debug("[QualityLoop] advisor evaluate skipped: %s", e)
                advisor_report = None

    targets = build_priority_rewrite_targets(
        titles,
        compliance_checklist=signals.get("compliance_checklist"),
        citation_failures=signals.get("citation_failures"),
        trap_issues=signals.get("trap_issues"),
        report=report,
        panel_issues=issues,
        source=source,
        advisor_report=advisor_report,
    )

    to_fix = pick_sections_to_fix(plan, generated, targets)
    program = state.get("document_type") or "wniosek dotacyjny"
    project_desc = state.get("project_description") or ""
    company_ctx = state.get("additional_context") or ""
    ext = state.get("external_context") if isinstance(state.get("external_context"), dict) else {}

    boost = (regulation_boost or state.get("regulation_boost") or "").strip()
    if not boost:
        # Lightweight inline boost from external_context (caller may pass full agent boost)
        snap_id = ext.get("regulation_snapshot_id")
        rules = ext.get("regulation_key_rules") or ext.get("key_rules") or []
        if rules:
            boost = (
                "[REGULATION SNAPSHOT v5.0 - QUALITY LOOP]:\n"
                + "\n".join(f"- {r}" for r in list(rules)[:8])
            )
        elif snap_id:
            boost = f"[REGULATION SNAPSHOT v5.0 - id={snap_id}]\nUżyj reguł z przypisanego snapshotu regulaminu."
        elif ext.get("required_sections"):
            boost = (
                "[REGULATION CONTEXT - required_sections]:\n"
                + "\n".join(f"- Wymagana sekcja: {s}" for s in list(ext.get("required_sections") or [])[:12])
            )

    fixed: List[str] = []
    for title in to_fix:
        notes = targets.get(title) or ["Podnieś jakość i spójność z całym wnioskiem."]
        new_text = rewrite_section_with_feedback(
            title=title,
            section_type=type_by_title.get(title, title),
            current_content=generated.get(title) or "",
            instructions=notes,
            program_name=program,
            project_description=project_desc,
            company_context=company_ctx,
            external_context=ext,
            regulation_boost=boost,
        )
        if new_text and len(new_text.strip()) > 40:
            generated[title] = new_text.strip()
            fixed.append(title)
            logger.info("[QualityLoop] rewritten section '%s' (%s notes)", title, len(notes))

    return {
        "generated_sections": generated,
        "fixed_sections": fixed,
        "targets": {k: v[:5] for k, v in targets.items() if k in to_fix},
        "source": source,
        "regulation_boost_used": bool(boost),
        "priority_signals": {
            "checklist_missing": len(
                (signals.get("compliance_checklist") or {}).get("missing_sections") or []
            )
            if isinstance(signals.get("compliance_checklist"), dict)
            else 0,
            "citation_items": len(signals.get("citation_failures") or [])
            if isinstance(signals.get("citation_failures"), list)
            else (1 if signals.get("citation_failures") else 0),
            "trap_items": (
                len((signals.get("trap_issues") or {}).get("detected") or [])
                if isinstance(signals.get("trap_issues"), dict)
                else len(signals.get("trap_issues") or [])
                if isinstance(signals.get("trap_issues"), list)
                else 0
            ),
        },
    }


def score_document_readiness(
    generated: Dict[str, str],
    plan: Sequence[Any],
    *,
    compliance_checklist: Any = None,
    citation_scores: Any = None,
    trap_issues: Any = None,
    regulation_context_present: Optional[bool] = None,
    external_context: Optional[dict] = None,
) -> Dict[str, Any]:
    """Heuristic readiness 0-100 for export soft-gate decisions, with grounding fields."""
    incomplete = []
    ok = 0
    for s in plan or []:
        t = s.get("title") if isinstance(s, dict) else str(s)
        c = (generated or {}).get(t) or ""
        if len(c) < DEFAULT_MIN_KEEP_SCORE or "[UZUPEŁNIĆ" in c:
            incomplete.append(t)
        else:
            ok += 1
    total = len(plan or [])
    score = int(round(100 * ok / max(total, 1))) if total else 0

    ext = external_context if isinstance(external_context, dict) else {}
    checklist = compliance_checklist if compliance_checklist is not None else ext.get("compliance_checklist")
    coverage: Optional[int] = None
    missing_sections: List[str] = []
    if isinstance(checklist, dict):
        if checklist.get("coverage_score") is not None:
            try:
                coverage = int(checklist.get("coverage_score"))
            except (TypeError, ValueError):
                coverage = None
        missing_sections = list(checklist.get("missing_sections") or [])
    elif checklist is not None:
        coverage = getattr(checklist, "coverage_score", None)
        missing_sections = list(getattr(checklist, "missing_sections", None) or [])

    # Aggregate citation mean
    cit_raw = citation_scores if citation_scores is not None else ext.get("citation_failures")
    citation_values: List[float] = []
    if isinstance(cit_raw, list):
        for item in cit_raw:
            if isinstance(item, dict):
                sc = item.get("overall_score", item.get("overall_citation_score"))
                if sc is not None:
                    try:
                        citation_values.append(float(sc))
                    except (TypeError, ValueError):
                        pass
            elif isinstance(item, (int, float)):
                citation_values.append(float(item))
    elif isinstance(cit_raw, dict):
        if "overall_score" in cit_raw or "overall_citation_score" in cit_raw:
            sc = cit_raw.get("overall_score", cit_raw.get("overall_citation_score"))
            try:
                citation_values.append(float(sc))
            except (TypeError, ValueError):
                pass
        else:
            for v in cit_raw.values():
                if isinstance(v, dict):
                    sc = v.get("overall_score", v.get("overall_citation_score"))
                    if sc is not None:
                        try:
                            citation_values.append(float(sc))
                        except (TypeError, ValueError):
                            pass
                elif isinstance(v, (int, float)):
                    citation_values.append(float(v))
    citation_mean = (
        round(sum(citation_values) / len(citation_values), 3) if citation_values else None
    )
    citation_ok = citation_mean is not None and citation_mean >= 0.72

    traps = trap_issues if trap_issues is not None else (
        ext.get("trap_issues") or ext.get("v5_grounding_certificate")
    )
    trap_risk = "unknown"
    trap_count = 0
    if isinstance(traps, dict):
        trap_risk = str(traps.get("risk_level") or traps.get("trap_risk") or "unknown")
        detected = traps.get("detected") or traps.get("traps") or []
        trap_count = len(detected) if isinstance(detected, list) else 0
    elif isinstance(traps, list):
        trap_count = len(traps)
        trap_risk = "medium" if trap_count else "low"

    if regulation_context_present is None:
        # Infer from external_context / generated content markers
        if ext.get("regulation_snapshot_id") or ext.get("required_sections"):
            regulation_context_present = True
        else:
            blob = " ".join((generated or {}).values())[:8000]
            regulation_context_present = (
                "[REGULATION SNAPSHOT" in blob
                or "Kontekst regulaminowy" in blob
                or bool(ext.get("v5_grounding_certificate"))
            )

    checklist_ok = coverage is None or coverage >= 70
    grounding_ok = bool(regulation_context_present) and checklist_ok and (
        citation_mean is None or citation_ok
    ) and trap_risk not in ("high", "critical")

    return {
        "score": score,
        "complete_sections": ok,
        "total": total,
        "incomplete": incomplete,
        # Grounding fields (P0 export soft-gate)
        "regulation_context_present": bool(regulation_context_present),
        "checklist_coverage": coverage,
        "checklist_ok": checklist_ok,
        "missing_sections": missing_sections[:12],
        "citation_mean": citation_mean,
        "citation_ok": citation_ok if citation_mean is not None else None,
        "trap_risk": trap_risk,
        "trap_count": trap_count,
        "grounding_ok": grounding_ok,
    }