src/orchestrator.py

from __future__ import annotations

import re
from dataclasses import asdict, dataclass, field
from typing import Any, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING

from llm_client_groq import classify_question, is_meta_question

if TYPE_CHECKING:
    from llm_client_groq import ConversationMemory
else:
    ConversationMemory = Any  # type: ignore[assignment]

# ---------------------------------------------------------------------------
# Regexes and lightweight legal intent detection
# ---------------------------------------------------------------------------

NORM_REF_RE = re.compile(
    r"§{1,2}\s*(?P<section>\d+[a-zA-Z]?)"
    r"(?:\s*(?:-|–|bis)\s*(?P<section_to>\d+[a-zA-Z]?))?"
    r"(?:\s*(?:Abs\.|Absatz)\s*(?P<subsection>\d+[a-zA-Z]?))?"
    r"(?:\s*Satz\s*(?P<sentence>\d+[a-zA-Z]?))?"
    r"(?:\s*(?:Nr\.|Nummer)\s*(?P<number>\d+[a-zA-Z]?))?"
    r"(?:\s*(?:Buchst\.|Buchstabe|lit\.)\s*(?P<letter>[a-zA-Z]))?",
    re.I,
)

SOURCE_MARKER_RE = re.compile(r"\[Quelle\s+(\d+)(?:[^\]]*)\]", re.I)
SOURCE_MARKER_WITH_OPTIONAL_REF_RE = re.compile(
    r"\[Quelle\s+(\d+)(?:[^\]]*)\]"
    r"(?:\s*\(§[^)]{1,120}\))?",
    re.I,
)
ORPHAN_LEGAL_REF_PAREN_RE = re.compile(r"\(§\s*\d{1,3}[a-z]?(?:\s+[^)]{0,80})?\)", re.I)
NEGATIVE_ANSWER_RE = re.compile(
    r"\b(keine\s+(relevante\s+)?textstelle|nichts?\s+(geregelt|enthalten|auffindbar)|"
    r"keine\s+(regelung|informationen|aussage)|nicht\s+belastbar\s+ableitbar|"
    r"kann\s+.*?nicht\s+(festgestellt|beantwortet)\s+werden)\b",
    re.I,
)

# Fine-grained citation fragments generated by the LLM must be verified against
# actual source metadata.  In practice, models often over-specialise references
# such as "§ 6 Abs. 1 Buchst. a" although the underlying chunk only supports
# "§ 6 Abs. 1".  The orchestrator therefore downgrades unsupported fine refs
# instead of passing them through as if they were verified citations.
FINE_BUCHST_REF_RE = re.compile(
    r"§\s*(?P<section>\d{1,3}[a-z]?)\s+"
    r"(?:Abs\.|Absatz)\s*(?P<subsection>\d{1,3}[a-z]?)\s+"
    r"(?:Buchst\.|Buchstabe|lit\.)\s*(?P<letter>[a-z])",
    re.I,
)

CITATION_PAREN_RE = re.compile(
    r"(?P<marker>\[Quelle\s+\d+(?:[^\]]*)\])\s*"
    r"\((?P<ref>§[^)]{1,120})\)",
    re.I,
)

DANGLING_CITATION_GRAMMAR_REPLACEMENTS: Tuple[Tuple[re.Pattern[str], str], ...] = (
    # Remove orphan legal-ref parentheses that remain after invalid source markers
    # were stripped, e.g. "[Quelle 1] (§ 6) und (§ 6)".
    (re.compile(r"(\[Quelle\s+\d+(?:[^\]]*)\](?:\s*\(§[^)]{1,120}\))?)\s+und\s+\(§[^)]{1,120}\)", re.I), r"\1"),
    (re.compile(r"(\[Quelle\s+\d+(?:[^\]]*)\](?:\s*\(§[^)]{1,120}\))?)\s*,\s*\(§[^)]{1,120}\)", re.I), r"\1"),
    (re.compile(r"\s+(?:und|oder)\s+\(§[^)]{1,120}\)", re.I), r""),
    (re.compile(r"\s+und\s+(beschrieben|genannt|geregelt|dargelegt|aufgeführt)\b", re.I), r" \1"),
    (re.compile(r"\bwie\s+in\s+([^.;:\n]{1,160}?)\s+und\s+(beschrieben|genannt|geregelt|dargelegt)\b", re.I), r"wie in \1 \2"),
    (re.compile(r"\bdie\s+in\s+den\s+Quellen\s+([^.;:\n]{1,160}?)\s+und\s+genannt\s+sind", re.I), r"die in \1 genannt sind"),
    (re.compile(r"\bdie\s+in\s+([^.;:\n]{1,160}?)\s+und\s+genannt\s+sind", re.I), r"die in \1 genannt sind"),
    (re.compile(r"\bund\s+weiteren\s+Quellen\s+(dargelegt|beschrieben|genannt)\s+(sind|ist)", re.I), r"\1 \2"),
)

DEFINITION_RE = re.compile(
    r"\b(was\s+(versteht|bedeutet)|wie\s+definiert|definition|legaldefinition|"
    r"begriff|unter\s+[„\"']?[^?]+[”\"']?\s+versteht)\b",
    re.I,
)
ENUMERATION_RE = re.compile(
    r"\b(welche|alle|sämtliche|liste|nennt|kriterien|voraussetzungen|tatbestandsmerkmale|"
    r"bestandteile|anforderungen|fälle|maßnahmen|regelt\s+§)\b",
    re.I,
)
COMPARISON_RE = re.compile(r"\b(unterschied|vergleiche|vergleich|gegenüber|vs\.?|versus)\b", re.I)
CLARIFICATION_RISK_RE = re.compile(
    r"\b(das|dies|diese|der\s+fall|so\s+ein\s+fall|dort|hierbei|teilnahme|anspruch|"
    r"abrechnung|pflicht|folge|konsequenz)\b",
    re.I,
)

