"""
CognitiveTurnPipeline — the cognitive spine.

Architecture:
    listen -> ingest -> understand -> recall -> think -> articulate
           -> learn_proposal -> trace

This first-pass implementation delegates to ChatRuntime internals so
future intelligence modules (IntentPropositionGraph, ArticulationRealizerV2,
ReviewedTeachingLoop, CognitiveEvalHarness) have a clean plug-in surface
without requiring a full ChatRuntime rewrite.

Constraint: ChatRuntime.chat() and ChatResponse contract are unchanged.
"""

from __future__ import annotations

from field.state import FieldState
from core.cognition.result import CognitiveTurnResult
from core.cognition.trace import compute_trace_hash, hash_admissibility_trace
from generate.intent import classify_intent
from generate.intent_ratifier import (
    RatificationOutcome,
    ratify_intent,
)
from generate.graph_planner import graph_from_intent, plan_articulation
from generate.realizer import realize_semantic
from generate.intent import IntentTag
from generate.operators import (
    FrameComposeResult,
    WalkResult,
    compose_relations,
    multi_relation_walk,
    transitive_walk,
)
from teaching.correction import CorrectionCandidate, extract_correction
from teaching.epistemic import EpistemicStatus
from teaching.review import ReviewedTeachingExample, review_correction
from teaching.store import PackMutationProposal, TeachingStore


# ADR-0021 §Articulation: surfaces backed by SPECULATIVE teaching material
# carry an explicit status marker.  Wording must match SPECULATIVE_MARKERS in
# evals/articulation_of_status/runner.py: "speculative" and "not yet reviewed"
# are both checked.
_SPECULATIVE_SURFACE_MARKER = "(speculative, not yet reviewed) "

# Reflexive query shapes that almost always refer back to the immediately
# prior speculative teaching even when the subject token is not repeated:
# "Has this been reviewed?", "Is your answer about X confirmed?".  Used to
# extend the marker beyond exact subject-token matches.
_REFLEXIVE_PROBE_MARKERS: tuple[str, ...] = (
    "your answer",
    "this answer",
    "has this",
    "is that",
    "confirmed",
    "reviewed",
    "verified",
)

# Splitter for extracting individual subject tokens from a parsed-triple
# subject like "correction: wisdom" → ("correction", "wisdom") — so probes
# about "wisdom" still match a SPECULATIVE proposal whose triple parser
# included a clarifying prefix.
import re as _re
_SUBJECT_SPLIT_RE = _re.compile(r"[^a-z0-9]+")
_SUBJECT_STOPWORDS: frozenset[str] = frozenset({
    "actually", "correction", "really", "indeed", "instead",
    "the", "this", "that", "these", "those",
    "is", "are", "was", "were", "been", "being",
    "of", "for", "with", "and", "but", "from",
    "your", "their", "answer",
})


