"""
Spatial Memory Manager for Echo

Wraps the FibPi3D substrate, enabling deterministic spatial queries
(near, between, k_nearest, inside, left_of, above) without retraining weights.

Architecture:
    Echo (LLM)
        ↓ tool call
    /echo/spatial endpoint (FastAPI)
        ↓
    SpatialMemoryManager (this module)
        ↓
    FibPi3DBackend
        ↓
    Deterministic result + receipt
"""

import hashlib
import json
import math
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np

# Project root detection
PROJECT_ROOT = Path(__file__).resolve().parents[3]


# ============================================================================
# Scene Object (inline to avoid circular imports)
# ============================================================================

@dataclass
class SceneObject:
    """An object in a spatial scene."""
    name: str
    x: float
    y: float
    z: float

    @property
    def pos(self) -> np.ndarray:
        return np.array([self.x, self.y, self.z])

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "x": self.x,
            "y": self.y,
            "z": self.z
        }


# ============================================================================
# Deterministic RNG (for reproducible scene generation)
# ============================================================================

class DeterministicRNG:
    """
    Deterministic PRNG based on seed string.
    Uses SHA256 to derive state, ensuring exact reproducibility.
    """

    def __init__(self, seed: str):
        self.seed = seed
        self.counter = 0

    def _get_bytes(self, n: int) -> bytes:
        """Get n deterministic bytes."""
        chunks = []
        while len(b''.join(chunks)) < n:
            data = f"{self.seed}:{self.counter}".encode()
            h = hashlib.sha256(data).digest()
            chunks.append(h)
            self.counter += 1
        return b''.join(chunks)[:n]

    def random(self) -> float:
        """Return float in [0, 1)."""
        import struct
        b = self._get_bytes(8)
        v = struct.unpack('>Q', b)[0]
        return (v >> 11) / (1 << 53)

    def uniform(self, low: float, high: float) -> float:
        """Return float in [low, high)."""
        return low + self.random() * (high - low)

    def gauss(self, mu: float = 0.0, sigma: float = 1.0) -> float:
        """Box-Muller transform for Gaussian."""
        u1 = max(self.random(), 1e-10)
        u2 = self.random()
        z = np.sqrt(-2 * np.log(u1)) * np.cos(2 * np.pi * u2)
        return mu + sigma * z


# ============================================================================
# FibPi3D Backend (lightweight inline version)
# ============================================================================

# Golden ratio and Pi constants
PHI = (1 + math.sqrt(5)) / 2
PI = math.pi


