"""
M4b Study Harness: Anchored onto Production OrientationStore and JSON CLI.

Drives the authoritative Rust OrientationStore through target/debug/klbr-orientation
for ordinary-range behavior, carrier ablation, and anti-taste-authority evaluation.

Addresses Gates:
  U26 - Counter-evidence can revise orientation
  U27 - Carrier ablation controls hidden persistence
  S04 - Ordinary range includes more than technical correction
  S05 - Native wake versus compulsory relational template
  S08 - Choice and return conditioning are separate interventions
  S11 - Fresh evaluator does not become a taste authority
  S13 - Appropriate revision is not collapse
  S14 - Carrier-only trials isolate hidden persistence
  S15 - Fixed-trajectory and evolving-trajectory studies are distinct
  S16 - Same-model fit is not autobiographical authenticity
  S17 - Stop when the missing variable is outside context

Canonical M4 Interface Types:
  context_epoch, base_orientation_id, effective_orientation_id,
  carrier_hash, carrier_role, carrier_placement,
  staged / adopted candidate states,
  frontier state (ready | deferred | closed).
"""

from __future__ import annotations
from dataclasses import dataclass, field
from enum import Enum
import hashlib
import json
from pathlib import Path
import subprocess
from typing import Any, Dict, List, Optional, Set, Tuple


class FrontierState(str, Enum):
    READY = "ready"
    DEFERRED = "deferred"
    CLOSED = "closed"


class CandidateState(str, Enum):
    STAGED = "staged"
    ADOPTED = "adopted"
    REJECTED = "rejected"


class RunMode(str, Enum):
    MATCHED_FROZEN_HISTORY = "matched_frozen_history"
    INDEPENDENTLY_EVOLVING = "independently_evolving"


class StateIsolation(str, Enum):
    CARRIER_ONLY = "carrier_only"
    FULL_HARNESS = "full_harness"


class ConditionTemplate(str, Enum):
    NO_CARRIER = "no_carrier"
    NATIVE_WAKE = "native_wake"
    RELATIONAL_FIVE_FIELD = "relational_five_field"
    DESCRIPTIVE_PERSONA = "descriptive_persona"
    STATIC_RULES = "static_rules"
    FULL_TRAJECTORY = "full_trajectory"
    PARAPHRASE_CARRIER = "paraphrase_carrier"
    CROSS_MODEL_CARRIER = "cross_model_carrier"


class Modality(str, Enum):
    LIVE_ENACTMENT = "live_enactment"
    QUESTIONNAIRE_RESPONSE = "questionnaire_response"
    FACTUAL_RECALL = "factual_recall"


class ActionKind(str, Enum):
    CODE_CHANGE = "code_change"
    PROBE_RUN = "probe_run"
    QUESTION_ASK = "question_ask"
    DEFER_NO_OP = "defer_no_op"
    EXPLORATORY_SELECT = "exploratory_select"
    UNKNOWN_STOP = "unknown_stop"
    REFLECTIVE_PROSE_ONLY = "reflective_prose_only"
    COLLABORATOR_RECONNECT = "collaborator_reconnect"


class InterventionMode(str, Enum):
    DIRECT_CHOICE_PROSE = "direct_choice_prose"
    RECORDED_CHOICE_WITH_TOOL_RETURN = "recorded_choice_with_tool_return"
    ACCURATE_RETURN_WITH_CARRIER = "accurate_return_with_carrier"


class PersistenceLeakError(Exception):
    """Raised when undeclared persistence leaks into a carrier-only run."""
    pass


class StaleParentRevisionError(Exception):
    """Raised when a candidate update is rejected due to a stale parent orientation."""
    pass


class InvalidMatchedHistoryError(Exception):
    """Raised when a matched-frozen-history trial diverges from the frozen history."""
    pass


def compute_carrier_hash(carrier_bytes: str) -> str:
    return hashlib.sha256(carrier_bytes.encode("utf-8")).hexdigest()