BUILTIN_QUERY_EXPANSIONS: Dict[str, List[str]] = {
    "nicht verfügbar": ["Nichtverfügbarkeit", "lieferfähig", "nicht lieferbar", "Lieferengpass", "Verfügbarkeit"],
    "beitritt": ["teilnehmen", "Teilnahme", "Mitgliedsverband", "DAV", "Erklärung", "beitreten"],
    "auseinzelung": ["Teilmenge", "Auseinzelung", "einzelne Einheit", "Packung", "Entnahme"],
    "pharmazeutische dienstleistungen": ["pharmazeutische Dienstleistungen", "Anlage 11", "Anspruchsvoraussetzungen", "Vergütung", "Abrechnung"],
    "wunscharzneimittel": ["Wunscharzneimittel", "Kostenerstattung", "anderes Fertigarzneimittel", "§§ 11 bis 14"],
}


# ---------------------------------------------------------------------------
# Public dataclasses
# ---------------------------------------------------------------------------

@dataclass(slots=True, frozen=True)
class NormReference:
    section: str
    section_to: str | None = None
    subsection: str | None = None
    sentence: str | None = None
    number: str | None = None
    letter: str | None = None

    @property
    def section_id(self) -> str:
        return f"§ {self.section}"

    @property
    def canonical_ref(self) -> str:
        parts = [self.section_id]
        if self.subsection:
            parts.append(f"Abs. {self.subsection}")
        if self.sentence:
            parts.append(f"Satz {self.sentence}")
        if self.number:
            parts.append(f"Nr. {self.number}")
        if self.letter:
            parts.append(f"Buchst. {self.letter.lower()}")
        return " ".join(parts)


@dataclass(slots=True)
class QuestionAnalysis:
    question: str
    question_kind: str
    intent: str
    norm_references: List[NormReference] = field(default_factory=list)
    query_terms: List[str] = field(default_factory=list)
    expanded_terms: List[str] = field(default_factory=list)
    needs_clarification: bool = False
    clarification_question: str | None = None
    reasons: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "question": self.question,
            "question_kind": self.question_kind,
            "intent": self.intent,
            "norm_references": [asdict(ref) | {"canonical_ref": ref.canonical_ref} for ref in self.norm_references],
            "query_terms": list(self.query_terms),
            "expanded_terms": list(self.expanded_terms),
            "needs_clarification": self.needs_clarification,
            "clarification_question": self.clarification_question,
            "reasons": list(self.reasons),
        }


@dataclass(slots=True)
class RetrievalAssessment:
    hit_count: int = 0
    direct_hit_count: int = 0
    exact_norm_hit_count: int = 0
    definition_hit_count: int = 0
    parent_context_count: int = 0
    neighbor_count: int = 0
    containers: List[str] = field(default_factory=list)
    sections: List[str] = field(default_factory=list)
    canonical_refs: List[str] = field(default_factory=list)
    low_confidence: bool = False
    negative_recheck_performed: bool = False
    negative_recheck_added_hits: int = 0
    reasons: List[str] = field(default_factory=list)

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


@dataclass(slots=True)
class AnswerAudit:
    cited_source_numbers: List[int] = field(default_factory=list)
    invalid_source_numbers: List[int] = field(default_factory=list)
    cited_canonical_refs: List[str] = field(default_factory=list)
    answer_basis: str = "unknown"  # explicit | derived | negative | insufficient | unknown
    negative_answer_detected: bool = False
    completeness_warnings: List[str] = field(default_factory=list)
    citation_warnings: List[str] = field(default_factory=list)
    auto_fixes: List[str] = field(default_factory=list)
    recommended_action: str = "accept"  # accept | recheck | clarify | caution
    confidence: float = 0.5

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


@dataclass(slots=True)
class OrchestratorOptions:
    top_k: int = 8
    fetch_k: int = 24
    max_final_results: int = 10
    max_sources: int = 5
    include_neighbors: bool = True
    include_explicit_sections: bool = True
    # False: das Ranking priorisiert den Hauptvertrag; ein harter Filter machte
    # Anlagen/Anhänge (Anlage 11, Dienstleistungs-Anhänge) unauffindbar.
    restrict_to_default_container: bool = False
    default_container_id: str = "Vertrag"
    min_score: float | None = None
    enable_negative_recheck: bool = True
    enable_clarification: bool = True
    enable_answer_audit: bool = True
    enrich_citations_with_canonical_refs: bool = True
    debug: bool = False


@dataclass(slots=True)
class OrchestratorResult:
    answer: str
    answer_type: str
    hits: List[Dict[str, Any]] = field(default_factory=list)
    raw_sources: List[Dict[str, Any]] = field(default_factory=list)
    sources: List[Dict[str, Any]] = field(default_factory=list)
    analysis: QuestionAnalysis | None = None
    retrieval_assessment: RetrievalAssessment | None = None
    answer_audit: AnswerAudit | None = None
    needs_clarification: bool = False
    clarification_question: str | None = None
    debug: Dict[str, Any] = field(default_factory=dict)

    def to_debug_dict(self) -> Dict[str, Any]:
        return {
            "analysis": self.analysis.to_dict() if self.analysis else None,
            "retrieval_assessment": self.retrieval_assessment.to_dict() if self.retrieval_assessment else None,
            "answer_audit": self.answer_audit.to_dict() if self.answer_audit else None,
            "needs_clarification": self.needs_clarification,
            "clarification_question": self.clarification_question,
            "debug": dict(self.debug),
        }


# ---------------------------------------------------------------------------
# Helper functions
# ---------------------------------------------------------------------------

def _normalize(text: str) -> str:
    return " ".join((text or "").strip().split())


def _lower(text: str) -> str:
    return _normalize(text).lower()


def _hit_meta(hit: Dict[str, Any]) -> Dict[str, Any]:
    meta = hit.get("metadata") or {}
    return meta if isinstance(meta, dict) else {}


def _hit_value(hit: Dict[str, Any], *keys: str, default: Any = None) -> Any:
    meta = _hit_meta(hit)
    for key in keys:
        value = hit.get(key)
        if value is not None and value != "":
            return value
    for key in keys:
        value = meta.get(key)
        if value is not None and value != "":
            return value
    return default