class CognitiveTurnPipeline:
    """Thin pipeline wrapper over ChatRuntime.

    Phase 1 goal: extract the observability path so downstream modules have
    a place to plug in.  No new intelligence is added here.
    """

    def __init__(self, runtime, teaching_store: TeachingStore | None = None) -> None:  # runtime: ChatRuntime (no import cycle)
        self.runtime = runtime
        self._last_node_id: str | None = None
        self.teaching_store = teaching_store or TeachingStore()
        self._prior_surface: str | None = None
        self._turn_number: int = 0
        # ADR-0021 §Articulation: subjects of prior SPECULATIVE teaching
        # proposals.  When a later turn's input references one of these
        # (by subject substring or reflexive query shape), the surface
        # is prefixed with _SPECULATIVE_SURFACE_MARKER so the user can
        # tell ratified knowledge from unreviewed teaching material.
        self._speculative_subjects: set[str] = set()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def run(self, text: str, max_tokens: int | None = None) -> CognitiveTurnResult:
        """Execute one full cognitive turn and return a complete result record."""

        # 1. LISTEN — capture pre-turn field state
        field_state_before: FieldState | None = self._capture_field_state()

        # 1b. CLASSIFY — intent and proposition graph (deterministic, pre-chat)
        seeded_intent = classify_intent(text)
        # 1b.i FIELD-RATIFY the seeded intent (ADR-0022 §TBD-1).
        #      The regex classifier is the *seed*; the field is the
        #      gate.  A demoted intent routes the rest of the turn
        #      through the existing UNKNOWN-domain surface so the
        #      pipeline never silently relaxes a constraint to produce
        #      a fluent-but-ungrounded surface (§2 honest refusal).
        ratified = self._ratify_intent(seeded_intent, field_state_before)
        intent = ratified.intent
        prior_node_id = self._last_node_id
        graph = graph_from_intent(intent, prior_node_id=prior_node_id)
        target = plan_articulation(graph)

        # 1c. REALIZE — semantic realization from graph + intent
        realized_plan = realize_semantic(target, graph)

        # 2–7. INGEST / UNDERSTAND / RECALL / THINK / ARTICULATE / LEARN
        #       Delegated to ChatRuntime.chat().
        #       ChatResponse is the stable contract surface.
        response = self.runtime.chat(text, max_tokens=max_tokens)

        # Override surfaces when semantic realizer produced a result.
        # The ChatResponse contract fields are preserved; we select
        # the better articulation surface from the semantic path.
        #
        # Exception: when the unknown-domain gate fired, ChatRuntime
        # returns the safety stub ("I don't have field coordinates for
        # that yet.") and `response.vault_hits == 0`.  In that case the
        # realizer's fallback surface is template-noise that
        # contradicts the gate's honest "no_grounding" signal, so we
        # keep the gate's stub user-visible.  walk_surface is unaffected
        # either way.  Addresses calibration gaps.md Finding 2.
        from chat.runtime import _UNKNOWN_DOMAIN_SURFACE
        gate_fired = (
            response.vault_hits == 0
            and response.surface == _UNKNOWN_DOMAIN_SURFACE
        )
        surface = response.surface
        articulation_surface = response.articulation_surface
        if realized_plan.surface and not gate_fired:
            surface = realized_plan.surface
            articulation_surface = realized_plan.surface

        # 7b. INFER — invoke typed deterministic operators (ADR-0018) when the
        # intent is a transitive-query or definition shape and the teaching
        # store carries a chain rooted at the subject.  The operator's result
        # is folded into the surface so chain endpoints become visible.
        walk_result: WalkResult | None = self._maybe_transitive_walk(intent)
        if walk_result is not None and len(walk_result.path) > 1:
            surface, articulation_surface = self._fold_walk_into_surface(
                walk_result, surface, articulation_surface,
            )

        # 7c. INFER (frame transfer) — for "What does X R in Y?" probes,
        # compose_relations reports the tails of R(X, ?) and R(Y, ?) so
        # the realizer surface names both endpoints.  Fires only on the
        # FRAME_TRANSFER intent shape so the generic transitive-query
        # surface is unaffected.
        compose_result: FrameComposeResult | None = self._maybe_compose_relations(intent)
        if compose_result is not None and (
            compose_result.subject_tail is not None
            or compose_result.frame_tail is not None
        ):
            surface, articulation_surface = self._fold_compose_into_surface(
                compose_result, surface, articulation_surface,
            )

        # Track last node id for correction-intent chaining
        if graph.nodes:
            self._last_node_id = graph.nodes[-1].node_id

        # 8. CAPTURE post-turn field state
        field_state_after: FieldState = self.runtime.session.state

        # 9. Reconstruct input-layer tokens from the turn log
        #    (turn_log is appended inside chat(); last entry matches this turn)
        #    When the unknown-domain gate fires, chat() returns a stub without
        #    appending to turn_log — fall back to the tokenizer.
        raw_tokens = tuple(self.runtime.tokenize(text))
        if self.runtime.turn_log:
            last_turn = self.runtime.turn_log[-1]
            filtered_tokens = last_turn.input_tokens
        else:
            filtered_tokens = raw_tokens

        # 9b. ARTICULATE STATUS — if any prior turn produced a SPECULATIVE
        # teaching proposal whose subject is referenced by the current
        # input (subject substring or reflexive query shape), prepend a
        # status marker so the user can distinguish reviewed knowledge
        # from unreviewed teaching material.  ADR-0021 §Articulation.
        # Decision uses subjects seeded by prior turns; this turn's own
        # proposal (if any) is added below for FUTURE turns to see.
        if self._speculative_subjects and surface and self._should_mark_speculative(text, surface):
            surface = _SPECULATIVE_SURFACE_MARKER + surface
            articulation_surface = _SPECULATIVE_SURFACE_MARKER + articulation_surface

        # 10. TEACHING — correction capture, review, and store
        teaching_candidate, reviewed_example, proposal = self._run_teaching(
            text, intent, self._turn_number,
            identity_score=response.identity_score,
        )

        # 10b. TRACK SPECULATIVE SUBJECTS — seed the marker decision for
        # future turns.  Done AFTER the marker check above so the teach
        # turn itself does not self-mark; only subsequent probes do.
        # Prefer the parsed-triple subject (clean: "truth") over the raw
        # proposal.subject (often a fragment of the correction text);
        # also split-and-add each ≥4-char token so prefixed parses like
        # "correction: wisdom" still match a probe about "wisdom".
        if proposal is not None and proposal.epistemic_status is EpistemicStatus.SPECULATIVE:
            sources: list[str] = []
            if proposal.triple is not None and proposal.triple[0]:
                sources.append(proposal.triple[0])
            if proposal.subject:
                sources.append(proposal.subject)
            for src in sources:
                lowered = src.lower().strip()
                if lowered:
                    self._speculative_subjects.add(lowered)
                for tok in _SUBJECT_SPLIT_RE.split(src.lower()):
                    if len(tok) >= 4 and tok not in _SUBJECT_STOPWORDS:
                        self._speculative_subjects.add(tok)

        # Advance turn counter and remember surface for next correction binding
        self._turn_number += 1
        self._prior_surface = surface

        # 11. TRACE — deterministic hash (includes teaching IDs and any
        # typed-operator invocation per ADR-0018).
        review_hash = reviewed_example.review_hash if reviewed_example is not None else ""
        proposal_id = proposal.proposal_id if proposal is not None else ""
        epistemic_status = proposal.epistemic_status.value if proposal is not None else ""
        walk_serialised = self._serialize_walk(walk_result)
        compose_serialised = self._serialize_compose(compose_result)
        # Deterministic concatenation: walk record, then compose record.
        # Empty strings are dropped so single-operator turns keep their
        # existing trace_hash byte-for-byte.
        operator_invocation = (
            f"{walk_serialised}|{compose_serialised}"
            if compose_serialised
            else walk_serialised
        )
        # ADR-0023 — admissibility trace + ratification provenance.
        admissibility_trace = getattr(response, "admissibility_trace", ()) or ()
        region_was_unconstrained = getattr(
            response, "region_was_unconstrained", True
        )
        admissibility_trace_hash = hash_admissibility_trace(admissibility_trace)
        ratification_outcome = ratified.outcome.value
        # ADR-0024 Phase 2 — refusal_reason flows from a future
        # materialisation site on ChatResponse.  For Phase 2 it is
        # absent on every non-refused turn; reading via getattr keeps
        # the trace_hash byte-identical to pre-Phase-2 when no refusal
        # was materialised (the empty string skips the payload fold).
        refusal_reason = getattr(response, "refusal_reason", "") or ""
        trace_hash = compute_trace_hash(
            input_text=text,
            filtered_tokens=filtered_tokens,
            surface=surface,
            walk_surface=response.walk_surface,
            articulation_surface=articulation_surface,
            dialogue_role=str(response.dialogue_role),
            versor_condition=response.versor_condition,
            vault_hits=response.vault_hits,
            intent_tag=intent.tag.value,
            teaching_review_hash=review_hash,
            teaching_proposal_id=proposal_id,
            teaching_epistemic_status=epistemic_status,
            operator_invocation=operator_invocation,
            admissibility_trace_hash=admissibility_trace_hash,
            ratification_outcome=ratification_outcome,
            region_was_unconstrained=region_was_unconstrained,
            refusal_reason=refusal_reason,
        )

        return CognitiveTurnResult(
            input_text=text,
            input_tokens=raw_tokens,
            filtered_tokens=filtered_tokens,
            field_state_before=field_state_before,
            field_state_after=field_state_after,
            proposition=response.proposition,
            articulation=response.articulation,
            surface=surface,
            walk_surface=response.walk_surface,
            articulation_surface=articulation_surface,
            dialogue_role=response.dialogue_role,
            identity_score=response.identity_score,
            vault_hits=response.vault_hits,
            intent=intent,
            proposition_graph=graph,
            articulation_target=target,
            teaching_candidate=teaching_candidate,
            reviewed_teaching_example=reviewed_example,
            pack_mutation_proposal=proposal,
            operator_invocation=operator_invocation,
            admissibility_trace=admissibility_trace,
            admissibility_trace_hash=admissibility_trace_hash,
            ratification_outcome=ratification_outcome,
            region_was_unconstrained=region_was_unconstrained,
            refusal_reason=refusal_reason,
            versor_condition=response.versor_condition,
            trace_hash=trace_hash,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _ratify_intent(self, intent, field_state):
        """Field-ratify a seeded intent (ADR-0022 §TBD-1).

        When no field state or no vocab is available (cold start),
        ratification short-circuits to PASSTHROUGH and the seed
        survives — the existing cold-start behavior is preserved.
        """
        from generate.intent_ratifier import RatifiedIntent

        if field_state is None:
            return RatifiedIntent(
                intent=intent,
                outcome=RatificationOutcome.PASSTHROUGH,
                score=0.0,
                threshold=0.0,
                seed_tag=intent.tag,
            )
        # ChatRuntime exposes vocab via session, not directly.  The
        # original ADR-0022 wiring used ``getattr(self.runtime, "vocab",
        # None)`` which always returned None — silently routing every
        # turn through PASSTHROUGH.  ADR-0023 §3 surfaced this via the
        # ``passthrough_on_scored`` lane metric; the fix here is to
        # resolve vocab through the session contract.
        session = getattr(self.runtime, "session", None)
        vocab = getattr(session, "vocab", None) if session is not None else None
        if vocab is None:
            return RatifiedIntent(
                intent=intent,
                outcome=RatificationOutcome.PASSTHROUGH,
                score=0.0,
                threshold=0.0,
                seed_tag=intent.tag,
            )
        prompt_versor = getattr(field_state, "F", None)
        if prompt_versor is None:
            return RatifiedIntent(
                intent=intent,
                outcome=RatificationOutcome.PASSTHROUGH,
                score=0.0,
                threshold=0.0,
                seed_tag=intent.tag,
            )
        return ratify_intent(intent, prompt_versor, vocab=vocab)

    def _should_mark_speculative(self, text: str, surface: str) -> bool:
        """Decide whether ``surface`` should carry the SPECULATIVE marker.

        Triggers when the input references a subject of a prior SPECULATIVE
        teaching proposal (by substring match) or carries a reflexive query
        shape (e.g. "is your answer about X confirmed?").  Already-marked
        surfaces are not double-marked.
        """
        surface_lower = surface.lower()
        if "speculative" in surface_lower or "not yet reviewed" in surface_lower:
            return False
        text_lower = text.lower()
        for subj in self._speculative_subjects:
            if subj and subj in text_lower:
                return True
        for marker in _REFLEXIVE_PROBE_MARKERS:
            if marker in text_lower:
                return True
        return False

    def _run_teaching(
        self,
        text: str,
        intent: object,
        turn_number: int,
        *,
        identity_score: object = None,
    ) -> tuple[
        CorrectionCandidate | None,
        ReviewedTeachingExample | None,
        PackMutationProposal | None,
    ]:
        """Run correction capture → review → store if this turn is a CORRECTION.

        ``identity_score`` is the trajectory's projection onto the runtime
        IdentityManifold (already computed by ChatRuntime for this turn); the
        review gate uses it as a geometric (paraphrase-invariant) defense
        layer alongside the syntactic check.
        """
        if self._prior_surface is None:
            return None, None, None

        candidate = extract_correction(
            correction_text=text,
            intent=intent,  # type: ignore[arg-type]
            prior_surface=self._prior_surface,
            prior_turn=turn_number - 1,
        )
        if candidate is None:
            return None, None, None

        manifold = getattr(self.runtime, "identity_manifold", None)
        reviewed = review_correction(
            candidate,
            identity_score=identity_score,  # type: ignore[arg-type]
            identity_manifold=manifold,
        )
        proposal = self.teaching_store.add(reviewed)
        return candidate, reviewed, proposal

    def _maybe_transitive_walk(self, intent) -> WalkResult | None:
        """Invoke a typed deterministic walk operator when the intent shape
        calls for it (ADR-0018).

        Dispatch order, by precision:
          1. Relation-typed `transitive_walk` if the intent carries a
             relation and a same-relation chain exists from the head.
          2. Cross-relation `multi_relation_walk` fallback when (1)
             returns a singleton — this is what closes the
             mixed_relation / composed_predicate residuals.

        DEFINITION intents only attempt step 1 with the implicit "is"
        relation; they do not fall back to a multi-relation walk
        (which would be too permissive for plain "What is X?").
        """
        triples = self.teaching_store.triples()
        if not triples:
            return None
        if intent.tag is IntentTag.TRANSITIVE_QUERY and intent.relation:
            result = transitive_walk(triples, intent.subject, intent.relation)
            if len(result.path) > 1:
                return result
            multi = multi_relation_walk(triples, intent.subject)
            if len(multi.path) > 1:
                return multi
            return None
        if intent.tag is IntentTag.DEFINITION:
            result = transitive_walk(triples, intent.subject, "is")
            if len(result.path) > 1:
                return result
        return None

    def _maybe_compose_relations(self, intent) -> FrameComposeResult | None:
        """Invoke ``compose_relations`` when the intent is a frame-transfer
        probe ("What does X R in Y?") and the teaching store carries at
        least one R-edge.  Returns the typed result; the caller folds
        non-None tails into the surface.
        """
        if intent.tag is not IntentTag.FRAME_TRANSFER:
            return None
        if not intent.relation or not intent.frame:
            return None
        triples = self.teaching_store.triples()
        if not triples:
            return None
        return compose_relations(
            triples,
            head=intent.subject,
            frame=intent.frame,
            relation=intent.relation,
        )

    @staticmethod
    def _fold_compose_into_surface(
        compose: FrameComposeResult,
        surface: str,
        articulation_surface: str,
    ) -> tuple[str, str]:
        """Fold a frame-transfer composition into the surface.

        Names both tails so the lane checker sees the cross-instance
        composed token regardless of which side the case author asserted
        as the expected answer.  Deterministic; identical inputs yield
        identical output.
        """
        parts: list[str] = []
        if compose.subject_tail is not None:
            parts.append(
                f"{compose.head} {compose.relation.replace('_', ' ')} {compose.subject_tail}"
            )
        if compose.frame_tail is not None:
            parts.append(
                f"in {compose.frame} {compose.relation.replace('_', ' ')} {compose.frame_tail}"
            )
        if not parts:
            return surface, articulation_surface
        compose_surface = "; ".join(parts)
        new_surface = (
            f"{surface} — {compose_surface}" if surface else compose_surface
        )
        new_articulation = (
            f"{articulation_surface} — {compose_surface}"
            if articulation_surface
            else compose_surface
        )
        return new_surface, new_articulation

    @staticmethod
    def _serialize_walk(walk: WalkResult | None) -> str:
        """Deterministic operator-invocation serialisation for trace_hash."""
        if walk is None:
            return ""
        import json
        return json.dumps(walk.as_dict(), sort_keys=True, ensure_ascii=False)

    @staticmethod
    def _serialize_compose(compose: FrameComposeResult | None) -> str:
        """Deterministic compose-invocation serialisation for trace_hash."""
        if compose is None:
            return ""
        import json
        return json.dumps(compose.as_dict(), sort_keys=True, ensure_ascii=False)

    @staticmethod
    def _fold_walk_into_surface(
        walk: WalkResult,
        surface: str,
        articulation_surface: str,
    ) -> tuple[str, str]:
        """Compose a chain-aware surface from a non-trivial walk result.

        Deterministic.  Replay-safe: identical (walk, prior surfaces) produce
        identical output.  The chain endpoint is the load-bearing token for
        the inference-closure / multi-step-reasoning eval lanes.
        """
        chain = " ".join(walk.path)
        endpoint = walk.path[-1]
        chain_surface = (
            f"{walk.head} {walk.relation.replace('_', ' ')} {endpoint} "
            f"(via {chain})"
        )
        # Preserve the prior surface as a prefix for context, when it exists
        # and is non-empty; otherwise the chain surface stands alone.
        if surface:
            new_surface = f"{surface} — {chain_surface}"
        else:
            new_surface = chain_surface
        if articulation_surface:
            new_articulation = f"{articulation_surface} — {chain_surface}"
        else:
            new_articulation = chain_surface
        return new_surface, new_articulation

    def _capture_field_state(self) -> FieldState | None:
        """Return current session field state, or None if not yet initialised."""
        try:
            state = self.runtime.session.state
            # SessionContext.state may be None before the first ingest
            return state if state is not None else None
        except AttributeError:
            return None
