"""Prediction loop — falsifiable belief revision via predict-and-test.

This is Phase 6E of the Rule 12 / CLS architecture. The prediction loop
makes the agent's beliefs *testable* by generating concrete predictions
about what the user will say or do next, then validating them against
incoming evidence on the next chat turn.

The engine is deliberately thin: all strategy lives in the
``memory-prediction`` skill (SKILL.md + config/prediction.yaml). This
module's only job is to:

  1. Read the recent context, build the LLM prompt, parse the structured
     prediction response, and write the resulting predictions to the
     graph as triples with ``prediction_outcome="pending"``.
  2. Read pending predictions, compare them against incoming
     observations via a second LLM call (validate mode), and update
     each pending triple's ``prediction_outcome`` to ``confirmed`` /
     ``surprising`` / unchanged.
  3. Walk pending predictions periodically and mark stale ones
     (older than ``max_pending_age_days``) as ``surprising`` so they
     decay rather than accumulating forever.

**The engine never crashes the chat loop.** Every operation wraps work
in try/except, returns a structured result dict, and degrades gracefully
when the LLM is missing or the skill config is unavailable. The
prediction loop is always optional — if it can't run, the agent still
works; only the falsifiability feedback is missing.

**Design parallels ``DreamPhaseHandler``**: skill-driven, two-mode LLM
calls, structured JSON outputs, instance-method API. Future Phase 6
sub-phases (affect, surfaces) will follow the same template.
"""

from __future__ import annotations

import json
import logging
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any

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

log = logging.getLogger(__name__)


_SKILL_NAME = "memory-prediction"

# These are documented as the only valid string values of
# ``prediction_outcome``. We do NOT enforce them as an enum — the
# field is a free-form string for forward compatibility — but new
# code should use these constants for consistency.
OUTCOME_PENDING = "pending"
OUTCOME_CONFIRMED = "confirmed"
OUTCOME_SURPRISING = "surprising"


def _now() -> datetime:
    return datetime.now(UTC)


def _empty_result(operation: str) -> dict:
    """Standard shape for a prediction-loop operation result."""
    return {
        "operation": operation,
        "success": True,
        "skill_loaded": False,
        "predictions_generated": 0,
        "predictions_validated": 0,
        "confirmed": 0,
        "surprising": 0,
        "still_pending": 0,
        "stale_pruned": 0,
        "model_tier": None,
        "duration_ms": 0,
        "skipped_reason": None,
        "error": None,
    }


# ── JSON parsing (shared shape with dream.py) ──────────────────────────────


def _extract_json(response: str) -> Any:
    """Extract a JSON object/array from an LLM response, fenced or bare.

    Mirrors ``dream.extract_json_from_response`` — duplicated here so the
    prediction loop has no dependency on the dream module beyond the
    shared graph and skill loader.
    """
    if not response or not isinstance(response, str):
        return None
    # Try fenced ```json block first
    import re
    fenced = re.search(r"```(?:json)?\s*(\{.*?\}|\[.*?\])\s*```", response, re.DOTALL)
    if fenced:
        try:
            return json.loads(fenced.group(1))
        except json.JSONDecodeError:
            pass
    # Bare JSON fallback -- use json.loads/raw_decode directly instead of
    # a naive depth scanner, because the depth scanner fails on JSON where
    # keys or values contain literal brace/bracket chars (e.g. {"{": null}).
    response = response.strip()
    for opener in ("{", "["):
        start = response.find(opener)
        if start == -1:
            continue
        candidate = response[start:]
        try:
            return json.loads(candidate)
        except json.JSONDecodeError:
            pass
        try:
            decoder = json.JSONDecoder()
            obj, _ = decoder.raw_decode(candidate)
            return obj
        except (json.JSONDecodeError, ValueError):
            pass
    return None


# ── PredictionEngine ───────────────────────────────────────────────────────


