#!/usr/bin/env python3
"""
reflex.py
---------
JUNK segments and Reflex Anchors for BloomOS.

JUNK = contextual entropy buffers between tasks
Reflex Anchors = embedded self-checks per task run

Based on HelixDNA paper: reflex agents, self-repair, contextual alignment.
"""

from __future__ import annotations

import os
import time
import hashlib
from dataclasses import dataclass, field
from typing import Dict, Any, Optional, List


@dataclass
class ReflexAnchor:
    """Embedded self-checkpoint for a task."""
    anchor_id: str
    timestamp: float
    context: str
    task_id: str
    anchor_type: str = "reflex"
    metadata: Dict[str, Any] = field(default_factory=dict)


def sha256_hex(data: bytes) -> str:
    """Compute SHA256 hex digest."""
    return hashlib.sha256(data).hexdigest()


def insert_reflex_anchor(
    task_id: str,
    context: str,
    anchor_type: str = "reflex",
    metadata: Optional[Dict[str, Any]] = None
) -> ReflexAnchor:
    """
    Create a reflex anchor for a task.

    Reflex anchors are embedded self-checks that allow the system
    to verify task integrity and provide self-repair capabilities.

    Args:
        task_id: Unique task identifier
        context: Task context/description
        anchor_type: Type of anchor (reflex, checkpoint, soul)
        metadata: Additional metadata

    Returns:
        ReflexAnchor object
    """
    anchor_id = sha256_hex(f"{task_id}:{context}:{time.time()}".encode())

    return ReflexAnchor(
        anchor_id=anchor_id,
        timestamp=time.time(),
        context=context,
        task_id=task_id,
        anchor_type=anchor_type,
        metadata=metadata or {}
    )


def add_junk_segment(payload: bytes, junk_size: int = 32) -> bytes:
    """
    Wrap payload with JUNK entropy segments.

    JUNK = contextual buffer that provides:
    - Entropy padding for alignment
    - Boundary markers for parsing
    - Self-repair checksums

    Args:
        payload: Raw task payload
        junk_size: Size of each JUNK segment in bytes

    Returns:
        Payload wrapped with JUNK segments
    """
    junk_prefix = os.urandom(junk_size)
    junk_suffix = os.urandom(junk_size)

    # Include checksum in suffix for verification
    checksum = sha256_hex(payload)[:8].encode()

    return junk_prefix + payload + checksum + junk_suffix


def extract_from_junk(wrapped: bytes, junk_size: int = 32) -> bytes:
    """
    Extract payload from JUNK-wrapped data.

    Args:
        wrapped: JUNK-wrapped payload
        junk_size: Size of each JUNK segment

    Returns:
        Original payload
    """
    # Remove prefix JUNK
    data = wrapped[junk_size:]
    # Remove checksum (8 bytes) and suffix JUNK
    data = data[:-(8 + junk_size)]
    return data


def verify_junk_integrity(wrapped: bytes, junk_size: int = 32) -> bool:
    """
    Verify JUNK-wrapped payload integrity using embedded checksum.

    Args:
        wrapped: JUNK-wrapped payload
        junk_size: Size of each JUNK segment

    Returns:
        True if checksum matches
    """
    try:
        # Extract payload
        payload = extract_from_junk(wrapped, junk_size)

        # Extract embedded checksum
        data_with_checksum = wrapped[junk_size:-(junk_size)]
        embedded_checksum = data_with_checksum[-8:].decode()

        # Compute actual checksum
        actual_checksum = sha256_hex(payload)[:8]

        return embedded_checksum == actual_checksum

    except Exception:
        return False


# ============================================================================
# ANCHOR STORAGE
# ============================================================================

_anchor_store: Dict[str, ReflexAnchor] = {}


def store_anchor(anchor: ReflexAnchor) -> None:
    """Store anchor in memory."""
    _anchor_store[anchor.anchor_id] = anchor


def get_anchor(anchor_id: str) -> Optional[ReflexAnchor]:
    """Retrieve anchor by ID."""
    return _anchor_store.get(anchor_id)


def get_task_anchors(task_id: str) -> List[ReflexAnchor]:
    """Get all anchors for a task."""
    return [a for a in _anchor_store.values() if a.task_id == task_id]


def clear_anchors() -> None:
    """Clear all stored anchors."""
    _anchor_store.clear()


# ============================================================================
# DEMO
# ============================================================================

def _demo():
    """Demonstrate reflex anchors and JUNK segments."""
    print("=" * 70)
    print("REFLEX ANCHORS & JUNK SEGMENTS DEMO")
    print("=" * 70)
    print()

    # Create reflex anchor
    anchor = insert_reflex_anchor(
        task_id="task_001",
        context="solve 2x + 5 = 13",
        anchor_type="reflex",
        metadata={"gate": "math_bias"}
    )
    store_anchor(anchor)

    print("REFLEX ANCHOR:")
    print(f"  ID: {anchor.anchor_id[:16]}...")
    print(f"  Task: {anchor.task_id}")
    print(f"  Context: {anchor.context}")
    print(f"  Type: {anchor.anchor_type}")
    print()

    # JUNK segment demo
    payload = b"Hello, HelixDNA!"
    wrapped = add_junk_segment(payload)

    print("JUNK SEGMENTS:")
    print(f"  Original: {len(payload)} bytes")
    print(f"  Wrapped:  {len(wrapped)} bytes")
    print(f"  Overhead: {len(wrapped) - len(payload)} bytes")
    print()

    # Verify integrity
    is_valid = verify_junk_integrity(wrapped)
    extracted = extract_from_junk(wrapped)

    print("INTEGRITY CHECK:")
    print(f"  Valid: {is_valid}")
    print(f"  Extracted matches: {extracted == payload}")
    print()

    print("=" * 70)


if __name__ == "__main__":
    _demo()