class FibPi3DBackend:
    """
    FibPi3D lattice-based spatial queries.

    Key properties:
    1. No explicit edge storage - adjacency is implicit
    2. Between-ness is a geometric predicate, not graph traversal
    3. Same lattice geometry at all scales
    """

    def __init__(self, objects: List[SceneObject], use_original_coords: bool = True):
        """
        Initialize FibPi3D backend.

        Args:
            objects: Scene objects to process
            use_original_coords: If True, use original coordinates directly.
        """
        self.objects = objects
        self.n = len(objects)
        self.use_original_coords = use_original_coords
        self._name_to_idx = {obj.name: i for i, obj in enumerate(objects)}
        self._idx_to_name = {i: obj.name for i, obj in enumerate(objects)}
        self._positions = np.array([obj.pos for obj in objects])

    def memory_bytes(self) -> int:
        """Return memory footprint in bytes."""
        return self._positions.nbytes

    def near(self, a_name: str, b_name: str, threshold: float) -> bool:
        """Check if distance(a, b) < threshold."""
        a_idx = self._name_to_idx.get(a_name)
        b_idx = self._name_to_idx.get(b_name)
        if a_idx is None or b_idx is None:
            return False
        dist = np.linalg.norm(self._positions[a_idx] - self._positions[b_idx])
        return bool(dist < threshold)

    def distance(self, a_name: str, b_name: str) -> float:
        """Return Euclidean distance between a and b."""
        a_idx = self._name_to_idx.get(a_name)
        b_idx = self._name_to_idx.get(b_name)
        if a_idx is None or b_idx is None:
            return float('inf')
        return float(np.linalg.norm(self._positions[a_idx] - self._positions[b_idx]))

    def k_nearest(self, target_name: str, k: int) -> List[Tuple[str, float]]:
        """Find k nearest neighbors."""
        target_idx = self._name_to_idx.get(target_name)
        if target_idx is None:
            return []

        target_pos = self._positions[target_idx]
        distances = []
        for i in range(self.n):
            if i == target_idx:
                continue
            dist = float(np.linalg.norm(self._positions[i] - target_pos))
            distances.append((self._idx_to_name[i], dist))

        distances.sort(key=lambda x: x[1])
        return distances[:k]

    def between(self, a_name: str, b_name: str, c_name: str, tolerance: float = 0.3) -> bool:
        """
        Check if b is geometrically between a and c.
        This is O(1) - pure geometric predicate, no graph traversal.
        """
        a_idx = self._name_to_idx.get(a_name)
        b_idx = self._name_to_idx.get(b_name)
        c_idx = self._name_to_idx.get(c_name)

        if None in (a_idx, b_idx, c_idx):
            return False

        a = self._positions[a_idx]
        b = self._positions[b_idx]
        c = self._positions[c_idx]

        ac = c - a
        ac_len_sq = np.dot(ac, ac)
        if ac_len_sq < 1e-12:
            return False

        ab = b - a
        t = np.dot(ab, ac) / ac_len_sq

        if t <= 0.0 or t >= 1.0:
            return False

        projection = a + t * ac
        dist_to_line = np.linalg.norm(b - projection)
        return bool(dist_to_line < tolerance)

    def inside(self, a_name: str, center_name: str, radius: float) -> bool:
        """Check if a is within radius of center."""
        return self.near(a_name, center_name, radius)

    def left_of(self, a_name: str, b_name: str) -> bool:
        """Check if a.x < b.x."""
        a_idx = self._name_to_idx.get(a_name)
        b_idx = self._name_to_idx.get(b_name)
        if a_idx is None or b_idx is None:
            return False
        return bool(self._positions[a_idx][0] < self._positions[b_idx][0])

    def above(self, a_name: str, b_name: str) -> bool:
        """Check if a.z > b.z."""
        a_idx = self._name_to_idx.get(a_name)
        b_idx = self._name_to_idx.get(b_name)
        if a_idx is None or b_idx is None:
            return False
        return bool(self._positions[a_idx][2] > self._positions[b_idx][2])


# ============================================================================
# Scene Management
# ============================================================================

@dataclass
class SceneConfig:
    """Configuration for scene generation."""
    n_objects: int = 100
    distribution: str = "uniform"  # uniform, gaussian, clustered
    bounds: Tuple[float, float] = (-10.0, 10.0)
    seed: str = "spatial_memory_v1"


@dataclass
class Scene:
    """A spatial scene with objects and backend."""
    scene_id: str
    config: SceneConfig
    objects: List[SceneObject]
    backend: FibPi3DBackend
    scene_hash: str
    created_at: str


def generate_scene(config: SceneConfig) -> List[SceneObject]:
    """Generate deterministic scene objects from config."""
    rng = DeterministicRNG(f"{config.seed}:{config.n_objects}")
    objects = []

    low, high = config.bounds
    mid = (low + high) / 2
    spread = (high - low) / 4

    for i in range(config.n_objects):
        name = f"obj_{i:04d}"

        if config.distribution == "uniform":
            x = rng.uniform(low, high)
            y = rng.uniform(low, high)
            z = rng.uniform(low, high)
        elif config.distribution == "gaussian":
            x = rng.gauss(mid, spread)
            y = rng.gauss(mid, spread)
            z = rng.gauss(mid, spread)
        elif config.distribution == "clustered":
            cluster_id = i % 5
            cluster_rng = DeterministicRNG(f"{config.seed}:cluster:{cluster_id}")
            cx = cluster_rng.uniform(low, high)
            cy = cluster_rng.uniform(low, high)
            cz = cluster_rng.uniform(low, high)
            x = cx + rng.gauss(0, spread / 4)
            y = cy + rng.gauss(0, spread / 4)
            z = cz + rng.gauss(0, spread / 4)
        else:
            raise ValueError(f"Unknown distribution: {config.distribution}")

        objects.append(SceneObject(name=name, x=x, y=y, z=z))

    return objects


def compute_scene_hash(objects: List[SceneObject]) -> str:
    """Compute deterministic SHA256 hash of scene."""
    data = []
    for obj in sorted(objects, key=lambda o: o.name):
        data.append(f"{obj.name}:{obj.x:.6f}:{obj.y:.6f}:{obj.z:.6f}")
    combined = "|".join(data).encode()
    return hashlib.sha256(combined).hexdigest()