class PredictionEngine:
    """Generate, validate, and prune testable predictions about user behavior.

    Construct once per process. Holds references to the knowledge graph,
    skill loader, and LLM provider. None of the references are required —
    when something is missing, the corresponding operation degrades
    gracefully.
    """

    def __init__(
        self,
        graph: KnowledgeGraph,
        skill_loader: SkillLoader | None = None,
        llm: LLMProvider | None = None,
        model_router: object | None = None,
        activity_logger: object | None = None,
    ) -> None:
        self._graph = graph
        self._skill_loader = skill_loader
        self._llm = llm
        self._model_router = model_router
        self._activity = activity_logger

    # ── Configuration ─────────────────────────────────────────────────────

    def _config(self) -> dict:
        """Load the prediction skill's config YAML, with degraded-mode defaults."""
        defaults = {
            "predictions_per_generation": 4,
            "min_generation_confidence": 0.55,
            "max_generation_confidence": 0.95,
            "max_pending_age_days": 14,
            "min_pending_floor": 2,
            "confirmation_confidence_boost": 0.05,
            "model_tier": "capable",
            "max_tokens_per_generation": 8000,
        }
        if self._skill_loader is None:
            return defaults
        try:
            cfg = self._skill_loader.load_skill_config_file(_SKILL_NAME, "prediction.yaml")
        except Exception:
            log.debug("predictions: load_skill_config_file failed", exc_info=True)
            return defaults
        if not isinstance(cfg, dict):
            return defaults
        merged = dict(defaults)
        for k, v in cfg.items():
            merged[k] = v
        return merged

    def _skill_body(self) -> str:
        """Read the SKILL.md body for the prediction skill."""
        if self._skill_loader is None:
            return ""
        try:
            skill = self._skill_loader.get_skill(_SKILL_NAME)
        except Exception:
            return ""
        return getattr(skill, "body", "") if skill else ""

    def _resolve_model(self, tier: str | None) -> str | None:
        if not tier or self._model_router is None:
            return None
        try:
            return self._model_router.resolve(tier)
        except Exception:
            return None

    # ── Mode A: Generate predictions ──────────────────────────────────────

    async def generate_predictions(
        self,
        context_text: str,
        *,
        max_count: int | None = None,
    ) -> dict:
        """Ask the LLM to generate N testable predictions from the context.

        ``context_text`` is the recent conversation summary (and any
        relevant graph state the caller wants to include). The LLM
        produces predictions in the documented schema; we filter by
        the configured confidence range and write each surviving
        prediction to the graph as a triple with
        ``prediction_outcome="pending"``.

        Returns a result dict with counts. Never raises.
        """
        result = _empty_result("generate")
        if not context_text:
            result["skipped_reason"] = "no_context"
            return result

        cfg = self._config()
        result["skill_loaded"] = self._skill_loader is not None
        result["model_tier"] = cfg.get("model_tier", "capable")

        if self._llm is None:
            result["skipped_reason"] = "no_llm_available"
            self._emit("predict_generate_skipped", "No LLM", result)
            return result

        target_count = int(max_count or cfg.get("predictions_per_generation", 4))
        target_count = max(1, min(target_count, 5))  # hard cap matches skill rule

        started = _now()
        try:
            payload = json.dumps({
                "mode": "generate",
                "context": context_text,
                "target_count": target_count,
            })
            response = await self._call_llm(payload, cfg)
            parsed = _extract_json(response)
        except Exception:
            log.exception("predictions: generate LLM call failed")
            result["success"] = False
            result["error"] = "llm_call_failed"
            return result

        if not isinstance(parsed, dict):
            result["skipped_reason"] = "llm_no_parseable_output"
            self._emit("predict_generate_no_output", "Unparseable LLM output", result, importance="high")
            return result

        predictions = parsed.get("predictions", [])
        if not isinstance(predictions, list):
            result["skipped_reason"] = "llm_predictions_not_list"
            return result

        min_conf = float(cfg.get("min_generation_confidence", 0.55))
        max_conf = float(cfg.get("max_generation_confidence", 0.95))

        written = 0
        for p in predictions:
            if written >= target_count:
                break
            if not isinstance(p, dict):
                continue
            try:
                conf = float(p.get("confidence", 0.0))
            except (TypeError, ValueError):
                conf = 0.0
            if not (min_conf <= conf <= max_conf):
                continue
            tid = self._write_prediction(p)
            if tid is not None:
                written += 1

        result["predictions_generated"] = written
        result["duration_ms"] = int((_now() - started).total_seconds() * 1000)
        self._emit(
            "predict_generated",
            f"Generated {written} pending predictions",
            result,
        )
        return result

    def _write_prediction(self, p: dict) -> int | None:
        """Write one prediction to the graph as a pending triple."""
        subject = p.get("subject")
        predicate = p.get("predicate")
        obj = p.get("object")
        if not subject or not predicate or obj is None:
            return None
        try:
            tid = self._graph.add_triple(
                subject=str(subject),
                predicate=str(predicate),
                obj=str(obj),
                subject_type=str(p.get("subject_type", "concept")),
                object_type=str(p.get("object_type", "concept")),
                knowledge_type=str(p.get("knowledge_type", "declarative")),
                epistemic_status="uncertain",  # predictions start uncertain
                confidence=float(p.get("confidence", 0.6)),
                authority="inferred",  # not user-confirmed; the agent guessed
                system_layer="neocortical",
                prediction_outcome=OUTCOME_PENDING,
            )
            # Stash the falsifiability hint + rationale on the triple
            # so the validation pass and the future review surface can
            # show the user *why* the agent predicted this.
            for t in self._graph._triples:  # noqa: SLF001
                if t.get("id") == tid:
                    notes = t.setdefault("notes", [])
                    if "falsifiability_hint" in p:
                        notes.append(f"falsifiability: {p['falsifiability_hint']}")
                    if "rationale" in p:
                        notes.append(f"rationale: {p['rationale']}")
                    break
            return tid
        except Exception:
            log.debug("predictions: add_triple failed for %s/%s/%s",
                      subject, predicate, obj, exc_info=True)
            return None

    # ── Mode B: Validate pending predictions ───────────────────────────────

    async def validate_predictions(
        self,
        observation_summary: str,
    ) -> dict:
        """Validate pending predictions against new observations.

        ``observation_summary`` is the agent's description of what just
        happened — typically a paragraph summarizing the latest user
        message and any newly extracted facts. The LLM compares it to
        the pending predictions and produces a verdict per prediction.

        Each ``confirmed`` verdict bumps the source triple's confidence
        and sets ``prediction_outcome=confirmed``. Each ``surprising``
        verdict marks the source triple as ``epistemic_status=uncertain``
        and sets ``prediction_outcome=surprising``. ``still_pending``
        verdicts are no-ops (the prediction stays in the graph).

        Returns a result dict with counts. Never raises.
        """
        result = _empty_result("validate")
        cfg = self._config()
        result["skill_loaded"] = self._skill_loader is not None
        result["model_tier"] = cfg.get("model_tier", "capable")

        pending = self.list_pending()
        if not pending:
            result["skipped_reason"] = "no_pending_predictions"
            return result

        if not observation_summary:
            result["skipped_reason"] = "no_observation_summary"
            return result

        if self._llm is None:
            result["skipped_reason"] = "no_llm_available"
            self._emit("predict_validate_skipped", "No LLM", result)
            return result

        started = _now()
        try:
            payload = json.dumps({
                "mode": "validate",
                "observation_summary": observation_summary,
                "pending_predictions": [
                    {
                        "id": t["id"],
                        "subject": self._graph._registry.get_name(t["subject"]),  # noqa: SLF001
                        "predicate": t["predicate"],
                        "object": self._graph._registry.get_name(t["object"]) if t.get("object") else t.get("object_literal"),  # noqa: SLF001
                        "knowledge_type": t.get("knowledge_type"),
                        "confidence": t.get("confidence"),
                    }
                    for t in pending
                ],
            })
            response = await self._call_llm(payload, cfg)
            parsed = _extract_json(response)
        except Exception:
            log.exception("predictions: validate LLM call failed")
            result["success"] = False
            result["error"] = "llm_call_failed"
            return result

        if not isinstance(parsed, dict):
            result["skipped_reason"] = "llm_no_parseable_output"
            return result

        verdicts = parsed.get("verdicts", [])
        if not isinstance(verdicts, list):
            return result

        boost = float(cfg.get("confirmation_confidence_boost", 0.05))
        for v in verdicts:
            if not isinstance(v, dict):
                continue
            applied = self._apply_verdict(v, boost)
            if applied == OUTCOME_CONFIRMED:
                result["confirmed"] += 1
            elif applied == OUTCOME_SURPRISING:
                result["surprising"] += 1
            else:
                result["still_pending"] += 1
            result["predictions_validated"] += 1

        result["duration_ms"] = int((_now() - started).total_seconds() * 1000)
        self._emit(
            "predict_validated",
            f"Validated {result['predictions_validated']}: {result['confirmed']} confirmed, {result['surprising']} surprising",
            result,
        )
        return result

    def _apply_verdict(self, v: dict, boost: float) -> str | None:
        """Apply one verdict to its target triple. Returns the outcome string."""
        try:
            tid = int(v.get("pending_triple_id"))
        except (TypeError, ValueError):
            return None
        outcome = v.get("outcome")
        if outcome not in (OUTCOME_CONFIRMED, OUTCOME_SURPRISING, "still_pending"):
            return None
        if outcome == "still_pending":
            return "still_pending"
        for t in self._graph._triples:  # noqa: SLF001
            if t.get("id") != tid:
                continue
            if t.get("prediction_outcome") != OUTCOME_PENDING:
                # Don't double-validate; respect prior outcomes
                return None
            t["prediction_outcome"] = outcome
            if outcome == OUTCOME_CONFIRMED:
                t["confidence"] = min(1.0, float(t.get("confidence", 0.5)) + boost)
                t["epistemic_status"] = "believed"
                t["last_verified"] = _now().isoformat()
            else:  # surprising
                t["epistemic_status"] = "uncertain"
            self._graph._dirty = True  # noqa: SLF001
            return outcome
        return None

    # ── Mode C: Prune stale pending predictions ───────────────────────────

    def prune_stale(self) -> dict:
        """Mark predictions older than the configured threshold as surprising.

        A prediction the agent made and that was never validated by
        incoming evidence is — by elapsed-time decay — a failed
        prediction. Marking it surprising rather than just deleting
        it preserves the trail of what the agent thought would happen.

        Returns a result dict; never raises.
        """
        result = _empty_result("prune_stale")
        cfg = self._config()
        max_age = int(cfg.get("max_pending_age_days", 14))
        floor = int(cfg.get("min_pending_floor", 2))
        cutoff = _now() - timedelta(days=max_age)

        pending = sorted(
            self.list_pending(),
            key=lambda t: t.get("source_date", ""),
        )
        if len(pending) <= floor:
            result["skipped_reason"] = "below_pending_floor"
            return result

        # Walk oldest-first, prune anything past cutoff, but keep at
        # least ``floor`` predictions alive.
        prunable = pending[: max(0, len(pending) - floor)]
        marked = 0
        for t in prunable:
            sd = t.get("source_date") or t.get("valid_from") or ""
            try:
                ts = datetime.fromisoformat(sd.replace("Z", "+00:00")) if sd else _now()
            except (ValueError, TypeError):
                ts = _now()
            if ts.tzinfo is None:
                ts = ts.replace(tzinfo=UTC)
            if ts < cutoff:
                t["prediction_outcome"] = OUTCOME_SURPRISING
                t["epistemic_status"] = "uncertain"
                self._graph._dirty = True  # noqa: SLF001
                marked += 1
        result["stale_pruned"] = marked
        if marked:
            self._emit(
                "predict_pruned_stale",
                f"Marked {marked} stale predictions as surprising",
                result,
            )
        return result

    # ── Read helpers ──────────────────────────────────────────────────────

    def list_pending(self) -> list[dict]:
        """Return all triples currently in the pending state."""
        return [
            t for t in self._graph._triples  # noqa: SLF001
            if t.get("status") == "active"
            and t.get("prediction_outcome") == OUTCOME_PENDING
        ]

    def stats(self) -> dict:
        """Return aggregate prediction-loop stats over all triples."""
        confirmed = surprising = pending = 0
        for t in self._graph._triples:  # noqa: SLF001
            outcome = t.get("prediction_outcome")
            if outcome == OUTCOME_PENDING:
                pending += 1
            elif outcome == OUTCOME_CONFIRMED:
                confirmed += 1
            elif outcome == OUTCOME_SURPRISING:
                surprising += 1
        total = confirmed + surprising
        accuracy = (confirmed / total) if total > 0 else None
        return {
            "pending": pending,
            "confirmed": confirmed,
            "surprising": surprising,
            "total_resolved": total,
            "accuracy": round(accuracy, 3) if accuracy is not None else None,
        }

    # ── Internal: LLM call + activity logging ─────────────────────────────

    async def _call_llm(self, payload: str, cfg: dict) -> str:
        """Call the LLM with the prediction skill body as the system prompt."""
        system = self._skill_body()
        if not system:
            log.warning("predictions: skill body missing, cannot run")
            return ""
        model = self._resolve_model(cfg.get("model_tier"))
        # Lazy import — keeps this module importable without the LLM layer
        # during tests that stub the provider.
        from agntspace.llm.base import Message

        response = await self._llm.complete(
            messages=[Message(role="user", content=payload)],
            system=system,
            model=model,
            max_tokens=int(cfg.get("max_tokens_per_generation", 8000)),
        )
        # complete() returns LLMResponse; older stubs may return a raw
        # string — support both by preferring `.content` when present.
        return getattr(response, "content", response)

    # ── Chat-loop integration (Phases 6E configurable modes) ────────────

    # Per-process turn counter for mode C (every_nth). Reset on server
    # restart. This is intentionally NOT persisted — a restart naturally
    # resets the cadence, and persisting it would require a DB round-trip
    # on every chat turn for zero user-visible benefit.
    _turn_counter: int = 0
    _last_generate_turn: int = 0

    async def chat_hook_post_response(
        self,
        context_text: str,
        observation_text: str,
    ) -> dict | None:
        """Called by Agent.chat() AFTER the response is sent.

        Implements the three prediction modes:

        - **every_turn** (A): always generates predictions from context_text
          AND validates pending predictions against observation_text.
        - **daily** (B): no-op — deep dream handles everything.
        - **every_nth** (C, default): generates every Nth turn, validates
          on the turn after a generation. N comes from the skill config.

        ``context_text`` is the recent conversation context (system prompt
        + last few turns). ``observation_text`` is the user's latest
        message — used as the observation summary for validation.

        Returns a result dict when it did something, None when it was a
        no-op. Never raises.
        """
        cfg = self._config()
        mode = str(cfg.get("prediction_mode", "every_nth"))
        nth = max(1, int(cfg.get("nth_turn", 5)))

        PredictionEngine._turn_counter += 1
        turn = PredictionEngine._turn_counter

        try:
            if mode == "daily":
                # Mode B: total no-op — deep dream owns the loop.
                return None

            if mode == "every_turn":
                # Mode A: validate first (against the new observation),
                # then generate new predictions from the full context.
                result: dict = {"mode": "every_turn", "turn": turn}
                if observation_text and self.list_pending():
                    v = await self.validate_predictions(observation_text)
                    result["validated"] = v
                if context_text:
                    g = await self.generate_predictions(context_text)
                    result["generated"] = g
                return result

            # Mode C: every_nth (default)
            turns_since_generate = turn - PredictionEngine._last_generate_turn
            result = {"mode": "every_nth", "turn": turn, "nth": nth}

            # On the turn RIGHT AFTER a generation, validate first.
            if (
                turns_since_generate == 1
                and PredictionEngine._last_generate_turn > 0
                and observation_text
                and self.list_pending()
            ):
                v = await self.validate_predictions(observation_text)
                result["validated"] = v

            # On every Nth turn, generate.
            if turns_since_generate >= nth:
                if context_text:
                    g = await self.generate_predictions(context_text)
                    result["generated"] = g
                    PredictionEngine._last_generate_turn = turn

            return result

        except Exception:
            log.debug("predictions: chat_hook_post_response failed", exc_info=True)
            return None

    @classmethod
    def reset_turn_counter(cls) -> None:
        """Reset the turn counter. Tests call this between cases."""
        cls._turn_counter = 0
        cls._last_generate_turn = 0

    def _emit(
        self,
        action: str,
        summary: str,
        details: dict | None = None,
        importance: str = "medium",
    ) -> None:
        """Emit a memory/predict event. Never raises."""
        if self._activity is None or not hasattr(self._activity, "log"):
            return
        try:
            self._activity.log("memory", action, summary, details or {}, importance=importance)
        except Exception:
            log.debug("predictions: activity.log failed", exc_info=True)
