#!/usr/bin/env python3
"""
TextEmbedder for Helix Memory Layer (WO#24).

Provides real text embeddings using:
1. sentence-transformers (preferred, if available)
2. HuggingFace AutoModel (fallback)
3. Deterministic FibPi3D hash (offline fallback)

See: ECHO_HELIX_MEMORY_LAYER_V1.md
"""

import hashlib
from typing import List, Optional, Tuple

import numpy as np


# Default model - small, fast, good quality
DEFAULT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"


def _try_sentence_transformers(model_name: str) -> Tuple[Optional[object], Optional[str]]:
    """Try loading a sentence-transformers model."""
    try:
        from sentence_transformers import SentenceTransformer
        model = SentenceTransformer(model_name)
        return model, "sentence_transformers"
    except Exception:
        return None, None


def _try_hf_model(model_name: str) -> Tuple[Optional[callable], Optional[str]]:
    """Try loading a HuggingFace model with mean pooling."""
    try:
        from transformers import AutoTokenizer, AutoModel
        import torch

        tok = AutoTokenizer.from_pretrained(model_name)
        mdl = AutoModel.from_pretrained(model_name)
        mdl.eval()

        def _encode(texts: List[str], batch_size: int = 16) -> np.ndarray:
            all_vecs = []
            with torch.no_grad():
                for i in range(0, len(texts), batch_size):
                    batch = texts[i:i + batch_size]
                    toks = tok(batch, padding=True, truncation=True, return_tensors="pt")
                    out = mdl(**toks)
                    last = out.last_hidden_state
                    mask = toks["attention_mask"].unsqueeze(-1)
                    summed = (last * mask).sum(dim=1)
                    counts = mask.sum(dim=1).clamp(min=1)
                    mean = (summed / counts).cpu().numpy()
                    all_vecs.append(mean)
            X = np.vstack(all_vecs).astype(np.float32)
            # Normalize
            norms = np.linalg.norm(X, axis=1, keepdims=True)
            X = X / np.maximum(norms, 1e-9)
            return X

        return _encode, "huggingface"
    except Exception:
        return None, None


def _fibpi3d_fallback(texts: List[str], dim: int = 384) -> np.ndarray:
    """
    Deterministic FibPi3D hash-based embedding (offline fallback).

    Uses golden ratio + Fibonacci sequence for spatial positioning.
    This is NOT semantic - just deterministic placeholder.
    """
    phi = 1.618033988749895
    fib = [1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610, 987]

    embeddings = []
    for text in texts:
        h = int(hashlib.sha256(text.encode()).hexdigest(), 16)

        coords = []
        for i in range(dim):
            # Extract bits from hash (cycle through if needed)
            shift = (i * 16) % 256
            angle = (h >> shift) & 0xFFFF
            radius = fib[i % len(fib)] * phi
            coord = (angle / 65536.0) * radius * ((-1) ** (i % 3))
            coords.append(coord)

        vec = np.array(coords, dtype=np.float32)
        # Normalize
        norm = np.linalg.norm(vec)
        if norm > 0:
            vec = vec / norm
        embeddings.append(vec)

    return np.vstack(embeddings)


class TextEmbedder:
    """
    Text embedding wrapper with multiple backends.

    Priority:
    1. sentence-transformers (best quality, fast)
    2. HuggingFace AutoModel (fallback)
    3. FibPi3D hash (deterministic offline fallback)
    """

    def __init__(self, model_name: Optional[str] = None, fallback_dim: int = 384):
        """
        Initialize TextEmbedder.

        Args:
            model_name: Model to use (default: all-MiniLM-L6-v2)
            fallback_dim: Dimension for FibPi3D fallback
        """
        self.model_name = model_name or DEFAULT_MODEL
        self.fallback_dim = fallback_dim
        self._model = None
        self._encode_fn = None
        self._backend = None
        self._dim = None

        self._init_model()

    def _init_model(self) -> None:
        """Initialize the embedding model."""
        # Try sentence-transformers first
        model, backend = _try_sentence_transformers(self.model_name)
        if model is not None:
            self._model = model
            self._backend = backend
            # Get dimension from model
            test_emb = model.encode(["test"], show_progress_bar=False)
            self._dim = test_emb.shape[1]
            return

        # Try HuggingFace
        encode_fn, backend = _try_hf_model(self.model_name)
        if encode_fn is not None:
            self._encode_fn = encode_fn
            self._backend = backend
            # Get dimension from test
            test_emb = encode_fn(["test"])
            self._dim = test_emb.shape[1]
            return

        # Fallback to FibPi3D
        self._backend = "fibpi3d_fallback"
        self._dim = self.fallback_dim

    @property
    def backend(self) -> str:
        """Return the active backend name."""
        return self._backend

    @property
    def dim(self) -> int:
        """Return the embedding dimension."""
        return self._dim

    def embed(self, text: str) -> np.ndarray:
        """
        Embed a single text.

        Args:
            text: Input text

        Returns:
            1D normalized embedding vector
        """
        return self.embed_batch([text])[0]

    def embed_batch(self, texts: List[str]) -> np.ndarray:
        """
        Embed multiple texts.

        Args:
            texts: List of input texts

        Returns:
            (N, D) array of normalized embeddings
        """
        if not texts:
            return np.zeros((0, self._dim), dtype=np.float32)

        if self._backend == "sentence_transformers":
            X = self._model.encode(
                texts,
                batch_size=64,
                normalize_embeddings=True,
                show_progress_bar=False
            )
            return np.asarray(X, dtype=np.float32)

        elif self._backend == "huggingface":
            return self._encode_fn(texts)

        else:  # fibpi3d_fallback
            return _fibpi3d_fallback(texts, dim=self._dim)

    def similarity(self, text1: str, text2: str) -> float:
        """
        Compute cosine similarity between two texts.

        Args:
            text1: First text
            text2: Second text

        Returns:
            Cosine similarity (0 to 1 for normalized vectors)
        """
        emb1 = self.embed(text1)
        emb2 = self.embed(text2)
        return float(np.dot(emb1, emb2))


if __name__ == "__main__":
    # Quick test
    print("=== TextEmbedder Test ===")

    embedder = TextEmbedder()
    print(f"Backend: {embedder.backend}")
    print(f"Dimension: {embedder.dim}")

    texts = [
        "I like cats.",
        "I love dogs.",
        "Quantum field theory is hard."
    ]

    print(f"\nEmbedding {len(texts)} texts...")
    embeddings = embedder.embed_batch(texts)
    print(f"Shape: {embeddings.shape}")

    # Test similarity
    print("\nSimilarity matrix:")
    for i, t1 in enumerate(texts):
        for j, t2 in enumerate(texts):
            sim = float(embeddings[i] @ embeddings[j])
            print(f"  [{i},{j}] {sim:.4f}", end="")
        print()

    # Query test
    query = "I enjoy puppies"
    q_emb = embedder.embed(query)
    sims = embeddings @ q_emb
    ranked = sorted(range(len(texts)), key=lambda i: -sims[i])
    print(f"\nQuery: '{query}'")
    print("Rankings:")
    for rank, idx in enumerate(ranked):
        print(f"  {rank+1}. [{sims[idx]:.4f}] {texts[idx]}")