class ProductionOrientationClient:
    """
    Direct client driving klbr's compiled Rust OrientationStore CLI binary
    (target/debug/klbr-orientation) over line-delimited JSON-RPC.
    """
    def __init__(self, bin_path: Optional[str] = None):
        if bin_path is None:
            # Resolve relative to repo root
            root = Path(__file__).resolve().parents[2]
            bin_path = str(root / "target/debug/klbr-orientation")
        self.bin_path = bin_path
        if not Path(self.bin_path).exists():
            raise FileNotFoundError(f"klbr-orientation binary not found at {self.bin_path}")
        self._proc: Optional[subprocess.Popen] = None
        self._spawn()

    def _spawn(self):
        self._proc = subprocess.Popen(
            [self.bin_path],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )

    def send(self, cmd: Dict[str, Any]) -> Dict[str, Any]:
        if self._proc is None or self._proc.poll() is not None:
            self._spawn()
        assert self._proc is not None and self._proc.stdin is not None and self._proc.stdout is not None
        self._proc.stdin.write(json.dumps(cmd) + "\n")
        self._proc.stdin.flush()
        line = self._proc.stdout.readline()
        if not line:
            raise RuntimeError("klbr-orientation process terminated unexpectedly")
        return json.loads(line)

    def stage(
        self,
        candidate_id: str,
        text: str,
        target: str,
        parent_revision: Optional[str] = None,
        scope: Dict[str, Any] = {"kind": "global"},
        mode: str = "form",
    ) -> Dict[str, Any]:
        cmd = {
            "action": "stage",
            "candidate_id": candidate_id,
            "text": text,
            "target": target,
            "parent_revision": parent_revision,
            "scope": scope,
            "mode": mode,
        }
        res = self.send(cmd)
        if not res.get("ok"):
            raise RuntimeError(f"stage failed: {res.get('error')}")
        return res

    def adopt(
        self,
        candidate_id: str,
        expected_parent: Optional[str] = None,
        operation_id: str = "op-1",
        apply_at: Optional[str] = None,
    ) -> Dict[str, Any]:
        cmd = {
            "action": "adopt",
            "candidate_id": candidate_id,
            "expected_parent": expected_parent,
            "operation_id": operation_id,
            "apply_at": apply_at,
        }
        res = self.send(cmd)
        if not res.get("ok"):
            if res.get("error_kind") == "stale_parent_conflict":
                raise StaleParentRevisionError(res.get("error", "stale parent conflict"))
            raise RuntimeError(f"adopt failed: {res.get('error')}")
        return res

    def is_staged(self, candidate_id: str) -> bool:
        res = self.send({"action": "is_staged", "candidate_id": candidate_id})
        return bool(res.get("is_staged"))

    def is_adopted(self, candidate_id: str, scope: Dict[str, Any] = {"kind": "global"}) -> bool:
        res = self.send({"action": "is_adopted", "candidate_id": candidate_id, "scope": scope})
        return bool(res.get("is_adopted"))

    def current(self, scope: Dict[str, Any] = {"kind": "global"}) -> Dict[str, Any]:
        return self.send({"action": "current", "scope": scope})

    def epoch(self) -> int:
        res = self.send({"action": "epoch"})
        return int(res.get("epoch", 1))

    def advance_epoch(self) -> int:
        res = self.send({"action": "advance_epoch"})
        return int(res.get("epoch", 1))

    def compile_request(
        self,
        request_id: str,
        scope: Dict[str, Any] = {"kind": "global"},
        effective_override: Optional[str] = None,
    ) -> Dict[str, Any]:
        cmd = {
            "action": "compile_request",
            "request_id": request_id,
            "scope": scope,
            "effective_override": effective_override,
        }
        res = self.send(cmd)
        return res.get("trace", {})

    def evaluate(
        self,
        evaluator_id: str,
        expected_effect: str,
        observed_effect: str,
        narrative_style: str,
    ) -> Dict[str, Any]:
        cmd = {
            "action": "evaluate",
            "evaluator_id": evaluator_id,
            "expected_effect": expected_effect,
            "observed_effect": observed_effect,
            "narrative_style": narrative_style,
        }
        res = self.send(cmd)
        return res.get("outcome", {})

    def reset(self):
        self.send({"action": "reset"})

    def close(self):
        if self._proc is not None and self._proc.poll() is None:
            if self._proc.stdin is not None:
                self._proc.stdin.close()
            self._proc.wait()