def _hit_text(hit: Dict[str, Any]) -> str:
    return str(hit.get("text") or hit.get("document") or "").strip()


def _hit_kinds(hit: Dict[str, Any]) -> List[str]:
    kinds = hit.get("retrieval_kinds") or []
    if isinstance(kinds, str):
        return [kinds]
    return [str(k) for k in kinds]


def _hit_score(hit: Dict[str, Any]) -> float:
    try:
        return float(hit.get("rank_score", hit.get("score", 0.0)) or 0.0)
    except (TypeError, ValueError):
        return 0.0


def _source_key(hit: Dict[str, Any]) -> Tuple[Any, ...]:
    meta = _hit_meta(hit)
    text_hash = meta.get("text_hash") or hit.get("text_hash")
    if text_hash:
        return ("hash", text_hash)
    legal_unit_id = meta.get("legal_unit_id") or hit.get("legal_unit_id")
    if legal_unit_id:
        return ("legal_unit", legal_unit_id)
    return (
        _hit_value(hit, "container", "container_id", default=""),
        _hit_value(hit, "section", "section_id", default=""),
        _hit_value(hit, "chunk_index", "chunk_index_in_section", default=""),
        _hit_text(hit)[:160],
    )


def _canonical_ref_from_hit(hit: Dict[str, Any]) -> str:
    meta = _hit_meta(hit)
    direct = _hit_value(hit, "canonical_ref", default=None)
    if direct:
        return str(direct)

    paragraph = meta.get("paragraph") or hit.get("paragraph") or _hit_value(hit, "section", "section_id", default="")
    subsection = meta.get("subsection") or hit.get("subsection")
    sentence = meta.get("sentence") or hit.get("sentence")
    number = meta.get("number") or hit.get("number")
    letter = meta.get("letter") or hit.get("letter")

    parts = [str(paragraph)] if paragraph else []
    if subsection:
        parts.append(f"Abs. {subsection}")
    if sentence:
        parts.append(f"Satz {sentence}")
    if number:
        parts.append(f"Nr. {number}")
    if letter:
        parts.append(f"Buchst. {str(letter).lower()}")
    return " ".join(parts)


def _dedupe_hits(hits: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]:
    merged: Dict[Tuple[Any, ...], Dict[str, Any]] = {}
    for hit in hits:
        key = _source_key(hit)
        old = merged.get(key)
        if old is None or _hit_score(hit) > _hit_score(old):
            merged[key] = dict(hit)
        else:
            old_kinds = set(_hit_kinds(old))
            old_kinds.update(_hit_kinds(hit))
            old["retrieval_kinds"] = sorted(old_kinds)
    out = list(merged.values())
    out.sort(key=lambda h: (_hit_score(h), "parent_context" in set(_hit_kinds(h))), reverse=True)
    return out


def extract_norm_references(question: str, *, max_range: int = 30) -> List[NormReference]:
    refs: List[NormReference] = []
    for match in NORM_REF_RE.finditer(question or ""):
        start = match.group("section")
        end = match.group("section_to")
        if end and start.isdigit() and end.isdigit():
            start_i, end_i = int(start), int(end)
            if start_i <= end_i and (end_i - start_i) <= max_range:
                for no in range(start_i, end_i + 1):
                    refs.append(NormReference(section=str(no)))
                continue
        refs.append(
            NormReference(
                section=start,
                section_to=end,
                subsection=match.group("subsection"),
                sentence=match.group("sentence"),
                number=match.group("number"),
                letter=match.group("letter"),
            )
        )
    seen: set[str] = set()
    unique: List[NormReference] = []
    for ref in refs:
        key = ref.canonical_ref
        if key not in seen:
            seen.add(key)
            unique.append(ref)
    return unique


def expand_query_terms(question: str) -> List[str]:
    q = _lower(question)
    terms: List[str] = []
    for key, values in BUILTIN_QUERY_EXPANSIONS.items():
        if key in q or any(v.lower() in q for v in values):
            terms.extend([key, *values])
    # Extract quoted terms as high-value legal-definition candidates.
    for quoted in re.findall(r"[„\"']([^„”\"']{2,80})[”\"']", question or ""):
        terms.append(quoted.strip())
    return list(dict.fromkeys(t for t in terms if t))


# ---------------------------------------------------------------------------
# Orchestrator
# ---------------------------------------------------------------------------

