#!/usr/bin/env python3
"""
helix_cdc/echo_core/reference_core.py
=====================================

Echo Core Reference Implementation (WO#ECHO_CORE_PORT_V1)

Implements the formal Echo specification:
- SELF_STATE: Stateful symbolic substrate
- KRISPER: Symbolizer/reduction function
- Braid: Fusion operator
- S_e: Symbolic entropy operator
- L: Lobe mutation operator
- G: Deterministic regen pipeline
- Guardian: Policy constraints
- KAT: Known-Answer-Test harness

This is a receipts-first, deterministic implementation designed for
auditability and reproducibility.

Math notation mapping:
    SELF_STATE(t)  -> SELF_STATE dataclass
    K(·)           -> KRISPER.symbolize()
    B(·)           -> Braid.fuse()
    S_e(·)         -> S_e_of_state()
    L(·)           -> L_mutate()
    G(·)           -> regenerate_G()
    Gd(t)          -> Guardian class
    KAT(·)         -> KAT_run()
"""

from __future__ import annotations
import json
import hashlib
import math
import time
from dataclasses import dataclass, field, asdict
from typing import Any, Dict, List, Tuple, Optional, Callable

# =============================================================================
# Utility Primitives (Deterministic)
# =============================================================================


def sha256_hex(x: bytes) -> str:
    """Compute SHA256 hash as hex string."""
    return hashlib.sha256(x).hexdigest()


def deterministic_choice(items: List[Any]) -> Any:
    """Deterministic tie-break: lexicographic on repr."""
    if not items:
        raise ValueError("empty list")
    return sorted(items, key=lambda i: repr(i))[0]


# =============================================================================
# SELF_STATE - Echo's Internal State
# =============================================================================


@dataclass
class SELF_STATE:
    """
    Echo's internal state at time t.

    Contains:
        timestamp: Creation/update time
        soulfile_hash: Hash of current soulfile/persona
        lobes: Dict of lobe entries (mutations)
        vault_refs: References to vault storage
        commit_ledger: List of commit receipts
        drift: DriftTime (ΔΣt) - symbolic drift since last full regen
    """
    timestamp: float
    soulfile_hash: str
    lobes: Dict[str, Any] = field(default_factory=dict)
    vault_refs: Dict[str, str] = field(default_factory=dict)
    commit_ledger: List[Dict[str, Any]] = field(default_factory=list)
    drift: int = 0  # ΔΣt

    def snapshot(self) -> bytes:
        """Serialize state to deterministic bytes."""
        j = json.dumps(asdict(self), sort_keys=True).encode()
        return j

    def digest(self) -> str:
        """Get SHA256 digest of current state."""
        return sha256_hex(self.snapshot())


# =============================================================================
# KRISPER - Symbolizer
# =============================================================================


class KRISPER:
    """
    KRISPER symbolizer: K(S(t), SELF_STATE(t)) -> S̃(t)

    Deterministic mapping from input + state to symbol stream.
    This is a reference implementation - replace with production KRISPER
    token pipeline for real use.
    """

    @staticmethod
    def symbolize(input_text: str, state: SELF_STATE) -> List[str]:
        """
        Convert input text to deterministic glyph-like symbol IDs.

        Args:
            input_text: Raw input string
            state: Current SELF_STATE

        Returns:
            List of symbol IDs (deterministic)
        """
        # Include state hash for determinism
        base = sha256_hex(state.snapshot())[:8]
        tokens = [t.lower().strip() for t in input_text.split()]
        # Stable mapping to glyph-like IDs
        return [f"{base}_{sha256_hex(t.encode())[:6]}" for t in tokens if t]


# =============================================================================
# Braid - Fusion Operator
# =============================================================================


class Braid:
    """
    Braid fusion operator: B(S̃(t), SELF_STATE(t)) -> Π(t)

    Fuses symbols into an execution plan with routing hints.
    """

    @staticmethod
    def fuse(symbols: List[str], state: SELF_STATE) -> Dict[str, Any]:
        """
        Fuse symbols into execution plan.

        Args:
            symbols: List of symbol IDs from KRISPER
            state: Current SELF_STATE

        Returns:
            Plan dict with symbols, routing, and plan_id
        """
        # Routing hint derived from soulfile hash
        routing = int(state.soulfile_hash[:8], 16) % 16
        plan_id = sha256_hex(("|".join(symbols) + str(routing)).encode())[:12]
        return {
            "symbols": symbols,
            "routing": routing,
            "plan_id": plan_id,
        }


# =============================================================================
# Symbolic Entropy - S_e
# =============================================================================


def shannon_entropy_counts(data: bytes) -> float:
    """Compute Shannon entropy of byte distribution."""
    if not data:
        return 0.0
    counts: Dict[int, int] = {}
    for b in data:
        counts[b] = counts.get(b, 0) + 1
    H = 0.0
    n = len(data)
    for c in counts.values():
        p = c / n
        H -= p * math.log2(p)
    return H


