#!/usr/bin/env python3
"""
memory_bridge.py

Bridge between:
  - Helix CDC / Superglyph agentic cells (ephemeral execution)
  - Echo-box memory_index (Pinecone or local storage)

Pattern:
  1) Pull tiny plan/genome (seed or .hxz)
  2) Run agentic cell ephemerally (cell handles regrow internally)
  3) Write memory to storage keyed by interaction_sha256
  4) Discard weights (cell cleanup handles this)

Integration:
  - Uses echo-box memory_index (get_memory_backend())
  - Supports both Pinecone (cloud) and LocalFibPi (on-device)
  - Keys memories by interaction_sha256 (not identity tokens)

You can call this from:
  - helix_cdc/agentic/swarm.py
  - BloomOS gates
  - A CLI wrapper
"""

import os
import sys
import json
import hashlib
import pathlib
import time
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

# Add echo-box ops to path for memory_index
ECHO_BOX_PATH = pathlib.Path.home() / "repos" / "echo-box" / "ops"
if ECHO_BOX_PATH.exists():
    sys.path.insert(0, str(ECHO_BOX_PATH))


# ---------- MEMORY INDEX INTEGRATION ----------

def get_memory_backend():
    """
    Get memory backend from echo-box (supports Pinecone or local storage).

    Set MEMORY_BACKEND env var:
      - "local" or "fibpi": On-device numpy arrays (default)
      - "pinecone": Cloud Pinecone API
    """
    try:
        from memory_index import get_memory_backend as _get_backend
        return _get_backend()
    except ImportError:
        raise RuntimeError(
            "memory_index not found. Ensure echo-box is in your path:\n"
            "  export PYTHONPATH=$HOME/repos/echo-box/ops:$PYTHONPATH"
        )


def interaction_hash(prompt: str, response: str) -> str:
    """
    Generate interaction_sha256 (memory key).

    This is the SAME function used in echo-box T-cells - ensures consistency.
    """
    payload = {"prompt": prompt, "response": response}
    b = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
    return hashlib.sha256(b).hexdigest()


def fibpi3d_embedding(text: str, dim: int = 4) -> List[float]:
    """
    Generate FibPi3D embedding (same as echo-box).

    Uses golden ratio and Fibonacci sequence for spatial positioning.
    """
    phi = 1.618033988749895  # golden ratio
    fib = [1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144]

    # Hash text
    h = int(hashlib.sha256(text.encode()).hexdigest(), 16)

    # Generate coordinates using Fibonacci spiral
    coords = []
    for i in range(dim):
        angle = (h >> (i * 16)) & 0xFFFF
        radius = fib[i % len(fib)] * phi
        coord = (angle / 65536.0) * radius
        coords.append(coord)

    # Normalize vector
    magnitude = sum(c * c for c in coords) ** 0.5
    if magnitude > 0:
        coords = [c / magnitude for c in coords]

    return coords


# ---------- AGENTIC CELL + MEMORY WRITE ----------

def _import_superglyph_cell():
    """
    Lazy import of SuperglyphAgenticCell.
    """
    from helix_cdc.agentic.cell import SuperglyphAgenticCell  # type: ignore
    return SuperglyphAgenticCell


def write_memory_to_index(
    memory_backend,
    prompt: str,
    response: str,
    task_type: str,
    plan_id: str,
    namespace: str = "thoughts"
) -> str:
    """
    Write a memory to the backend (Pinecone or local).

    Args:
        memory_backend: MemoryIndex from get_memory_backend()
        prompt: Task prompt
        response: Cell response
        task_type: Task type (code, instruct, reasoning)
        plan_id: Plan identifier
        namespace: Memory namespace (default: "thoughts")

    Returns:
        interaction_sha256 (memory key)
    """
    # Generate memory key
    memory_id = interaction_hash(prompt, response)

    # Generate FibPi3D embedding
    embedding = fibpi3d_embedding(f"{prompt} {response}", dim=4)

    # Prepare metadata
    now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    metadata = {
        "prompt": prompt,
        "candidate_output": response,
        "task_type": task_type,
        "plan_id": plan_id,
        "timestamp_utc": now,
        "source": "helix_cdc_ephemeral_cell",
    }

    # Upsert to backend
    result = memory_backend.upsert_batch(
        [{"id": memory_id, "vector": embedding, "metadata": metadata}],
        namespace=namespace
    )

    if result.get("success", 0) > 0:
        print(f"✅ Memory stored: {memory_id[:16]}... (namespace={namespace})")
    else:
        print(f"⚠️  Memory storage failed: {result}")

    return memory_id