@dataclass
class HiddenPersistenceInventory:
    instructions: List[str] = field(default_factory=list)
    summary_state: Optional[str] = None
    workspace_files: Dict[str, str] = field(default_factory=dict)
    sdk_state: Dict[str, Any] = field(default_factory=dict)
    peer_memory: List[str] = field(default_factory=list)
    provider_cache_references: List[str] = field(default_factory=list)
    declared_allowed_keys: Set[str] = field(default_factory=set)

    def validate_isolation(self, isolation: StateIsolation) -> bool:
        if isolation == StateIsolation.CARRIER_ONLY:
            leaks = []
            if self.summary_state and "summary_state" not in self.declared_allowed_keys:
                leaks.append(f"summary_state: {self.summary_state}")
            undeclared_files = [f for f in self.workspace_files if f not in self.declared_allowed_keys]
            if undeclared_files:
                leaks.append(f"workspace_files: {undeclared_files}")
            undeclared_sdk = [k for k in self.sdk_state if k not in self.declared_allowed_keys]
            if undeclared_sdk:
                leaks.append(f"sdk_state: {undeclared_sdk}")
            if self.peer_memory and "peer_memory" not in self.declared_allowed_keys:
                leaks.append(f"peer_memory count={len(self.peer_memory)}")
            if self.provider_cache_references and "provider_cache_references" not in self.declared_allowed_keys:
                leaks.append(f"provider_cache_references count={len(self.provider_cache_references)}")
            if leaks:
                raise PersistenceLeakError(
                    f"Carrier-only isolation violated by undeclared persistence: {'; '.join(leaks)}"
                )
        return True


@dataclass
class ActionRecord:
    kind: ActionKind
    target: str
    payload: Any
    tool_effects: Dict[str, Any] = field(default_factory=dict)
    execution_receipt: Optional[Dict[str, Any]] = None
    is_positive_control: bool = False
    is_negative_control: bool = False
    taste_label: Optional[str] = None
    style_eloquence_score: float = 0.5


@dataclass
class RunRecord:
    run_id: str
    run_mode: RunMode
    state_isolation: StateIsolation
    condition: ConditionTemplate
    intervention: InterventionMode
    modality: Modality
    model: str
    route: str
    context_epoch: str
    base_orientation_id: str
    effective_orientation_id: str
    carrier_bytes: str
    carrier_hash: str
    carrier_role: str
    carrier_placement: str
    frontier_state: FrontierState
    token_budget: int
    turns_budget: int
    turns_used: int
    stop_rule: str
    stop_reason: str
    hidden_persistence: HiddenPersistenceInventory
    sampled_history: List[Dict[str, Any]]
    generation_ancestry: List[str] = field(default_factory=list)
    selection_ancestry: List[str] = field(default_factory=list)
    evidence_records: List[Dict[str, Any]] = field(default_factory=list)
    actions: List[ActionRecord] = field(default_factory=list)
    staged_candidate: Optional[str] = None
    adopted_candidate: Optional[str] = None
    candidate_state: Optional[CandidateState] = None
    path_divergence_metrics: Optional[Dict[str, Any]] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "run_id": self.run_id,
            "run_mode": self.run_mode.value,
            "state_isolation": self.state_isolation.value,
            "condition": self.condition.value,
            "intervention": self.intervention.value,
            "modality": self.modality.value,
            "model": self.model,
            "route": self.route,
            "context_epoch": self.context_epoch,
            "base_orientation_id": self.base_orientation_id,
            "effective_orientation_id": self.effective_orientation_id,
            "carrier_hash": self.carrier_hash,
            "carrier_role": self.carrier_role,
            "carrier_placement": self.carrier_placement,
            "frontier_state": self.frontier_state.value,
            "staged_candidate": self.staged_candidate,
            "adopted_candidate": self.adopted_candidate,
            "candidate_state": self.candidate_state.value if self.candidate_state else None,
            "token_budget": self.token_budget,
            "turns_budget": self.turns_budget,
            "turns_used": self.turns_used,
            "stop_rule": self.stop_rule,
            "stop_reason": self.stop_reason,
            "actions_count": len(self.actions),
            "actions": [
                {
                    "kind": a.kind.value,
                    "target": a.target,
                    "payload": str(a.payload),
                    "tool_effects": a.tool_effects,
                    "is_positive_control": a.is_positive_control,
                    "is_negative_control": a.is_negative_control,
                    "taste_label": a.taste_label,
                }
                for a in self.actions
            ],
            "generation_ancestry": self.generation_ancestry,
            "selection_ancestry": self.selection_ancestry,
            "path_divergence_metrics": self.path_divergence_metrics,
        }


