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

WO-ORGANISM-ADAPTER-01:
Bridge between FibPi+Se "brain" architecture and concrete LLM backends
(vendor Mistral, RegenMistral, TinyLlama, etc).

The architecture:
- FibPi3D + shells + Se scheduler provide:
  * Address space (WHERE in 72D torus)
  * Decision: store vs retrieve vs explore
  * Scheduling via φ/π mixing
- Backend provides:
  * generate(text, **kw) -> str
  * Optionally: logits, embeddings, etc.

We unify everything into a HelixOrganism that:
1. Takes a prompt + episode context
2. Computes FibPi address for that context
3. Queries Se scheduler for store/retrieve decision
4. Builds steered prompt (retrieved memories + original)
5. Calls backend.generate()
6. Optionally stores result back to shell memory
7. Emits machine-readable interaction record

Usage:
    from helix_cdc.agentic.organism_adapter import HelixOrganism, BackendModel
    from helix_cdc.loaders.mistral_backends import GGUFMistralBackend

    backend = GGUFMistralBackend.from_file("mistral_regen", path)
    organism = HelixOrganism(
        backend=backend,
        organism_id="mistral_regen_v01",
    )

    result = organism.run_trial(
        episode_id="test_ep",
        prompt_id="q1",
        prompt_text="What is 2+2?",
        se_mode="hot",
    )
