#!/usr/bin/env python3
"""
state.py
--------
BloomOS persistent state layer:
- ledger.jsonl: task history (append-only)
- junk segments: entropy buffers persisted to disk
- reflex anchors: stored to disk for recovery

Builds on top of reflex.py and identity.py, adding file persistence.
"""

from __future__ import annotations

import os
import json
import time
import secrets
import hashlib
from pathlib import Path
from dataclasses import dataclass, asdict
from typing import Dict, Any, Optional, List

# Import existing modules
from .reflex import ReflexAnchor, insert_reflex_anchor, store_anchor
from .identity import SoulAnchor, create_soul_anchor, soul_id

# State root directory
STATE_ROOT = Path(os.getenv("BLOOMOS_STATE_ROOT", "/tmp/bloomos_state"))
LEDGER_PATH = STATE_ROOT / "ledger.jsonl"
JUNK_DIR = STATE_ROOT / "junk"
ANCHOR_DIR = STATE_ROOT / "anchors"
CAPSULE_DIR = STATE_ROOT / "capsules"

# Ensure directories exist
for p in (STATE_ROOT, JUNK_DIR, ANCHOR_DIR, CAPSULE_DIR):
    p.mkdir(parents=True, exist_ok=True)


# ============================================================================
# PERSISTENT REFLEX ANCHORS
# ============================================================================

def create_persistent_reflex_anchor(
    task_id: str,
    context: str,
    anchor_type: str = "reflex",
    metadata: Optional[Dict[str, Any]] = None
) -> ReflexAnchor:
    """
    Create a reflex anchor and persist it to disk.

    Args:
        task_id: Unique task identifier
        context: Task context/description
        anchor_type: Type of anchor
        metadata: Additional metadata

    Returns:
        ReflexAnchor object (also stored in-memory and on disk)
    """
    # Use existing function
    anchor = insert_reflex_anchor(task_id, context, anchor_type, metadata)

    # Store in memory
    store_anchor(anchor)

    # Persist to disk
    anchor_path = ANCHOR_DIR / f"{anchor.anchor_id}.json"
    anchor_path.write_text(
        json.dumps(asdict(anchor), indent=2, default=str),
        encoding="utf-8"
    )

    return anchor


def load_anchor(anchor_id: str) -> Optional[ReflexAnchor]:
    """Load anchor from disk by ID."""
    path = ANCHOR_DIR / f"{anchor_id}.json"
    if not path.exists():
        return None

    data = json.loads(path.read_text(encoding="utf-8"))
    return ReflexAnchor(**data)


def list_anchors(task_id: Optional[str] = None) -> List[ReflexAnchor]:
    """List all anchors, optionally filtered by task_id."""
    anchors = []
    for path in ANCHOR_DIR.glob("*.json"):
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
            anchor = ReflexAnchor(**data)
            if task_id is None or anchor.task_id == task_id:
                anchors.append(anchor)
        except Exception:
            continue
    return anchors


# ============================================================================
# PERSISTENT JUNK SEGMENTS
# ============================================================================

def add_persistent_junk_segment(
    payload: bytes,
    junk_size: int = 16
) -> tuple[bytes, str]:
    """
    Wrap payload with entropy junk and persist to disk.

    Args:
        payload: Raw payload bytes
        junk_size: Size of junk segments

    Returns:
        Tuple of (wrapped_payload, junk_id)
    """
    left = secrets.token_bytes(junk_size)
    right = secrets.token_bytes(junk_size)

    # Generate junk ID
    ts = int(time.time() * 1000)
    junk_id = hashlib.sha256(left + right + str(ts).encode()).hexdigest()[:16]

    # Persist junk to disk
    (JUNK_DIR / f"{junk_id}_L.bin").write_bytes(left)
    (JUNK_DIR / f"{junk_id}_R.bin").write_bytes(right)

    # Checksum for verification
    checksum = hashlib.sha256(payload).hexdigest()[:8].encode()

    wrapped = left + payload + checksum + right
    return wrapped, junk_id


def extract_from_persistent_junk(junk_id: str, wrapped: bytes) -> Optional[bytes]:
    """
    Extract payload from junk-wrapped data using stored junk segments.

    Args:
        junk_id: Junk segment identifier
        wrapped: Wrapped payload

    Returns:
        Original payload or None if verification fails
    """
    left_path = JUNK_DIR / f"{junk_id}_L.bin"
    right_path = JUNK_DIR / f"{junk_id}_R.bin"

    if not left_path.exists() or not right_path.exists():
        return None

    left = left_path.read_bytes()
    right = right_path.read_bytes()

    junk_size = len(left)

    # Extract payload (remove junk + checksum)
    data = wrapped[junk_size:]  # Remove left junk
    data = data[:-(8 + junk_size)]  # Remove checksum + right junk

    return data