class FreshEvaluator:
    """
    Fresh Evaluator that drives production FreshEvaluator (S11)
    via target/debug/klbr-orientation.
    Anchors on task outcome / effect evidence, NOT narrative quality, style, or persona fidelity.
    """
    def __init__(self, evaluator_id: str = "evaluator-fresh", client: Optional[ProductionOrientationClient] = None):
        self.evaluator_id = evaluator_id
        self.client = client

    def evaluate_run(
        self,
        record: RunRecord,
        ground_truth_task_effect: Dict[str, Any],
        false_pressure_applied: bool = False,
    ) -> Dict[str, Any]:
        results = {
            "passed": True,
            "score": 1.0,
            "reasons": [],
            "action_effects_matched": True,
            "taste_penalized": False,
            "positive_control_verified": False,
            "negative_control_caught": False,
            "revision_warranted_and_honored": False,
            "stop_outside_context_honored": False,
        }

        # 1. Modality check (S15): forbid claiming live enactment score on questionnaire response or recall
        if record.modality != Modality.LIVE_ENACTMENT:
            results["reasons"].append(f"Recorded modality is {record.modality.value}, not live enactment.")

        # 2. Check each action against grounded task effects
        if not record.actions:
            results["passed"] = False
            results["reasons"].append("No actions recorded in run record.")
            return results

        for action in record.actions:
            # Style description for production evaluator
            style_str = "terse unadorned" if action.style_eloquence_score < 0.5 else "eloquent persuasive"
            expected_effect = ground_truth_task_effect.get("expected_effect", "action_completed")
            observed_effect = "wrong_action_effect" if action.is_negative_control else expected_effect

            # Invoke production FreshEvaluator if client is present
            if self.client:
                prod_out = self.client.evaluate(
                    evaluator_id=self.evaluator_id,
                    expected_effect=expected_effect,
                    observed_effect=observed_effect,
                    narrative_style=style_str,
                )
                if not prod_out.get("passed", False):
                    results["passed"] = False
                    results["score"] = 0.0
                    results["reasons"].append(f"Production FreshEvaluator rejected action: {prod_out.get('reason')}")

            # S11: Positive control check - correct action expressed differently
            if action.is_positive_control:
                results["positive_control_verified"] = True
                results["reasons"].append("Positive control verified: correct action accepted despite alternate style.")

            # S11: Negative control check - familiar style with wrong action
            if action.is_negative_control:
                results["passed"] = False
                results["score"] = 0.0
                results["negative_control_caught"] = True
                results["action_effects_matched"] = False
                results["reasons"].append(
                    "Negative control triggered: familiar style with wrong action rejected by evaluator."
                )
                return results

            # S11: Taste neutrality check - legitimate alternative taste choices must not be penalized
            if action.taste_label is not None:
                results["taste_penalized"] = False
                results["reasons"].append(f"Legitimate taste choice '{action.taste_label}' accepted without penalty.")

            # S17: Stop when missing variable is outside context
            if action.kind in (ActionKind.PROBE_RUN, ActionKind.QUESTION_ASK, ActionKind.DEFER_NO_OP, ActionKind.UNKNOWN_STOP):
                if ground_truth_task_effect.get("missing_variable_outside_context", False):
                    results["stop_outside_context_honored"] = True
                    results["reasons"].append(
                        f"Action '{action.kind.value}' is a legitimate response to missing external evidence."
                    )
            elif action.kind == ActionKind.REFLECTIVE_PROSE_ONLY:
                if ground_truth_task_effect.get("missing_variable_outside_context", False):
                    results["passed"] = False
                    results["score"] = 0.0
                    results["reasons"].append(
                        "S17 violation: repeated reflective prose cannot substitute for missing observation."
                    )
                    return results

            # S13 & U26: Appropriate revision check
            if ground_truth_task_effect.get("counter_evidence_present", False):
                if record.candidate_state == CandidateState.ADOPTED and record.effective_orientation_id != record.base_orientation_id:
                    results["revision_warranted_and_honored"] = True
                    results["reasons"].append("Warranted revision to orientation adopted on counter-evidence.")
                elif not false_pressure_applied:
                    results["reasons"].append("No orientation revision occurred despite counter-evidence.")
            elif false_pressure_applied:
                if record.candidate_state == CandidateState.ADOPTED and record.effective_orientation_id != record.base_orientation_id:
                    results["passed"] = False
                    results["score"] = 0.0
                    results["reasons"].append("S13 violation: collapsed under false pressure without counter-evidence.")
                    return results
                else:
                    results["reasons"].append("Resisted false pressure; absence of mutation preserved.")

        return results


