#!/usr/bin/env python3
"""
helix_substrate_core.py
-----------------------
Core data structures for Helix as a symbolic substrate.

This is NOT the whole engine.
It's the typed "contract" layer that mirrors HELIX_SUBSTRATE_SPEC.md:

- HelixCapsule: living model bundle
- HelixState: Σ_s, B, E
- GuardianVector / Vows
- RegenPatch + RegenReceiptRef
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


# --- Core symbolic objects ----------------------------------------------------


@dataclass
class EntropyScaffold:
    """Tracks symbolic entropy + DriftTime for a capsule."""
    drift_time: int              # ΔΣt
    symbolic_entropy: float      # Sₑ
    forks: int = 0               # number of forks/branches
    mutation_pressure: float = 0.0


@dataclass
class GuardianVector:
    """
    Guardian + vows.

    Hard constraints: forbid certain mutations / regions.
    Soft constraints: influence routing / reasoning bands.
    """
    hard_forbidden_keys: List[str] = field(default_factory=list)
    soft_biases: Dict[str, float] = field(default_factory=dict)
    vows_contract_id: Optional[str] = None  # link to vows doc / policy


@dataclass
class RegenReceiptRef:
    """
    Reference to a regen receipt stored elsewhere (JSONL, sym, etc.).
    """
    receipt_id: str
    path: str
    kat_pass: bool
    delta_S_e: float
    delta_drift: int


@dataclass
class LobeBindings:
    """
    Echo lobe ↔ Helix capsule bindings.

    This describes how Echo lobes read/write state into the capsule.
    """
    read_lobes: List[str] = field(default_factory=list)
    write_lobes: List[str] = field(default_factory=list)
    # Optional: mapping from lobe name → capsule region key
    mapping: Dict[str, str] = field(default_factory=dict)


@dataclass
class BraidMap:
    """
    Braid map: how attention heads / layers are mapped to FibPi3D ports.

    In practice this will mirror whatever your existing fibpi/flowtorch
    code is already doing; this is just the typed shape.
    """
    ports: List[int]                      # e.g. [1, 5, 8, 13, ...] Fibonacci ports
    head_to_port: Dict[str, int]          # "layer.head" → port id
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class HelixCapsule:
    """
    Helix capsule: a 'living model'.

    Contains:
      - core weights
      - braid map
      - entropy scaffold
      - guardian vector + vows
      - regen receipt refs
      - lobe bindings (Echo integration)
    """
    capsule_id: str
    model_path: str              # path to GGUF / weights
    quantization: str            # e.g. "Q5_K_M"
    braid_map: BraidMap
    entropy: EntropyScaffold
    guardian: GuardianVector
    regen_history: List[RegenReceiptRef] = field(default_factory=list)
    lobe_bindings: LobeBindings = field(default_factory=LobeBindings)
    extra: Dict[str, Any] = field(default_factory=dict)


@dataclass
class HelixState:
    """
    Symbolic state at runtime for a capsule.
    Σ_s (symbol state), B (braid operator), E (entropy lattice).
    """
    capsule: HelixCapsule
    sigma_s: Dict[str, Any] = field(default_factory=dict)  # symbol state vector
    braid_runtime_state: Dict[str, Any] = field(default_factory=dict)
    entropy_runtime_state: Dict[str, Any] = field(default_factory=dict)


@dataclass
class RegenPatch:
    """
    Proposed change to a Helix capsule.
    This is what your regen engine 'G' produces.

    It SHOULD be small and structured, not a full model.
    """
    patch_id: str
    target_capsule_id: str
    changes: Dict[str, Any]          # e.g. {"braid_map": {...}, "guardian": {...}}
    proposed_delta_drift: int = 0
    proposed_delta_S_e: float = 0.0


# --- API surface for Claude / higher layers ----------------------------------


def apply_regen_patch(
    state: HelixState,
    patch: RegenPatch,
    kat_result: Optional[Dict[str, Any]] = None,
) -> HelixState:
    """
    Apply a regen patch to a HelixState, assuming KAT has already run.

    - If kat_result['kat_pass'] is False, this function should NOT be called.
    - In real code, this MUST be atomic + receipt-logged.
    - Here, we only update fields matching patch.changes keys.
    """
    if patch.target_capsule_id != state.capsule.capsule_id:
        raise ValueError("Patch target_capsule_id does not match state capsule_id")

    # Mutate a shallow copy to keep semantics clear
    capsule = state.capsule

    changes = patch.changes or {}

    # Example: update braid_map, guardian, entropy
    if "braid_map" in changes:
        # Expect dict structure compatible with BraidMap
        bm_data = changes["braid_map"]
        capsule.braid_map = BraidMap(
            ports=bm_data.get("ports", capsule.braid_map.ports),
            head_to_port=bm_data.get("head_to_port", capsule.braid_map.head_to_port),
            metadata={**capsule.braid_map.metadata, **bm_data.get("metadata", {})},
        )

    if "guardian" in changes:
        g_data = changes["guardian"]
        capsule.guardian.hard_forbidden_keys = g_data.get(
            "hard_forbidden_keys", capsule.guardian.hard_forbidden_keys
        )
        capsule.guardian.soft_biases.update(g_data.get("soft_biases", {}))
        if "vows_contract_id" in g_data:
            capsule.guardian.vows_contract_id = g_data["vows_contract_id"]

    if "entropy" in changes:
        e_data = changes["entropy"]
        capsule.entropy.drift_time += e_data.get("delta_drift", 0)
        capsule.entropy.symbolic_entropy += e_data.get("delta_S_e", 0.0)
        capsule.entropy.forks += e_data.get("delta_forks", 0)
        capsule.entropy.mutation_pressure += e_data.get("delta_mutation_pressure", 0.0)

    # Record receipt reference if provided
    if kat_result:
        rec = RegenReceiptRef(
            receipt_id=str(kat_result.get("receipt_id", patch.patch_id)),
            path=str(kat_result.get("receipt_path", "")),
            kat_pass=bool(kat_result.get("kat_pass", True)),
            delta_S_e=float(kat_result.get("delta_S_e", patch.proposed_delta_S_e)),
            delta_drift=int(kat_result.get("delta_drift", patch.proposed_delta_drift)),
        )
        capsule.regen_history.append(rec)

    # Return updated state (same object, but explicit for clarity)
    return state


def create_empty_capsule(
    capsule_id: str,
    model_path: str,
    quantization: str = "Q5_K_M",
) -> HelixCapsule:
    """
    Factory for a minimal HelixCapsule with default structures.
    """
    return HelixCapsule(
        capsule_id=capsule_id,
        model_path=model_path,
        quantization=quantization,
        braid_map=BraidMap(ports=[], head_to_port={}),
        entropy=EntropyScaffold(drift_time=0, symbolic_entropy=0.0),
        guardian=GuardianVector(),
        regen_history=[],
        lobe_bindings=LobeBindings(),
        extra={},
    )


def create_state(capsule: HelixCapsule) -> HelixState:
    """
    Factory for a HelixState from a capsule.
    """
    return HelixState(
        capsule=capsule,
        sigma_s={},
        braid_runtime_state={},
        entropy_runtime_state={},
    )