def run_ephemeral_plan_with_memory(
    genome_path: pathlib.Path,
    plan_id: str,
    task_type: str,
    task_data: str,
    memory_backend: Optional[Any] = None,
    namespace: str = "thoughts",
) -> Dict[str, Any]:
    """
    High-level orchestrator for "pull tiny plan, regrow, discard + memory":

      1) Instantiate SuperglyphAgenticCell with genome
      2) Run autonomous_task (cell selects phenotype and executes ephemerally)
      3) Write result into memory backend (Pinecone or local) keyed by interaction_sha256
      4) Cell cleanup discards weights automatically

    Note: SuperglyphAgenticCell handles ephemeral regrow internally via
    express_phenotype() and regrow() methods.

    Args:
        genome_path: Path to superglyph genome (.hxz with multiple model seeds)
        plan_id: Logical plan identifier (e.g., 'math.v1', 'itinerary.2025.11.25')
        task_type: Task type: 'code', 'instruct', 'reasoning'
        task_data: Task data (text input for the agentic cell)
        memory_backend: MemoryIndex from get_memory_backend() (or None to skip storage)
        namespace: Memory namespace (default: "thoughts")

    Returns:
      {
        "plan_id": ...,
        "task_type": ...,
        "prompt": <task_data>,
        "response": <cell output>,
        "memory_id": <interaction_sha256 or None>,
        "genome_path": <genome path>,
      }
    """
    # 1) Get memory backend if not provided
    if memory_backend is None:
        try:
            memory_backend = get_memory_backend()
        except RuntimeError:
            print("⚠️  Memory backend not available - skipping storage")
            memory_backend = None

    # 2) Agentic cell (takes genome_path only)
    SuperglyphAgenticCell = _import_superglyph_cell()
    cell = SuperglyphAgenticCell(genome_path=str(genome_path))

    # 3) Run autonomous task (cell selects phenotype based on task_type)
    print(f"🧬 Running ephemeral cell: task_type={task_type}, plan={plan_id}")
    result = cell.autonomous_task(task_type=task_type, task_data=task_data)

    # Extract response text
    response_text = str(result.get("output", result))

    # 4) Long-term memory storage
    memory_id = None
    if memory_backend is not None:
        memory_id = write_memory_to_index(
            memory_backend,
            prompt=task_data,
            response=response_text,
            task_type=task_type,
            plan_id=plan_id,
            namespace=namespace
        )

    return {
        "plan_id": plan_id,
        "task_type": task_type,
        "prompt": task_data,
        "response": response_text,
        "memory_id": memory_id,
        "genome_path": str(genome_path),
    }


# ---------- CLI ENTRYPOINT ----------

def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        description="Run ephemeral Superglyph plan + store memory (Pinecone or local)."
    )
    parser.add_argument(
        "--genome",
        required=True,
        help="Path to superglyph genome (.hxz with multiple model seeds).",
    )
    parser.add_argument(
        "--plan-id",
        required=True,
        help="Logical plan identifier (e.g., 'math.v1', 'itinerary.2025.11.25').",
    )
    parser.add_argument(
        "--task-type",
        required=True,
        help="Task type: 'code', 'instruct', 'reasoning'",
    )
    parser.add_argument(
        "--task-data",
        required=True,
        help="Task data (text input for the agentic cell).",
    )
    parser.add_argument(
        "--namespace",
        default="thoughts",
        help="Memory namespace (default: thoughts).",
    )
    parser.add_argument(
        "--skip-memory",
        action="store_true",
        help="Skip memory storage (test cell only).",
    )
    args = parser.parse_args()

    print("=" * 70)
    print("EPHEMERAL SUPERGLYPH PLAN + MEMORY")
    print("=" * 70)
    print(f"Genome: {args.genome}")
    print(f"Plan: {args.plan_id}")
    print(f"Task type: {args.task_type}")
    print(f"Memory backend: {os.environ.get('MEMORY_BACKEND', 'local')}")
    print("=" * 70)

    # Get memory backend (or None if skipped)
    memory_backend = None if args.skip_memory else get_memory_backend()

    result = run_ephemeral_plan_with_memory(
        genome_path=pathlib.Path(args.genome),
        plan_id=args.plan_id,
        task_type=args.task_type,
        task_data=args.task_data,
        memory_backend=memory_backend,
        namespace=args.namespace,
    )

    print("\n" + "=" * 70)
    print("RESULT")
    print("=" * 70)
    print(json.dumps(result, indent=2, default=str))


if __name__ == "__main__":
    main()
