#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

"""Embedding generation helpers.

Supports multiple providers. The customer supplies their own API key
for client-side generation, or uses Veculo's server-side embedding
endpoint (billed separately).
"""

from __future__ import annotations
from typing import Optional
from veculo.errors import VeculoError


class EmbeddingProvider:
    """Base class for embedding providers."""

    def embed(self, text: str) -> list[float]:
        raise NotImplementedError

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        return [self.embed(t) for t in texts]


class OpenAIEmbeddings(EmbeddingProvider):
    """Generate embeddings using OpenAI's API.

    Requires: pip install openai

    Example::
        embedder = OpenAIEmbeddings(api_key="sk-...")
        vector = embedder.embed("Hello world")
    """

    def __init__(
        self,
        api_key: str | None = None,
        model: str = "text-embedding-3-small",
    ):
        try:
            import openai
        except ImportError:
            raise VeculoError(
                "OpenAI support requires the 'openai' package. "
                "Install it with: pip install 'veculo[openai]'"
            )
        import os
        self._client = openai.OpenAI(api_key=api_key or os.environ.get("OPENAI_API_KEY"))
        self._model = model

    def embed(self, text: str) -> list[float]:
        response = self._client.embeddings.create(input=text, model=self._model)
        return response.data[0].embedding

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        response = self._client.embeddings.create(input=texts, model=self._model)
        return [d.embedding for d in sorted(response.data, key=lambda d: d.index)]


class VertexAIEmbeddings(EmbeddingProvider):
    """Generate embeddings using Google Vertex AI.

    Requires: pip install google-cloud-aiplatform
    Uses Application Default Credentials or explicit credentials.

    Example::
        embedder = VertexAIEmbeddings(project="my-project")
        vector = embedder.embed("Hello world")
    """

    def __init__(
        self,
        project: str | None = None,
        location: str = "us-central1",
        model: str = "text-embedding-005",
    ):
        try:
            from vertexai.language_models import TextEmbeddingModel
            import vertexai
        except ImportError:
            raise VeculoError(
                "Vertex AI support requires 'google-cloud-aiplatform'. "
                "Install it with: pip install 'veculo[vertexai]'"
            )
        import os
        project = project or os.environ.get("GOOGLE_CLOUD_PROJECT")
        vertexai.init(project=project, location=location)
        self._model = TextEmbeddingModel.from_pretrained(model)

    def embed(self, text: str) -> list[float]:
        embeddings = self._model.get_embeddings([text])
        return embeddings[0].values

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        embeddings = self._model.get_embeddings(texts)
        return [e.values for e in embeddings]


class SentenceTransformerEmbeddings(EmbeddingProvider):
    """Generate embeddings using local sentence-transformers models.

    Requires: pip install sentence-transformers
    No API key needed — runs locally.

    Example::
        embedder = SentenceTransformerEmbeddings()  # defaults to all-MiniLM-L6-v2
        vector = embedder.embed("Hello world")
    """

    def __init__(self, model: str = "all-MiniLM-L6-v2"):
        try:
            from sentence_transformers import SentenceTransformer
        except ImportError:
            raise VeculoError(
                "Local embeddings require 'sentence-transformers'. "
                "Install it with: pip install 'veculo[local]'"
            )
        self._model = SentenceTransformer(model)

    def embed(self, text: str) -> list[float]:
        return self._model.encode(text).tolist()

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        return self._model.encode(texts).tolist()