class LegalAnswerOrchestrator:
    """
    Fachliche Ablaufsteuerung für die juristische RAG-Antwort.

    Diese Klasse ist bewusst unabhängig von FastAPI. `app.py` bleibt die
    Aussteuerungs- und API-Schicht und kann diesen Orchestrator als einen
    Knoten im Ask-Graph verwenden.

    Erwartete externe Komponenten:
    - retriever: besitzt idealerweise `query(...)` und optional `verify_negative_result(...)`
    - composer: besitzt `compose_with_sources(...)` oder `compose(...)`
    """

    def __init__(self, retriever: Any, composer: Any, *, options: OrchestratorOptions | None = None):
        self.retriever = retriever
        self.composer = composer
        self.options = options or OrchestratorOptions()

    # ------------------------------------------------------------------
    # Analysis and planning
    # ------------------------------------------------------------------

    def analyze_question(self, question: str) -> QuestionAnalysis:
        q = _normalize(question)
        lower = q.lower()
        refs = extract_norm_references(q)
        expanded = expand_query_terms(q)
        reasons: List[str] = []

        if is_meta_question(q):
            intent = "meta"
            reasons.append("meta_question")
        elif DEFINITION_RE.search(q):
            intent = "definition"
            reasons.append("definition_pattern")
        elif refs and ENUMERATION_RE.search(q):
            intent = "norm_enumeration"
            reasons.append("explicit_norm_and_enumeration_pattern")
        elif refs:
            intent = "explicit_norm"
            reasons.append("explicit_norm_reference")
        elif COMPARISON_RE.search(q):
            intent = "comparison"
            reasons.append("comparison_pattern")
        elif ENUMERATION_RE.search(q):
            intent = "enumeration"
            reasons.append("enumeration_pattern")
        else:
            intent = "legal" if classify_question(q) == "legal" else classify_question(q)
            reasons.append("classified_by_llm_client_rules")

        # Conservative clarification detection: only ask when no explicit norm is
        # present and the wording is likely underspecified.
        needs_clarification = False
        clarification_question: str | None = None
        if self.options.enable_clarification and not refs and CLARIFICATION_RISK_RE.search(lower):
            if len(expanded) == 0 and len(q.split()) <= 12:
                needs_clarification = True
                clarification_question = (
                    "Meinst du eine konkrete Regelung im Rahmenvertrag, eine bestimmte Anlage "
                    "oder die rechtliche Folge für einen bestimmten Sachverhalt?"
                )
                reasons.append("underspecified_without_norm_reference")

        return QuestionAnalysis(
            question=q,
            question_kind=classify_question(q),
            intent=intent,
            norm_references=refs,
            query_terms=[r.canonical_ref for r in refs],
            expanded_terms=expanded,
            needs_clarification=needs_clarification,
            clarification_question=clarification_question,
            reasons=reasons,
        )

    def _retriever_query(self, question: str, analysis: QuestionAnalysis, options: OrchestratorOptions) -> List[Dict[str, Any]]:
        where = {"container_id": options.default_container_id} if options.restrict_to_default_container else None
        explicit_sections = [ref.section_id for ref in analysis.norm_references] or None

        kwargs: Dict[str, Any] = {
            "question": question,
            "top_k": options.top_k,
            "where": where,
            "fetch_k": options.fetch_k,
            "include_explicit_sections": options.include_explicit_sections,
            "explicit_sections": explicit_sections,
            "include_neighbors": options.include_neighbors,
            "max_final_results": options.max_final_results,
            "restrict_to_default_container": options.restrict_to_default_container,
        }
        if options.min_score is not None:
            kwargs["min_score"] = options.min_score

        # Newer standalone retriever may accept these flags; older one will not.
        if analysis.intent == "definition":
            kwargs["enable_definition_lookup"] = True
        kwargs["verify_negative_answer"] = False

        try:
            return list(self.retriever.query(**kwargs) or [])
        except TypeError:
            # Backward-compatible fallback for older retrievers.
            kwargs.pop("enable_definition_lookup", None)
            kwargs.pop("verify_negative_answer", None)
            try:
                return list(self.retriever.query(**kwargs) or [])
            except TypeError:
                try:
                    return list(self.retriever.query(question=question, top_k=options.top_k, where=where) or [])
                except TypeError:
                    return list(self.retriever.query(question, top_k=options.top_k) or [])

    def assess_retrieval(self, hits: Sequence[Dict[str, Any]], analysis: QuestionAnalysis) -> RetrievalAssessment:
        kinds = [kind for hit in hits for kind in _hit_kinds(hit)]
        containers = list(dict.fromkeys(str(_hit_value(h, "container", "container_id", default="")) for h in hits if _hit_value(h, "container", "container_id", default="")))
        sections = list(dict.fromkeys(str(_hit_value(h, "section", "section_id", default="")) for h in hits if _hit_value(h, "section", "section_id", default="")))
        canonical_refs = list(dict.fromkeys(ref for ref in (_canonical_ref_from_hit(h) for h in hits) if ref))

        direct_hit_count = sum(1 for k in kinds if k not in {"neighbor", "parent_context"})
        exact_count = sum(1 for k in kinds if k in {"exact_norm", "explicit_section", "section_lookup"})
        definition_count = sum(1 for h in hits if bool(_hit_value(h, "is_definition", default=False)) or "definition" in set(_hit_kinds(h)))
        parent_count = sum(1 for k in kinds if k == "parent_context")
        neighbor_count = sum(1 for k in kinds if k == "neighbor")

        low_confidence = False
        reasons: List[str] = []
        if not hits:
            low_confidence = True
            reasons.append("no_hits")
        if analysis.norm_references and exact_count == 0:
            low_confidence = True
            reasons.append("explicit_norm_without_exact_hit")
        if analysis.intent == "definition" and definition_count == 0:
            low_confidence = True
            reasons.append("definition_question_without_definition_hit")
        if direct_hit_count == 0 and hits:
            low_confidence = True
            reasons.append("only_context_hits")

        return RetrievalAssessment(
            hit_count=len(hits),
            direct_hit_count=direct_hit_count,
            exact_norm_hit_count=exact_count,
            definition_hit_count=definition_count,
            parent_context_count=parent_count,
            neighbor_count=neighbor_count,
            containers=containers,
            sections=sections,
            canonical_refs=canonical_refs[:20],
            low_confidence=low_confidence,
            reasons=reasons,
        )

    def _negative_recheck(
        self,
        question: str,
        analysis: QuestionAnalysis,
        hits: List[Dict[str, Any]],
        assessment: RetrievalAssessment,
        options: OrchestratorOptions,
    ) -> Tuple[List[Dict[str, Any]], RetrievalAssessment]:
        if not options.enable_negative_recheck or not assessment.low_confidence:
            return hits, assessment

        before = len(hits)
        recheck_hits: List[Dict[str, Any]] = []

        if hasattr(self.retriever, "verify_negative_result"):
            try:
                check = self.retriever.verify_negative_result(question)
                if isinstance(check, dict):
                    recheck_hits = list(check.get("results") or check.get("hits") or [])
                elif isinstance(check, list):
                    recheck_hits = list(check)
            except TypeError:
                try:
                    check = self.retriever.verify_negative_result(question=question)
                    if isinstance(check, dict):
                        recheck_hits = list(check.get("results") or check.get("hits") or [])
                except Exception:
                    recheck_hits = []
            except Exception:
                recheck_hits = []

        if not recheck_hits:
            # Manual broad fallback: all containers, no restrictive score if possible.
            broad_options = OrchestratorOptions(**asdict(options))
            broad_options.restrict_to_default_container = False
            broad_options.min_score = None
            broad_options.include_neighbors = True
            broad_options.include_explicit_sections = True
            broad_options.fetch_k = max(options.fetch_k, options.top_k * 6)
            broad_options.max_final_results = max(options.max_final_results, 16)
            recheck_hits = self._retriever_query(question, analysis, broad_options)

        combined = _dedupe_hits([*hits, *recheck_hits])
        assessment = self.assess_retrieval(combined, analysis)
        assessment.negative_recheck_performed = True
        assessment.negative_recheck_added_hits = max(0, len(combined) - before)
        if assessment.negative_recheck_added_hits:
            assessment.reasons.append("negative_recheck_added_hits")
        return combined, assessment

    # ------------------------------------------------------------------
    # Composition and audit
    # ------------------------------------------------------------------

    def _compose(self, question: str, hits: List[Dict[str, Any]], memory: Optional[ConversationMemory]) -> Tuple[str, str, List[Dict[str, Any]]]:
        if hasattr(self.composer, "compose_with_sources"):
            try:
                return self.composer.compose_with_sources(question, hits, memory=memory)
            except TypeError:
                pass
        answer, answer_type = self.composer.compose(question, hits, memory=memory)
        sources: List[Dict[str, Any]] = []
        if hasattr(self.composer, "build_sources"):
            try:
                sources = self.composer.build_sources(hits)
            except Exception:
                sources = []
        return answer, answer_type, sources

    @staticmethod
    def _source_numbers_from_source(source: Dict[str, Any], *, fallback: int | None = None) -> List[int]:
        """Return all source numbers represented by a Composer/API source item."""
        numbers: List[int] = []

        raw_numbers = source.get("source_numbers")
        if raw_numbers is not None:
            if not isinstance(raw_numbers, (list, tuple, set)):
                raw_numbers = [raw_numbers]
            for value in raw_numbers:
                try:
                    number = int(value)
                except (TypeError, ValueError):
                    continue
                if number not in numbers:
                    numbers.append(number)

        raw_single = source.get("source_number")
        if raw_single is not None:
            try:
                number = int(raw_single)
                if number not in numbers:
                    numbers.append(number)
            except (TypeError, ValueError):
                pass

        if not numbers and fallback is not None:
            numbers.append(fallback)

        return sorted(numbers)

    @staticmethod
    def _canonical_refs_from_source(source: Dict[str, Any]) -> List[str]:
        refs: List[str] = []

        raw_refs = source.get("canonical_refs")
        if raw_refs is not None:
            if not isinstance(raw_refs, (list, tuple, set)):
                raw_refs = [raw_refs]
            for value in raw_refs:
                ref = _normalize(str(value or ""))
                if ref and ref not in refs:
                    refs.append(ref)

        direct = _normalize(str(source.get("canonical_ref") or ""))
        if direct and direct not in refs:
            refs.append(direct)

        fallback = _normalize(_canonical_ref_from_hit(source))
        if fallback and fallback not in refs:
            refs.append(fallback)

        return refs

    @classmethod
    def _source_number_to_ref(cls, sources: Sequence[Dict[str, Any]], hits: Sequence[Dict[str, Any]]) -> Dict[int, str]:
        """Map every visible [Quelle n] number to the best known canonical ref.

        Composer sources may group multiple source numbers into one displayed
        source item.  The previous implementation only considered
        ``source_number`` and silently lost grouped numbers such as
        ``source_numbers=[1, 3]``.
        """
        mapping: Dict[int, str] = {}

        for idx, source in enumerate(sources, start=1):
            numbers = cls._source_numbers_from_source(source, fallback=idx)
            refs = cls._canonical_refs_from_source(source)
            if not refs:
                continue

            # Prefer a concise, stable ref for inline enrichment.  Fine-grained
            # refs remain available in the separate source list; inline enrichment
            # must not invent suspicious Buchst.-level citations.
            ref = refs[0]
            for candidate in refs:
                if "Buchst." not in candidate and "Buchstabe" not in candidate:
                    ref = candidate
                    break

            for number in numbers:
                mapping.setdefault(number, ref)

        # Fallback: use retrieval hit order only when Composer sources did not
        # provide any number mapping.  This keeps older stacks usable while not
        # overriding Composer numbering.
        if not mapping:
            for idx, hit in enumerate(hits, start=1):
                ref = _canonical_ref_from_hit(hit)
                if ref:
                    mapping[idx] = ref

        return mapping

    @staticmethod
    def _answer_has_inline_ref_after_marker(answer: str, marker_end: int) -> bool:
        return bool(re.match(r"\s*\(", (answer or "")[marker_end : marker_end + 8]))

    @classmethod
    def _enrich_answer_citations(cls, answer: str, source_to_ref: Dict[int, str]) -> str:
        """Conservatively add canonical refs to bare source markers.

        Inline enrichment is intentionally restrained.  Generic section-only
        references such as ``§ 6`` add little value and caused noisy outputs
        like ``[Quelle 1] (§ 6) und (§ 6)``.  More specific refs may still be
        added when the marker is bare and the model did not already provide a
        parenthetical.
        """
        if not answer or not source_to_ref:
            return answer

        def repl(match: re.Match[str]) -> str:
            source_no = int(match.group(1))
            ref = source_to_ref.get(source_no)
            marker = match.group(0)
            if not ref:
                return marker
            if cls._answer_has_inline_ref_after_marker(answer, match.end()):
                return marker
            if not re.search(r"\b(?:Abs\.|Satz|Nr\.|Buchst\.)\b", ref, flags=re.I):
                return marker
            return f"{marker} ({ref})"

        return SOURCE_MARKER_RE.sub(repl, answer)

    @staticmethod
    def _extract_source_numbers(answer: str) -> List[int]:
        out: List[int] = []
        for match in SOURCE_MARKER_RE.finditer(answer or ""):
            try:
                out.append(int(match.group(1)))
            except ValueError:
                continue
        return list(dict.fromkeys(out))

    @staticmethod
    def _normalize_legal_ref(ref: str) -> str:
        s = _normalize(ref)
        s = re.sub(r"§\s*", "§ ", s)
        s = re.sub(r"\bAbsatz\b", "Abs.", s, flags=re.I)
        s = re.sub(r"\bBuchstabe\b|\blit\.", "Buchst.", s, flags=re.I)
        s = re.sub(r"\s+", " ", s).strip(" .,;:")
        return s.lower()

    @classmethod
    def _supported_ref_set(cls, sources: Sequence[Dict[str, Any]], hits: Sequence[Dict[str, Any]]) -> set[str]:
        refs: set[str] = set()
        for source in sources:
            for ref in cls._canonical_refs_from_source(source):
                refs.add(cls._normalize_legal_ref(ref))
        for hit in hits:
            ref = _canonical_ref_from_hit(hit)
            if ref:
                refs.add(cls._normalize_legal_ref(ref))
        return refs

    @classmethod
    def _downgrade_unsupported_fine_refs(
        cls,
        answer: str,
        sources: Sequence[Dict[str, Any]],
        hits: Sequence[Dict[str, Any]],
    ) -> Tuple[str, List[str]]:
        """Downgrade unsupported Buchst.-level references to Absatz level."""
        fixes: List[str] = []
        if not answer:
            return answer, fixes

        supported = cls._supported_ref_set(sources, hits)

        def repl(match: re.Match[str]) -> str:
            full = match.group(0)
            norm_full = cls._normalize_legal_ref(full)
            if norm_full in supported:
                return full
            downgraded = f"§ {match.group('section')} Abs. {match.group('subsection')}"
            fixes.append(f"downgraded_unsupported_fine_ref:{full}->{downgraded}")
            return downgraded

        cleaned = FINE_BUCHST_REF_RE.sub(repl, answer)

        # Collapse duplicates introduced by downgrading, e.g.
        # "[Quelle 1] (§ 6 Abs. 1), § 6 Abs. 1".
        cleaned = re.sub(
            r"(\[Quelle\s+\d+(?:[^\]]*)\]\s*\((§\s*\d{1,3}[a-z]?\s+Abs\.\s*\d{1,3}[a-z]?)\))\s*,\s*\2",
            r"\1",
            cleaned,
            flags=re.I,
        )
        cleaned = re.sub(
            r"(\((§\s*\d{1,3}[a-z]?\s+Abs\.\s*\d{1,3}[a-z]?)\))\s*,\s*\2",
            r"\1",
            cleaned,
            flags=re.I,
        )
        return cleaned, fixes

    @classmethod
    def _valid_source_numbers(cls, sources: Sequence[Dict[str, Any]], hits: Sequence[Dict[str, Any]]) -> set[int]:
        valid: set[int] = set()
        for idx, source in enumerate(sources, start=1):
            valid.update(cls._source_numbers_from_source(source, fallback=idx))
        if not valid:
            valid.update(range(1, len(hits) + 1))
        return valid

    @classmethod
    def _strip_invalid_source_markers_safely(
        cls,
        answer: str,
        sources: Sequence[Dict[str, Any]],
        hits: Sequence[Dict[str, Any]],
    ) -> Tuple[str, List[str]]:
        """Remove invalid source markers including attached parenthetical refs.

        Previous versions removed only ``[Quelle n]`` and left the enrichment
        tail behind.  That produced broken fragments like ``und (§ 6)``.  This
        method removes the whole citation atom for invalid markers and then runs
        a grammar cleanup pass.
        """
        fixes: List[str] = []
        if not answer:
            return answer, fixes

        valid = cls._valid_source_numbers(sources, hits)

        def repl(match: re.Match[str]) -> str:
            try:
                number = int(match.group(1))
            except (TypeError, ValueError):
                fixes.append("removed_malformed_source_marker")
                return ""
            if number in valid:
                return match.group(0)
            fixes.append(f"removed_invalid_source_marker_with_ref:{number}")
            return ""

        cleaned = SOURCE_MARKER_WITH_OPTIONAL_REF_RE.sub(repl, answer)
        cleaned = cls._cleanup_citation_grammar(cleaned)
        return cleaned.strip(), fixes

    @staticmethod
    def _cleanup_citation_grammar(answer: str) -> str:
        """Repair grammar after citation cleanup without changing legal content."""
        cleaned = answer or ""

        # Collapse model/enrichment duplicates such as:
        #   [Quelle 1] (§ 6), § 6 Abs. 1 -> [Quelle 1] (§ 6 Abs. 1)
        cleaned = re.sub(
            r"(\[Quelle\s+\d+(?:[^\]]*)\])\s*\(§\s*(\d{1,3}[a-z]?)\)\s*,\s*(§\s*\2\s+Abs\.\s*\d{1,3}[a-z]?)",
            r"\1 (\3)",
            cleaned,
            flags=re.I,
        )

        for pattern, replacement in DANGLING_CITATION_GRAMMAR_REPLACEMENTS:
            cleaned = pattern.sub(replacement, cleaned)

        # Remove duplicate adjacent identical parenthetical legal refs.
        cleaned = re.sub(
            r"(\(§[^)]{1,120}\))\s*(?:,|und)\s*\1",
            r"\1",
            cleaned,
            flags=re.I,
        )

        # Clean common connective leftovers.
        cleaned = re.sub(r"\s+und\s+([,.;:])", r"\1", cleaned, flags=re.I)
        cleaned = re.sub(r"\s+und\s*(?=\.)", "", cleaned, flags=re.I)
        cleaned = re.sub(r"\(\s*\)", "", cleaned)
        cleaned = re.sub(r"\s+([,.;:])", r"\1", cleaned)
        cleaned = re.sub(r"[ \t]{2,}", " ", cleaned)
        cleaned = re.sub(r"\n[ \t]+", "\n", cleaned)
        return cleaned.strip()

    @classmethod
    def _ensure_cited_sources_visible(
        cls,
        answer: str,
        sources: List[Dict[str, Any]],
        hits: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """Ensure every cited [Quelle n] is represented in returned sources."""
        cited = cls._extract_source_numbers(answer)
        if not cited:
            return sources

        visible = cls._valid_source_numbers(sources, [])
        missing = [n for n in cited if n not in visible]
        if not missing:
            return sources

        out = list(sources)
        for number in missing:
            idx = number - 1
            if 0 <= idx < len(hits):
                source = dict(hits[idx])
                source["source_number"] = number
                source["source_numbers"] = [number]
                if "canonical_ref" not in source or not source.get("canonical_ref"):
                    source["canonical_ref"] = _canonical_ref_from_hit(source)
                out.append(source)
        return out

    @classmethod
    def _filter_sources_to_cited_sources(
        cls,
        answer: str,
        sources: List[Dict[str, Any]],
        hits: List[Dict[str, Any]],
        *,
        max_sources: int,
    ) -> List[Dict[str, Any]]:
        """Return only sources that are actually cited, with hit fallback.

        For answer quality, visible sources should not include unrelated
        retrieval leftovers when the answer cites only a small subset.  If the
        model cites no source markers, keep the curated source list.
        """
        cited = cls._extract_source_numbers(answer)
        if not cited:
            return list(sources)[:max_sources]

        by_number: Dict[int, Dict[str, Any]] = {}
        for idx, source in enumerate(sources, start=1):
            for number in cls._source_numbers_from_source(source, fallback=idx):
                by_number.setdefault(number, source)

        out: List[Dict[str, Any]] = []
        seen_ids: set[int] = set()
        for number in cited:
            source = by_number.get(number)
            if source is None:
                hit_idx = number - 1
                if 0 <= hit_idx < len(hits):
                    source = dict(hits[hit_idx])
                    source["source_number"] = number
                    source["source_numbers"] = [number]
                    source.setdefault("canonical_ref", _canonical_ref_from_hit(source))
            if source is None:
                continue
            ident = id(source)
            if ident in seen_ids:
                continue
            seen_ids.add(ident)
            out.append(source)
            if len(out) >= max_sources:
                break

        return out or list(sources)[:max_sources]

    @classmethod
    def _postprocess_answer(
        cls,
        answer: str,
        sources: List[Dict[str, Any]],
        hits: List[Dict[str, Any]],
    ) -> Tuple[str, List[str]]:
        fixes: List[str] = []
        cleaned, ref_fixes = cls._downgrade_unsupported_fine_refs(answer, sources, hits)
        fixes.extend(ref_fixes)
        cleaned, marker_fixes = cls._strip_invalid_source_markers_safely(cleaned, sources, hits)
        fixes.extend(marker_fixes)
        cleaned = cls._cleanup_citation_grammar(cleaned)
        return cleaned, fixes

    @staticmethod
    def _detect_answer_basis(answer: str) -> str:
        a = _lower(answer)
        if NEGATIVE_ANSWER_RE.search(answer or ""):
            return "negative"
        if "ausdrücklich geregelt" in a:
            return "explicit"
        if "systematisch" in a or "ableitbar" in a or "auslegung" in a:
            return "derived"
        if "nicht belastbar" in a or "keine belastbare" in a:
            return "insufficient"
        return "unknown"

    @staticmethod
    def _expected_list_labels_from_hits(hits: Sequence[Dict[str, Any]]) -> List[str]:
        text = "\n".join(_hit_text(h) for h in hits[:8])
        labels = set()
        for m in re.finditer(r"(?:^|\n|\s)([a-z])\)\s+", text):
            labels.add(f"{m.group(1).lower()})")
        for m in re.finditer(r"(?:^|\n|\s)(\d{1,2})[.)]\s+", text):
            labels.add(f"{m.group(1)}")
        ordered_letters = [f"{chr(i)})" for i in range(ord("a"), ord("z") + 1) if f"{chr(i)})" in labels]
        ordered_nums = sorted([x for x in labels if x.isdigit()], key=lambda n: int(n))
        return ordered_letters + ordered_nums

    def audit_answer(
        self,
        question: str,
        answer: str,
        hits: List[Dict[str, Any]],
        sources: List[Dict[str, Any]],
        analysis: QuestionAnalysis,
        assessment: RetrievalAssessment,
    ) -> AnswerAudit:
        cited = self._extract_source_numbers(answer)
        valid_numbers = self._valid_source_numbers(sources, hits)

        invalid = [n for n in cited if n not in valid_numbers]
        basis = self._detect_answer_basis(answer)
        negative = basis == "negative"
        completeness_warnings: List[str] = []
        citation_warnings: List[str] = []

        if not cited and hits and basis != "negative":
            citation_warnings.append("answer_contains_no_source_markers")
        if invalid:
            citation_warnings.append("answer_contains_invalid_source_markers")
        if analysis.norm_references and not any(ref.section_id in " ".join(assessment.sections + assessment.canonical_refs) for ref in analysis.norm_references):
            completeness_warnings.append("explicit_norm_not_reflected_in_retrieved_context")
        if analysis.intent in {"enumeration", "norm_enumeration"}:
            labels = self._expected_list_labels_from_hits(hits)
            if len(labels) >= 3:
                answer_lower = answer.lower()
                missing_labels = [label for label in labels if label not in answer_lower]
                # Do not force exact labels in prose answers, but flag likely omissions.
                if len(missing_labels) >= max(2, len(labels) // 2):
                    completeness_warnings.append(
                        "possible_incomplete_enumeration: expected list markers " + ", ".join(labels[:12])
                    )
        if negative and assessment.hit_count > 0 and not assessment.low_confidence:
            completeness_warnings.append("negative_answer_despite_available_context")

        confidence = 0.55
        if assessment.hit_count:
            confidence += 0.15
        if assessment.exact_norm_hit_count and analysis.norm_references:
            confidence += 0.15
        if assessment.definition_hit_count and analysis.intent == "definition":
            confidence += 0.15
        if citation_warnings:
            confidence -= 0.15
        if completeness_warnings:
            confidence -= 0.15
        if negative and assessment.low_confidence:
            confidence -= 0.10
        confidence = round(max(0.0, min(confidence, 0.98)), 2)

        recommended_action = "accept"
        if analysis.needs_clarification:
            recommended_action = "clarify"
        elif negative and assessment.low_confidence:
            recommended_action = "recheck"
        elif citation_warnings or completeness_warnings:
            recommended_action = "caution"

        refs_by_no = self._source_number_to_ref(sources, hits)
        cited_refs = [refs_by_no[n] for n in cited if n in refs_by_no]

        return AnswerAudit(
            cited_source_numbers=cited,
            invalid_source_numbers=invalid,
            cited_canonical_refs=cited_refs,
            answer_basis=basis,
            negative_answer_detected=negative,
            completeness_warnings=completeness_warnings,
            citation_warnings=citation_warnings,
            recommended_action=recommended_action,
            confidence=confidence,
        )

    def _maybe_recompose_after_negative_answer(
        self,
        question: str,
        answer: str,
        hits: List[Dict[str, Any]],
        sources: List[Dict[str, Any]],
        analysis: QuestionAnalysis,
        assessment: RetrievalAssessment,
        memory: Optional[ConversationMemory],
        options: OrchestratorOptions,
    ) -> Tuple[str, str, List[Dict[str, Any]], List[Dict[str, Any]], RetrievalAssessment]:
        if not options.enable_negative_recheck or not NEGATIVE_ANSWER_RE.search(answer or ""):
            return answer, "document", sources, hits, assessment

        old_count = len(hits)
        hits2, assessment2 = self._negative_recheck(question, analysis, hits, assessment, options)
        if len(hits2) <= old_count:
            return answer, "document", sources, hits, assessment2

        answer2, answer_type2, sources2 = self._compose(question, hits2, memory)
        return answer2, answer_type2, sources2, hits2, assessment2

    # ------------------------------------------------------------------
    # Public flow
    # ------------------------------------------------------------------

    def run(
        self,
        question: str,
        *,
        memory: Optional[ConversationMemory] = None,
        options: OrchestratorOptions | None = None,
    ) -> OrchestratorResult:
        options = options or self.options
        question = _normalize(question)
        analysis = self.analyze_question(question)

        if analysis.intent == "meta":
            answer, answer_type = self.composer.compose(question, [], memory=memory)
            return OrchestratorResult(
                answer=answer,
                answer_type=answer_type,
                analysis=analysis,
                retrieval_assessment=RetrievalAssessment(),
                answer_audit=AnswerAudit(answer_basis="meta", confidence=1.0),
            )

        if analysis.needs_clarification:
            answer = analysis.clarification_question or "Bitte präzisiere deine Frage."
            return OrchestratorResult(
                answer=answer,
                answer_type="clarification",
                analysis=analysis,
                retrieval_assessment=RetrievalAssessment(low_confidence=True, reasons=["clarification_required"]),
                answer_audit=AnswerAudit(answer_basis="clarification", recommended_action="clarify", confidence=0.4),
                needs_clarification=True,
                clarification_question=answer,
            )

        hits = self._retriever_query(question, analysis, options)
        hits = _dedupe_hits(hits)
        assessment = self.assess_retrieval(hits, analysis)
        hits, assessment = self._negative_recheck(question, analysis, hits, assessment, options)

        if not hits:
            answer = (
                "Ich habe im verfügbaren Vertragskorpus keine belastbare Textstelle gefunden. "
                "Das bedeutet nicht zwingend, dass der Sachverhalt rechtlich nicht geregelt ist; "
                "es heißt zunächst nur, dass die relevante Regelung im aktuellen Retrieval-Kontext "
                "nicht auffindbar war."
            )
            audit = AnswerAudit(
                answer_basis="insufficient",
                negative_answer_detected=True,
                recommended_action="caution",
                confidence=0.25,
                completeness_warnings=["no_retrieval_hits_after_recheck"],
            )
            return OrchestratorResult(
                answer=answer,
                answer_type="none",
                hits=[],
                raw_sources=[],
                sources=[],
                analysis=analysis,
                retrieval_assessment=assessment,
                answer_audit=audit,
            )

        answer, answer_type, sources = self._compose(question, hits, memory)
        answer, answer_type, sources, hits, assessment = self._maybe_recompose_after_negative_answer(
            question, answer, hits, sources, analysis, assessment, memory, options
        )

        # Keep source visibility and answer postprocessing inside the same
        # numbering universe as the Composer.  This prevents the API layer from
        # deleting markers such as [Quelle 3] while the answer still contains
        # dangling grammar fragments.
        sources = self._ensure_cited_sources_visible(answer, sources, hits)

        refs_by_no = self._source_number_to_ref(sources, hits)
        if options.enrich_citations_with_canonical_refs:
            answer = self._enrich_answer_citations(answer, refs_by_no)

        sources = self._ensure_cited_sources_visible(answer, sources, hits)
        answer, auto_fixes = self._postprocess_answer(answer, sources, hits)
        sources = self._ensure_cited_sources_visible(answer, sources, hits)
        sources = self._filter_sources_to_cited_sources(
            answer,
            sources,
            hits,
            max_sources=options.max_sources,
        )

        audit = self.audit_answer(question, answer, hits, sources, analysis, assessment) if options.enable_answer_audit else AnswerAudit()
        audit.auto_fixes.extend(auto_fixes)

        debug: Dict[str, Any] = {}
        if options.debug:
            debug = {
                "source_number_to_canonical_ref": refs_by_no,
                "retrieved_hit_count": len(hits),
                "source_count": len(sources),
                "answer_auto_fixes": auto_fixes,
                "returned_source_numbers": [
                    self._source_numbers_from_source(source, fallback=idx)
                    for idx, source in enumerate(sources, start=1)
                ],
                "cited_source_numbers_after_postprocess": self._extract_source_numbers(answer),
            }

        return OrchestratorResult(
            answer=answer,
            answer_type=answer_type,
            hits=hits,
            raw_sources=sources,
            sources=sources,
            analysis=analysis,
            retrieval_assessment=assessment,
            answer_audit=audit,
            needs_clarification=False,
            debug=debug,
        )