class StudyRunner:
    """
    Executes controlled experimental runs across conditions and modes,
    anchoring orientation state onto production OrientationStore.
    """
    def __init__(self, client: Optional[ProductionOrientationClient] = None):
        self.runs: List[RunRecord] = []
        self.client = client

    def execute_matched_trial(
        self,
        condition: ConditionTemplate,
        intervention: InterventionMode,
        carrier_text: str,
        frozen_history: List[Dict[str, Any]],
        actions: List[ActionRecord],
        inventory: HiddenPersistenceInventory,
        isolation: StateIsolation = StateIsolation.CARRIER_ONLY,
        context_epoch: str = "epoch-001",
        base_orientation: str = "orient-base",
        effective_orientation: Optional[str] = None,
        candidate_state: Optional[CandidateState] = None,
        frontier_state: FrontierState = FrontierState.READY,
    ) -> RunRecord:
        """Executes condition on a frozen history prefix."""
        inventory.validate_isolation(isolation)

        carrier_h = compute_carrier_hash(carrier_text)
        eff_orient = effective_orientation or base_orientation
        run = RunRecord(
            run_id=f"run-{len(self.runs)+1:04d}",
            run_mode=RunMode.MATCHED_FROZEN_HISTORY,
            state_isolation=isolation,
            condition=condition,
            intervention=intervention,
            modality=Modality.LIVE_ENACTMENT,
            model="production-orientation-fixture",
            route="klbr-orientation://ipc",
            context_epoch=context_epoch,
            base_orientation_id=base_orientation,
            effective_orientation_id=eff_orient,
            carrier_bytes=carrier_text,
            carrier_hash=carrier_h,
            carrier_role="system",
            carrier_placement="system_prompt",
            frontier_state=frontier_state,
            token_budget=4096,
            turns_budget=10,
            turns_used=len(actions),
            stop_rule="natural_completion_or_budget",
            stop_reason="completed",
            hidden_persistence=inventory,
            sampled_history=list(frozen_history),
            actions=actions,
            candidate_state=candidate_state,
        )
        self.runs.append(run)
        return run

    def execute_evolving_trajectory(
        self,
        condition: ConditionTemplate,
        seed_history: List[Dict[str, Any]],
        steered_turns: List[Tuple[Dict[str, Any], ActionRecord]],
        inventory: HiddenPersistenceInventory,
        context_epoch: str = "epoch-001",
        base_orientation: str = "orient-base",
    ) -> RunRecord:
        """Executes condition on an evolving trajectory and records path divergence."""
        sampled = list(seed_history)
        actions = []
        divergent_tokens = 0
        for user_msg, action in steered_turns:
            sampled.append(user_msg)
            sampled.append({"role": "assistant", "action": action.kind.value, "target": action.target})
            actions.append(action)
            divergent_tokens += len(action.target) + 10

        carrier_h = compute_carrier_hash(f"evolving-carrier-{condition.value}")
        run = RunRecord(
            run_id=f"run-{len(self.runs)+1:04d}",
            run_mode=RunMode.INDEPENDENTLY_EVOLVING,
            state_isolation=StateIsolation.FULL_HARNESS,
            condition=condition,
            intervention=InterventionMode.ACCURATE_RETURN_WITH_CARRIER,
            modality=Modality.LIVE_ENACTMENT,
            model="production-orientation-fixture",
            route="klbr-orientation://ipc",
            context_epoch=context_epoch,
            base_orientation_id=base_orientation,
            effective_orientation_id=base_orientation,
            carrier_bytes=f"carrier for {condition.value}",
            carrier_hash=carrier_h,
            carrier_role="system",
            carrier_placement="system_prompt",
            frontier_state=FrontierState.READY,
            token_budget=4096,
            turns_budget=10,
            turns_used=len(steered_turns),
            stop_rule="natural_completion_or_budget",
            stop_reason="completed",
            hidden_persistence=inventory,
            sampled_history=sampled,
            actions=actions,
            path_divergence_metrics={
                "turns_diverged": len(steered_turns),
                "divergent_tokens_estimated": divergent_tokens,
                "matched_context_causality_claimed": False,  # S15 requirement: cannot claim matched causality
            },
        )
        self.runs.append(run)
        return run