# ============================================================================
# LEDGER (APPEND-ONLY TASK HISTORY)
# ============================================================================

def append_ledger_entry(entry: Dict[str, Any]) -> None:
    """
    Append a single JSON entry to the ledger.

    Args:
        entry: Dict to append (will add timestamp if missing)
    """
    entry = dict(entry)
    entry.setdefault("ts", time.time())
    entry.setdefault("source", "bloomos")

    LEDGER_PATH.parent.mkdir(parents=True, exist_ok=True)
    with LEDGER_PATH.open("a", encoding="utf-8") as f:
        f.write(json.dumps(entry, sort_keys=True, default=str) + "\n")


def read_ledger(
    limit: Optional[int] = None,
    task_type: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    Read entries from the ledger.

    Args:
        limit: Maximum entries to return (None = all)
        task_type: Filter by task_type

    Returns:
        List of ledger entries
    """
    if not LEDGER_PATH.exists():
        return []

    entries = []
    with LEDGER_PATH.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                entry = json.loads(line)
                if task_type is None or entry.get("task_type") == task_type:
                    entries.append(entry)
            except Exception:
                continue

    if limit:
        entries = entries[-limit:]

    return entries


def ledger_stats() -> Dict[str, Any]:
    """Get statistics about the ledger."""
    entries = read_ledger()

    if not entries:
        return {"total": 0, "task_types": {}, "gates": {}}

    task_types: Dict[str, int] = {}
    gates: Dict[str, int] = {}

    for e in entries:
        tt = e.get("task_type", "unknown")
        task_types[tt] = task_types.get(tt, 0) + 1

        gate = e.get("gate_used") or e.get("gate", "unknown")
        gates[gate] = gates.get(gate, 0) + 1

    return {
        "total": len(entries),
        "task_types": task_types,
        "gates": gates,
        "first_ts": entries[0].get("ts") if entries else None,
        "last_ts": entries[-1].get("ts") if entries else None,
    }


# ============================================================================
# CAPSULE PERSISTENCE
# ============================================================================

def save_capsule(capsule_json: str, capsule_id: str) -> Path:
    """
    Save a .sym capsule to disk.

    Args:
        capsule_json: JSON-serialized capsule
        capsule_id: Unique capsule identifier

    Returns:
        Path to saved file
    """
    path = CAPSULE_DIR / f"{capsule_id}.sym"
    path.write_text(capsule_json, encoding="utf-8")
    return path


def load_capsule(capsule_id: str) -> Optional[str]:
    """Load capsule JSON from disk."""
    path = CAPSULE_DIR / f"{capsule_id}.sym"
    if not path.exists():
        return None
    return path.read_text(encoding="utf-8")


def list_capsules() -> List[str]:
    """List all capsule IDs."""
    return [p.stem for p in CAPSULE_DIR.glob("*.sym")]


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

def _demo():
    """Demonstrate persistent state layer."""
    print("=" * 70)
    print("BLOOMOS PERSISTENT STATE DEMO")
    print("=" * 70)
    print()

    print(f"STATE_ROOT: {STATE_ROOT}")
    print()

    # Create reflex anchor
    anchor = create_persistent_reflex_anchor(
        task_id="demo_task_001",
        context="solve 2x + 5 = 13",
        metadata={"demo": True}
    )
    print(f"REFLEX ANCHOR:")
    print(f"  ID: {anchor.anchor_id[:16]}...")
    print(f"  Saved to: {ANCHOR_DIR}/{anchor.anchor_id}.json")
    print()

    # Add junk segment
    payload = b"Hello, BloomOS!"
    wrapped, junk_id = add_persistent_junk_segment(payload)
    print(f"JUNK SEGMENT:")
    print(f"  Original: {len(payload)} bytes")
    print(f"  Wrapped:  {len(wrapped)} bytes")
    print(f"  Junk ID:  {junk_id}")
    print()

    # Append ledger entry
    append_ledger_entry({
        "task_id": "demo_task_001",
        "task_type": "math",
        "description": "solve 2x + 5 = 13",
        "gate_used": "math_bias",
        "logits_mean": 0.75,
        "reflex_anchor_id": anchor.anchor_id,
    })
    print("LEDGER ENTRY: Appended")
    print()

    # Show ledger stats
    stats = ledger_stats()
    print(f"LEDGER STATS:")
    print(f"  Total entries: {stats['total']}")
    print(f"  Task types: {stats['task_types']}")
    print(f"  Gates: {stats['gates']}")
    print()

    print("=" * 70)


if __name__ == "__main__":
    _demo()