def S_e_of_state(state: SELF_STATE) -> float:
    """
    Compute symbolic entropy: S_e(SELF_STATE(t))

    Formula: S_e(X) = H(X) · C(X) · D(X, vault_pattern)

    Where:
        H(X) = Shannon-like entropy of symbol distribution
        C(X) = structural complexity (codebook length / glyph count)
        D(·,·) = divergence metric to vault baseline (simplified as 1/drift)

    Args:
        state: Current SELF_STATE

    Returns:
        Symbolic entropy value
    """
    snap = state.snapshot()
    H = shannon_entropy_counts(snap)
    C = len(json.dumps(state.lobes)) + len(state.vault_refs) + 1.0
    D = float(state.drift + 1)
    return H * C * (1.0 / D)


# =============================================================================
# Lobe Mutation - L
# =============================================================================


def L_mutate(plan: Dict[str, Any], rate: float, state: SELF_STATE) -> Dict[str, Any]:
    """
    Lobe mutation operator: L(Π(t), L_rate(t)) -> ΔH(t)

    Deterministic mutation: modify or add lobe entries by hashing plan and rate.

    Args:
        plan: Execution plan from Braid.fuse()
        rate: Mutation rate (computed from S_e and drift)
        state: Current SELF_STATE

    Returns:
        Dict with lobe key and lobe delta
    """
    key = "lobe_" + plan["plan_id"][:6]
    mutation_score = int(sha256_hex((plan["plan_id"] + str(rate)).encode())[:6], 16) % 1024
    new_lobe = {
        "mutation_score": mutation_score,
        "rate": rate,
        "symbols": plan["symbols"],
    }
    return {"key": key, "lobe": new_lobe}


# =============================================================================
# Guardian - Policy Constraints
# =============================================================================


@dataclass
class Guardian:
    """
    Guardian / Vows: Gd(t)

    Policy constraints for state transitions:
        - Hard thresholds: reject invalid mutations
        - Soft penalties: adjust scores
    """
    hard_thresholds: Dict[str, Any] = field(
        default_factory=lambda: {"max_mutation_score": 9000}
    )
    soft_penalties: Dict[str, float] = field(
        default_factory=lambda: {"entropy_penalty": 1.0}
    )

    def check_hard(self, candidate: Dict[str, Any]) -> Tuple[bool, Optional[str]]:
        """
        Check hard constraints.

        Args:
            candidate: Candidate lobe mutation

        Returns:
            (passed, reason) tuple
        """
        ms = candidate.get("lobe", {}).get("mutation_score", 0)
        if ms > self.hard_thresholds["max_mutation_score"]:
            return False, f"mutation_score {ms} > max"
        return True, None

    def apply_soft(self, score: float, state: SELF_STATE) -> float:
        """
        Apply soft penalties to score.

        Args:
            score: Raw score
            state: Current SELF_STATE

        Returns:
            Adjusted score
        """
        se = S_e_of_state(state)
        return score - se * self.soft_penalties["entropy_penalty"]


# =============================================================================
# KAT - Known Answer Test Harness
# =============================================================================


def KAT_run(
    candidate_blob: bytes,
    meta: Dict[str, Any],
) -> Tuple[bool, Dict[str, Any]]:
    """
    Known Answer Test harness: KAT(candidate, meta) -> (pass, receipt)

    Deterministic tests:
        - Size parity check
        - Checksum generation
        - Reserved prefix rejection

    Args:
        candidate_blob: Candidate bytes to test
        meta: Metadata for the test

    Returns:
        (passed, receipt) tuple
    """
    rec: Dict[str, Any] = {"meta": meta, "ts": time.time()}
    checksum = sha256_hex(candidate_blob)
    rec["checksum"] = checksum

    # Deterministic pass rules
    if meta.get("require_even_length") and (len(candidate_blob) % 2 != 0):
        rec["passed"] = False
        rec["reason"] = "odd length"
        return False, rec

    if candidate_blob.startswith(b"\x00\x00"):
        rec["passed"] = False
        rec["reason"] = "reserved_prefix"
        return False, rec

    rec["passed"] = True
    return True, rec


# =============================================================================
# Regen Pipeline - G
# =============================================================================