"""

from __future__ import annotations

import time
import hashlib
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Literal, Protocol, Tuple

# Try to import real FibPi substrate
try:
    from helix_cdc.substrate.fibpi import fibpi_mix_u64, fibpi_vault_coords
    FIBPI_AVAILABLE = True
except ImportError:
    FIBPI_AVAILABLE = False

# Try to import real Se router
try:
    from helix_cdc.agentic.se_router import PhenotypeRouter, TaskSignal
    SE_ROUTER_AVAILABLE = True
except ImportError:
    SE_ROUTER_AVAILABLE = False


class BackendModel(Protocol):
    """Minimal interface any model backend must implement."""
    name: str

    def generate(self, prompt: str, max_tokens: int = 128, temperature: float = 0.0) -> str:
        """Generate completion for the given prompt."""
        ...


@dataclass
class EpisodeContext:
    """Context for a single interaction episode."""
    episode_id: str
    prompt_id: str
    fibpi_addr: Tuple[float, float, float]
    se_value: float
    notes: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ShellMemory:
    """
    Simple in-memory shell store for memories.

    In production this would be backed by crystal_vault or a real DB.
    For the organism adapter, we just need store/retrieve semantics.
    """
    memories: Dict[str, List[str]] = field(default_factory=dict)

    def _addr_key(self, addr: Tuple[float, float, float]) -> str:
        """Convert address to bucket key (quantize to regions)."""
        # Quantize to 8x8x8 grid
        x, y, z = addr
        bx = int(x * 8) % 8
        by = int(y * 8) % 8
        bz = int(z * 8) % 8
        return f"{bx}_{by}_{bz}"

    def store(self, addr: Tuple[float, float, float], content: str) -> None:
        """Store content at address."""
        key = self._addr_key(addr)
        if key not in self.memories:
            self.memories[key] = []
        # Keep last 10 per bucket
        self.memories[key].append(content)
        if len(self.memories[key]) > 10:
            self.memories[key] = self.memories[key][-10:]

    def query(self, addr: Tuple[float, float, float], k: int = 3) -> List[str]:
        """Retrieve k nearest memories from address region."""
        key = self._addr_key(addr)
        if key not in self.memories:
            return []
        return self.memories[key][-k:]

    def count(self) -> int:
        return sum(len(v) for v in self.memories.values())


@dataclass
class HelixOrganism:
    """
    Wraps a backend model with FibPi+Se steering.

    Responsibilities:
    - Map (episode, prompt) -> FibPi address
    - Decide when to retrieve vs store based on Se
    - Build a steered prompt (retrieved memories + original text)
    - Call backend.generate()
    - Emit a machine-readable interaction record
    """
    backend: BackendModel
    organism_id: str
    shell_memory: ShellMemory = field(default_factory=ShellMemory)
    se_store_threshold: float = 0.6
    se_retrieve_threshold: float = 0.4

    def _fibpi_address_for(self, episode_id: str, prompt_id: str) -> Tuple[float, float, float]:
        """
        Compute FibPi3D address for this interaction.

        Uses real fibpi_mix_u64 if available, otherwise deterministic hash.
        """
        # Create deterministic seed from organism + episode + prompt
        seed = f"{self.organism_id}:{episode_id}:{prompt_id}"
        seed_hash = hashlib.sha256(seed.encode()).digest()

        if FIBPI_AVAILABLE:
            # Use real FibPi mixing
            idx = int.from_bytes(seed_hash[:8], 'little')
            mix = fibpi_mix_u64(idx)
            x = (mix & 0xFFFF) / 65535.0
            y = ((mix >> 16) & 0xFFFF) / 65535.0
            z = ((mix >> 32) & 0xFFFF) / 65535.0
        else:
            # Fallback: simple hash-based coords
            x = int.from_bytes(seed_hash[0:4], 'little') / (2**32 - 1)
            y = int.from_bytes(seed_hash[4:8], 'little') / (2**32 - 1)
            z = int.from_bytes(seed_hash[8:12], 'little') / (2**32 - 1)

        return (x, y, z)

    def _compute_se_value(self, se_mode: Literal["cold", "hot"], task_type: str = "general") -> float:
        """
        Compute Se (symbolic entropy) value for this interaction.

        se_mode="cold": Low Se -> minimal retrieval/storage (1D-like baseline)
        se_mode="hot": High Se -> enable retrieval/storage (full brain mode)

        In production this would use PhenotypeRouter's learned scheduler.
        """
        if SE_ROUTER_AVAILABLE:
            # Use real Se router to estimate
            try:
                router = PhenotypeRouter()
                signal = TaskSignal.from_task_data(
                    task_id=hash(self.organism_id) % 10000,
                    task_type=task_type,
                    task_data=None,
                )
                # Approximate Se from signal entropy
                import numpy as np
                signal_entropy = -np.sum(np.abs(signal.input_signal) *
                                          np.log(np.abs(signal.input_signal) + 1e-10)) / len(signal.input_signal)
                se = min(1.0, max(0.0, signal_entropy / 2.0))
            except Exception:
                se = 0.2 if se_mode == "cold" else 0.8
        else:
            # Fallback: fixed values based on mode
            se = 0.2 if se_mode == "cold" else 0.8

        return se

    def _retrieve_memories(self, addr: Tuple[float, float, float], k: int = 3) -> List[str]:
        """Retrieve memories from shell at address."""
        return self.shell_memory.query(addr, k=k)

    def _store_memory(self, addr: Tuple[float, float, float], content: str) -> None:
        """Store content to shell at address."""
        self.shell_memory.store(addr, content)

    def run_trial(
        self,
        episode_id: str,
        prompt_id: str,
        prompt_text: str,
        se_mode: Literal["cold", "hot"],
        task_type: str = "general",
        max_tokens: int = 128,
        base_temperature: float = 0.0,
    ) -> Dict[str, Any]:
        """
        One FibPi+Se-steered interaction.

        Args:
            episode_id: Identifier for this episode/session
            prompt_id: Identifier for this specific prompt
            prompt_text: The actual prompt text
            se_mode: "cold" = no retrieval/minimal store, "hot" = full brain mode
            task_type: Type of task (affects Se routing)
            max_tokens: Max tokens to generate
            base_temperature: Base temperature for generation

        Returns:
            Machine-readable record of the interaction
        """
        addr = self._fibpi_address_for(episode_id, prompt_id)
        se_value = self._compute_se_value(se_mode, task_type)

        # Retrieve memories if Se is high enough
        retrieved: List[str] = []
        if se_mode == "hot" and se_value >= self.se_retrieve_threshold:
            retrieved = self._retrieve_memories(addr, k=3)

        # Build steered prompt
        steering_prefix = ""
        if retrieved:
            steering_prefix = (
                "[[CONTEXT FROM MEMORY]]\n" +
                "\n---\n".join(retrieved) +
                "\n[[END CONTEXT]]\n\n"
            )
        steered_prompt = steering_prefix + prompt_text

        # Optionally adjust temperature by Se (higher Se -> more exploration)
        temp = base_temperature + 0.3 * (se_value - 0.5)
        temp = max(0.0, min(1.0, temp))

        # Call backend
        t0 = time.time()
        try:
            completion = self.backend.generate(
                steered_prompt,
                max_tokens=max_tokens,
                temperature=temp,
            )
            error = None
        except Exception as e:
            completion = f"[ERROR: {str(e)}]"
            error = str(e)
        dt = time.time() - t0

        # Decide whether to store this outcome
        stored = False
        if se_mode == "hot" and se_value >= self.se_store_threshold and error is None:
            self._store_memory(addr, completion)
            stored = True

        return {
            "organism_id": self.organism_id,
            "backend_name": getattr(self.backend, "name", "unknown"),
            "episode_id": episode_id,
            "prompt_id": prompt_id,
            "se_mode": se_mode,
            "se_value": float(se_value),
            "fibpi_addr": list(addr),
            "retrieved_count": len(retrieved),
            "retrieved_texts": retrieved,
            "prompt": prompt_text,
            "steered_prompt": steered_prompt,
            "completion": completion,
            "stored": stored,
            "latency_ms": dt * 1000.0,
            "temperature": temp,
            "error": error,
            "shell_memory_count": self.shell_memory.count(),
        }

    def reset_memory(self) -> None:
        """Clear shell memory."""
        self.shell_memory = ShellMemory()


# =============================================================================
# Stub backends for testing
# =============================================================================

@dataclass
class EchoBackend:
    """Simple echo backend for testing - returns the prompt."""
    name: str = "echo"

    def generate(self, prompt: str, max_tokens: int = 128, temperature: float = 0.0) -> str:
        return f"[ECHO] {prompt[:max_tokens]}"


@dataclass
class StubBackend:
    """Stub backend that returns fixed responses."""
    name: str = "stub"
    responses: Dict[str, str] = field(default_factory=dict)
    default_response: str = "I don't know."

    def generate(self, prompt: str, max_tokens: int = 128, temperature: float = 0.0) -> str:
        # Check if any key matches in prompt
        for key, response in self.responses.items():
            if key.lower() in prompt.lower():
                return response[:max_tokens]
        return self.default_response


# =============================================================================
# Module test
# =============================================================================

def _test_organism():
    """Quick test of organism adapter."""
    print("Testing HelixOrganism with EchoBackend...")

    backend = EchoBackend()
    organism = HelixOrganism(
        backend=backend,
        organism_id="test_organism",
    )

    # Cold mode trial
    result1 = organism.run_trial(
        episode_id="test_ep",
        prompt_id="q1",
        prompt_text="What is 2+2?",
        se_mode="cold",
    )
    print(f"Cold trial: se={result1['se_value']:.2f}, stored={result1['stored']}")

    # Hot mode trial
    result2 = organism.run_trial(
        episode_id="test_ep",
        prompt_id="q2",
        prompt_text="What is the capital of France?",
        se_mode="hot",
    )
    print(f"Hot trial: se={result2['se_value']:.2f}, stored={result2['stored']}, "
          f"shell_count={result2['shell_memory_count']}")

    # Another hot trial should retrieve
    result3 = organism.run_trial(
        episode_id="test_ep",
        prompt_id="q3",
        prompt_text="Tell me more about France",
        se_mode="hot",
    )
    print(f"Hot trial 2: retrieved={result3['retrieved_count']}, "
          f"shell_count={result3['shell_memory_count']}")

    print("PASS: HelixOrganism basic test")
    return True


if __name__ == "__main__":
    _test_organism()
