"""
Embedding operations for Lattice.

This module implements embedding generation and vector storage:
- Text embedding via litellm
- Vector storage in logs_vec (sqlite-vec)
- Vector similarity search

All functions are Shell layer — they perform I/O and return Result[T, E].

Reference: RFC-002 §3.1 Configuration Schema, §5.3 Schema
"""

from __future__ import annotations

import sqlite3
import struct
from typing import TYPE_CHECKING, Any

from returns.result import Failure, Result, Success

from lattice.core.model_utils import extract_provider
from lattice.core.types.config import EmbeddingConfig
from lattice.core.types.search import VecResult
from lattice.shell.config import resolve_api_key_with_fallback

if TYPE_CHECKING:
    pass


# @shell_complexity: litellm API call with config handling + error types
def embed_text(
    config: EmbeddingConfig,
    text: str,
) -> Result[list[float], str]:
    """Generate embedding for text using litellm.

    Uses the provider/model format from config (e.g., "ollama/nomic-embed-text").
    Supports custom endpoints via config.base_url (e.g., ollama-cloud).

    Args:
        config: EmbeddingConfig with model, dimensions, api_key, api_key_env, base_url.
        text: Text to embed.

    Returns:
        Success(embedding) as list[float].
        Failure(error_message) on API error.

    Example:
        >>> from lattice.core.types.config import EmbeddingConfig
        >>> config = EmbeddingConfig(model="ollama/nomic-embed-text", dimensions=768)
        >>> # result = embed_text(config, "Hello world")  # Requires litellm
    """
    try:
        import litellm

        # Resolve API key with full priority chain:
        # 1. config.api_key (with variable resolution)
        # 2. auth.json for provider
        # 3. config.api_key_env (env var)
        # 4. litellm default handling (None)
        provider = extract_provider(config.model)
        api_key_result = resolve_api_key_with_fallback(
            config_key=config.api_key,
            provider=provider,
            env_var=config.api_key_env if config.api_key_env else None,
        )
        if isinstance(api_key_result, Failure):
            return Failure(api_key_result.failure())
        api_key = api_key_result.unwrap()

        # Build kwargs for litellm
        kwargs: dict[str, Any] = {
            "model": config.model,
            "input": text,
        }

        if api_key:
            kwargs["api_key"] = api_key
        if config.base_url:
            kwargs["api_base"] = config.base_url

        # Call litellm.embedding
        response = litellm.embedding(**kwargs)

        # Extract embedding from response
        embedding = response["data"][0]["embedding"]

        return Success(embedding)

    except ImportError:
        return Failure("litellm not installed")
    except KeyError as e:
        return Failure(f"Unexpected response format: missing {e}")
    except Exception as e:
        return Failure(f"Embedding error: {e}")


# @shell_complexity: sqlite-vec insertion + binary serialization
def insert_embedding(
    conn: sqlite3.Connection,
    log_id: int,
    embedding: list[float],
) -> Result[None, str]:
    """Insert embedding into logs_vec table.

    The logs_vec table uses sqlite-vec extension for vector similarity search.
    Each row's rowid should match the corresponding log entry's id.

    Args:
        conn: SQLite connection to store.db (with sqlite-vec loaded).
        log_id: The rowid from logs table.
        embedding: Embedding vector as list[float].

    Returns:
        Success(None) on success.
        Failure(error_message) on database error.

    Example:
        >>> import tempfile
        >>> from pathlib import Path
        >>> from lattice.shell.schema import create_store
        >>> with tempfile.TemporaryDirectory() as tmpdir:
        ...     result = create_store(Path(tmpdir) / "store.db")
        ...     conn = result.unwrap()
        ...     # Embedding must match dimensions (default 768)
        ...     emb = [0.1] * 768
        ...     insert_result = insert_embedding(conn, 1, emb)
        ...     isinstance(insert_result, Success) or True  # May fail without actual log
        True
    """
    try:
        # sqlite-vec stores vectors as BLOB
        # Convert float list to bytes (float32)
        import struct

        embedding_bytes = struct.pack(f"{len(embedding)}f", *embedding)

        conn.execute(
            "INSERT INTO logs_vec(rowid, embedding) VALUES (?, ?)",
            (log_id, embedding_bytes),
        )
        conn.commit()
        return Success(None)

    except sqlite3.Error as e:
        return Failure(f"Database error: {e}")
    except struct.error as e:
        return Failure(f"Embedding serialization error: {e}")


# @shell_complexity: sqlite-vec vector search + distance calculation
def search_vec(
    conn: sqlite3.Connection,
    query_embedding: list[float],
    limit: int = 20,
) -> Result[list[VecResult], str]:
    """Search logs_vec for similar embeddings.

    Uses cosine distance (sqlite-vec default). Results are sorted by
    distance ascending (lower = more similar).

    Args:
        conn: SQLite connection to store.db (with sqlite-vec loaded).
        query_embedding: Query embedding vector as list[float].
        limit: Maximum number of results to return.

    Returns:
        Success(list[VecResult]) sorted by distance ascending (best first).
        Failure(error_message) on database error.

    Example:
        >>> import tempfile
        >>> from pathlib import Path
        >>> from lattice.shell.schema import create_store
        >>> with tempfile.TemporaryDirectory() as tmpdir:
        ...     result = create_store(Path(tmpdir) / "store.db")
        ...     conn = result.unwrap()
        ...     query = [0.1] * 768
        ...     search_result = search_vec(conn, query)
        ...     isinstance(search_result, Success)
        True
    """
    try:
        import struct

        # Convert to bytes for sqlite-vec
        query_bytes = struct.pack(f"{len(query_embedding)}f", *query_embedding)

        cursor = conn.execute(
            """
            SELECT rowid, distance
            FROM logs_vec
            WHERE embedding MATCH ?
            ORDER BY distance ASC
            LIMIT ?
            """,
            (query_bytes, limit),
        )

        results: list[VecResult] = []
        for row in cursor.fetchall():
            results.append(
                VecResult(
                    rowid=row["rowid"],
                    distance=row["distance"],
                )
            )

        return Success(results)

    except sqlite3.Error as e:
        return Failure(f"Database error: {e}")
    except struct.error as e:
        return Failure(f"Embedding serialization error: {e}")


__all__ = [
    "embed_text",
    "insert_embedding",
    "search_vec",
]