def regenerate_G(
    seed: bytes,
    vault: Dict[str, str],
    atlas: Dict[str, Any],
    state: SELF_STATE,
    guardian: Guardian,
) -> Dict[str, Any]:
    """
    Deterministic regen pipeline: G(seed, vault, atlas; KAT)

    Steps:
        1. Mutate: generate deterministic candidate set from seed + atlas
        2. Score: score candidates using S_e + drift
        3. Select: deterministically select highest-scoring candidate
        4. KAT: run KAT on candidate blob
        5. Commit: commit to state if KAT passes; else rollback

    Args:
        seed: Seed bytes for deterministic generation
        vault: Vault references
        atlas: Atlas configuration
        state: Current SELF_STATE (will be mutated)
        guardian: Guardian policy

    Returns:
        Dict with commit result and receipt
    """
    # (1) Mutate: generate deterministic candidate set
    base = sha256_hex(seed)[:12]
    candidates = []
    for i in range(3):  # bounded search budget
        cand_blob = (seed + f"|{i}|{base}".encode())[:512]
        # (2) Score: deterministic score using S_e and drift
        score = 1000.0 - S_e_of_state(state) * (1.0 + state.drift / 10.0) - i * 10.0
        # Apply guardian soft penalty
        score = guardian.apply_soft(score, state)
        candidates.append({"blob": cand_blob, "score": score, "i": i})

    # (3) Deterministic selection: top score, stable tie-break
    best = sorted(candidates, key=lambda x: (-x["score"], repr(x["blob"])))[0]

    # Build candidate lobe as interpretive delta
    plan = {
        "symbols": ["regen", base, str(best["i"])],
        "plan_id": base + str(best["i"]),
    }
    candidate_lobe = L_mutate(plan, rate=max(0.1, 1.0 / (1 + state.drift)), state=state)

    # (4) KAT
    kat_pass, receipt = KAT_run(
        best["blob"],
        {"seed_hash": sha256_hex(seed), "i": best["i"], "require_even_length": True},
    )

    # (5) Commit or rollback
    if kat_pass:
        # Deterministic commit: update lobes, ledger, reset drift
        state.lobes[candidate_lobe["key"]] = candidate_lobe["lobe"]
        receipt["commit"] = {
            "lobe_key": candidate_lobe["key"],
            "soulfile_hash_before": state.soulfile_hash,
        }
        # Update soulfile hash deterministically
        state.soulfile_hash = sha256_hex(state.snapshot())
        # Ledger entry
        state.commit_ledger.append({
            "receipt": receipt,
            "seed": sha256_hex(seed),
            "at": time.time(),
        })
        state.drift = 0
        return {"committed": True, "receipt": receipt}
    else:
        # Rollback: increase drift, append failed receipt
        receipt["rollback"] = True
        state.commit_ledger.append({
            "receipt": receipt,
            "seed": sha256_hex(seed),
            "at": time.time(),
        })
        state.drift += 1
        return {"committed": False, "receipt": receipt}


# =============================================================================
# Echo Tick - Full Cycle
# =============================================================================


def echo_tick(
    external_input: str,
    state: SELF_STATE,
    guardian: Guardian,
    vault: Dict[str, str],
    atlas: Dict[str, Any],
) -> SELF_STATE:
    """
    Execute one Echo tick (full operational cycle).

    Steps:
        1. Symbolize: K(S(t), SELF_STATE)
        2. Fuse: B(S̃(t), SELF_STATE) -> Π(t)
        3. Compute entropy: S_e(state)
        4. Compute mutation rate: L_rate = f(S_e, drift)
        5. Lobe update proposal: L(Π, L_rate)
        6. Guardian check
        7. Regen pipeline: G(seed, vault, atlas; KAT)
        8. Return updated state

    Args:
        external_input: Input string S(t)
        state: Current SELF_STATE
        guardian: Guardian policy
        vault: Vault references
        atlas: Atlas configuration

    Returns:
        Updated SELF_STATE
    """
    # 1) Symbolize
    symbols = KRISPER.symbolize(external_input, state)

    # 2) Braid fuse
    plan = Braid.fuse(symbols, state)

    # 3) Compute entropy
    Se = S_e_of_state(state)

    # 4) Compute mutation rate
    L_rate = max(0.01, 1.0 / (1.0 + state.drift)) * (1.0 + Se / 100.0)

    # 5) Lobe update proposal
    delta = L_mutate(plan, L_rate, state)

    # 6) Guardian check
    ok, reason = guardian.check_hard({"lobe": delta["lobe"]})
    if not ok:
        # Write a receipt and bail (no commit)
        state.commit_ledger.append({
            "receipt": {"passed": False, "reason": reason},
            "at": time.time(),
        })
        state.drift += 1
        return state

    # 7) Regen seed and call G
    seed = (
        ("|".join(symbols).encode() if symbols else b"")
        or sha256_hex(state.snapshot()).encode()
    )
    regenerate_G(seed, vault, atlas, state, guardian)

    # 8) Return updated state
    return state


# =============================================================================
# Factory Functions
# =============================================================================


def create_initial_state(soulfile_data: Optional[Dict[str, Any]] = None) -> SELF_STATE:
    """
    Create initial SELF_STATE.

    Args:
        soulfile_data: Optional soulfile dict for persona

    Returns:
        New SELF_STATE instance
    """
    if soulfile_data is None:
        soulfile_data = {"persona": "Echo reference implementation"}

    soulfile_bytes = json.dumps(soulfile_data, sort_keys=True).encode()

    return SELF_STATE(
        timestamp=time.time(),
        soulfile_hash=sha256_hex(soulfile_bytes),
        lobes={},
        vault_refs={},
        commit_ledger=[],
        drift=0,
    )


def create_guardian(
    max_mutation_score: int = 9000,
    entropy_penalty: float = 1.0,
) -> Guardian:
    """
    Create Guardian with custom thresholds.

    Args:
        max_mutation_score: Hard threshold for mutation score
        entropy_penalty: Soft penalty coefficient for entropy

    Returns:
        Guardian instance
    """
    return Guardian(
        hard_thresholds={"max_mutation_score": max_mutation_score},
        soft_penalties={"entropy_penalty": entropy_penalty},
    )