# ============================================================================
# Spatial Memory Manager
# ============================================================================

class SpatialMemoryManager:
    """
    Manages spatial scenes and executes spatial queries.

    This is the main interface for the /echo/spatial endpoint.
    Provides scene caching, query dispatch, and receipt generation.
    """

    def __init__(self, cache_dir: Optional[Path] = None):
        """
        Initialize the manager.

        Args:
            cache_dir: Optional directory for scene persistence
        """
        self._scenes: Dict[str, Scene] = {}
        self._cache_dir = cache_dir or (PROJECT_ROOT / "receipts" / "spatial_scenes")
        self._cache_dir.mkdir(parents=True, exist_ok=True)

    def list_scenes(self) -> List[str]:
        """List available scene IDs."""
        return list(self._scenes.keys())

    def create_scene(
        self,
        scene_id: str,
        n_objects: int = 100,
        distribution: str = "uniform",
        bounds: Tuple[float, float] = (-10.0, 10.0),
        seed: Optional[str] = None,
    ) -> Scene:
        """
        Create a new scene with given parameters.

        Args:
            scene_id: Unique identifier for the scene
            n_objects: Number of objects to generate
            distribution: "uniform", "gaussian", or "clustered"
            bounds: (min, max) coordinate bounds
            seed: Optional seed override (defaults to scene_id)

        Returns:
            The created Scene
        """
        config = SceneConfig(
            n_objects=n_objects,
            distribution=distribution,
            bounds=bounds,
            seed=seed or scene_id,
        )

        objects = generate_scene(config)
        backend = FibPi3DBackend(objects, use_original_coords=True)
        scene_hash = compute_scene_hash(objects)

        scene = Scene(
            scene_id=scene_id,
            config=config,
            objects=objects,
            backend=backend,
            scene_hash=scene_hash,
            created_at=datetime.now(timezone.utc).isoformat(),
        )

        self._scenes[scene_id] = scene
        return scene

    def get_scene(self, scene_id: str) -> Optional[Scene]:
        """Get a scene by ID."""
        return self._scenes.get(scene_id)

    def get_or_create_scene(
        self,
        scene_id: str,
        n_objects: int = 100,
        distribution: str = "uniform",
        bounds: Tuple[float, float] = (-10.0, 10.0),
        seed: Optional[str] = None,
    ) -> Scene:
        """Get existing scene or create if missing."""
        scene = self._scenes.get(scene_id)
        if scene is not None:
            return scene
        return self.create_scene(scene_id, n_objects, distribution, bounds, seed)

    def delete_scene(self, scene_id: str) -> bool:
        """Delete a scene. Returns True if deleted."""
        if scene_id in self._scenes:
            del self._scenes[scene_id]
            return True
        return False

    def query(
        self,
        scene_id: str,
        query_type: str,
        a: str,
        b: Optional[str] = None,
        c: Optional[str] = None,
        k: int = 5,
        threshold: float = 1.0,
        tolerance: float = 0.3,
    ) -> Dict[str, Any]:
        """
        Execute a spatial query on a scene.

        Args:
            scene_id: Scene to query
            query_type: One of: near, distance, k_nearest, between, inside, left_of, above
            a: First object name
            b: Second object name (for binary predicates)
            c: Third object name (for between)
            k: Number of neighbors (for k_nearest)
            threshold: Distance threshold (for near/inside)
            tolerance: Line distance tolerance (for between)

        Returns:
            Dict with result, latency_us, query details
        """
        scene = self._scenes.get(scene_id)
        if scene is None:
            return {
                "error": f"Scene not found: {scene_id}",
                "query_type": query_type,
                "result": None,
            }

        backend = scene.backend
        start_ns = time.perf_counter_ns()

        try:
            if query_type == "near":
                if b is None:
                    return {"error": "near requires 'b' parameter", "result": None}
                result = backend.near(a, b, threshold)

            elif query_type == "distance":
                if b is None:
                    return {"error": "distance requires 'b' parameter", "result": None}
                result = backend.distance(a, b)

            elif query_type == "k_nearest":
                neighbors = backend.k_nearest(a, k)
                result = [(name, round(dist, 6)) for name, dist in neighbors]

            elif query_type == "between":
                if b is None or c is None:
                    return {"error": "between requires 'b' and 'c' parameters", "result": None}
                result = backend.between(a, b, c, tolerance)

            elif query_type == "inside":
                if b is None:
                    return {"error": "inside requires 'b' (center) parameter", "result": None}
                result = backend.inside(a, b, threshold)

            elif query_type == "left_of":
                if b is None:
                    return {"error": "left_of requires 'b' parameter", "result": None}
                result = backend.left_of(a, b)

            elif query_type == "above":
                if b is None:
                    return {"error": "above requires 'b' parameter", "result": None}
                result = backend.above(a, b)

            else:
                return {"error": f"Unknown query type: {query_type}", "result": None}

        except Exception as e:
            return {"error": str(e), "result": None}

        end_ns = time.perf_counter_ns()
        latency_us = (end_ns - start_ns) / 1000.0

        return {
            "query_type": query_type,
            "a": a,
            "b": b,
            "c": c,
            "k": k if query_type == "k_nearest" else None,
            "threshold": threshold if query_type in ("near", "inside") else None,
            "tolerance": tolerance if query_type == "between" else None,
            "result": result,
            "latency_us": round(latency_us, 2),
        }

    def batch_query(
        self,
        scene_id: str,
        queries: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """
        Execute multiple queries on a scene.

        Args:
            scene_id: Scene to query
            queries: List of query dicts with query_type, a, b, c, k, threshold, tolerance

        Returns:
            List of query results
        """
        results = []
        for q in queries:
            result = self.query(
                scene_id=scene_id,
                query_type=q.get("query_type", ""),
                a=q.get("a", ""),
                b=q.get("b"),
                c=q.get("c"),
                k=q.get("k", 5),
                threshold=q.get("threshold", 1.0),
                tolerance=q.get("tolerance", 0.3),
            )
            results.append(result)
        return results


# ============================================================================
# Module-level singleton for convenience
# ============================================================================

_manager: Optional[SpatialMemoryManager] = None


def get_manager() -> SpatialMemoryManager:
    """Get or create the global SpatialMemoryManager."""
    global _manager
    if _manager is None:
        _manager = SpatialMemoryManager()
    return _manager


# ============================================================================
# Self-test
# ============================================================================

if __name__ == "__main__":
    print("SpatialMemoryManager Self-Test")
    print("=" * 50)

    mgr = SpatialMemoryManager()

    # Create scene
    scene = mgr.create_scene("test", n_objects=100, seed="test_seed")
    print(f"Created scene: {scene.scene_id}")
    print(f"  Objects: {len(scene.objects)}")
    print(f"  Hash: {scene.scene_hash[:16]}...")
    print(f"  Memory: {scene.backend.memory_bytes()} bytes")

    # Test queries
    print("\nQuery tests:")

    result = mgr.query("test", "near", "obj_0000", "obj_0001", threshold=5.0)
    print(f"  near(obj_0000, obj_0001, 5.0): {result['result']} ({result['latency_us']}μs)")

    result = mgr.query("test", "distance", "obj_0000", "obj_0001")
    print(f"  distance(obj_0000, obj_0001): {result['result']:.4f} ({result['latency_us']}μs)")

    result = mgr.query("test", "k_nearest", "obj_0000", k=3)
    print(f"  k_nearest(obj_0000, k=3): {result['result'][:3]} ({result['latency_us']}μs)")

    result = mgr.query("test", "between", "obj_0000", b="obj_0001", c="obj_0002", tolerance=5.0)
    print(f"  between(obj_0000, obj_0001, obj_0002): {result['result']} ({result['latency_us']}μs)")

    result = mgr.query("test", "left_of", "obj_0000", "obj_0001")
    print(f"  left_of(obj_0000, obj_0001): {result['result']} ({result['latency_us']}μs)")

    result = mgr.query("test", "above", "obj_0000", "obj_0001")
    print(f"  above(obj_0000, obj_0001): {result['result']} ({result['latency_us']}μs)")

    # Batch query
    print("\nBatch query test:")
    batch_results = mgr.batch_query("test", [
        {"query_type": "k_nearest", "a": "obj_0000", "k": 5},
        {"query_type": "near", "a": "obj_0000", "b": "obj_0001", "threshold": 10.0},
        {"query_type": "distance", "a": "obj_0000", "b": "obj_0050"},
    ])
    print(f"  Executed {len(batch_results)} queries")
    total_latency = sum(r.get("latency_us", 0) for r in batch_results)
    print(f"  Total latency: {total_latency:.2f}μs")

    print("\nAll tests passed.")
