"""Embedding providers for semantic section search.

Supports Gemini (text-embedding-004), OpenAI (text-embedding-3-small),
and sentence-transformers (fully offline, no API key required).

Auto-detection priority (first available wins):
    1. JDOCMUNCH_EMBEDDING_PROVIDER env var (gemini/openai/sentence-transformers/none)
    2. GOOGLE_API_KEY → Gemini
    3. OPENAI_API_KEY → OpenAI
    4. sentence-transformers installed → local offline model

Set JDOCMUNCH_EMBEDDING_PROVIDER=none to disable all embedding.
"""

import math
import os
from typing import Optional


# ---------------------------------------------------------------------------
# Text preparation
# ---------------------------------------------------------------------------

def _section_embed_text(section) -> str:
    """Build the text to embed for a section.

    Prepends title so short-titled sections (e.g. "Emotional Consequences"
    followed by a bullet list) still get a semantically rich embedding.
    """
    parts = [section.title]
    if section.summary and section.summary != section.title:
        parts.append(section.summary)
    if section.content:
        parts.append(section.content[:1000])
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# Cosine similarity (pure Python — no numpy dependency)
# ---------------------------------------------------------------------------

def cosine_similarity(a: list, b: list) -> float:
    """Cosine similarity between two float vectors."""
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0
    return dot / (norm_a * norm_b)


# ---------------------------------------------------------------------------
# Provider detection
# ---------------------------------------------------------------------------

def _sentence_transformers_available() -> bool:
    """Return True if sentence-transformers is importable."""
    try:
        import sentence_transformers  # noqa: F401
        return True
    except ImportError:
        return False


def get_provider_name() -> Optional[str]:
    """Return the active provider name, or None if embeddings are disabled."""
    explicit = os.environ.get("JDOCMUNCH_EMBEDDING_PROVIDER", "").lower().strip()
    if explicit == "gemini":
        return "gemini"
    if explicit == "openai":
        return "openai"
    if explicit in ("sentence-transformers", "sentence_transformers", "local"):
        return "sentence-transformers"
    if explicit == "none":
        return None
    # Auto-detect: cloud providers first, then offline fallback
    if os.environ.get("GOOGLE_API_KEY"):
        return "gemini"
    if os.environ.get("OPENAI_API_KEY"):
        return "openai"
    if _sentence_transformers_available():
        return "sentence-transformers"
    return None


# ---------------------------------------------------------------------------
# Gemini provider
# ---------------------------------------------------------------------------

class _GeminiProvider:
    """Embed via Google Gemini text-embedding-004 (768 dims)."""

    MODEL = "models/text-embedding-004"
    BATCH_SIZE = 50  # conservative to avoid rate limits

    def __init__(self):
        import google.generativeai as genai
        genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
        self._genai = genai

    def embed_texts(self, texts: list, task_type: str = "retrieval_document") -> list:
        embeddings = []
        for text in texts:
            try:
                result = self._genai.embed_content(
                    model=self.MODEL,
                    content=text,
                    task_type=task_type,
                )
                embeddings.append(result["embedding"])
            except Exception:
                embeddings.append([])
        return embeddings


# ---------------------------------------------------------------------------
# OpenAI provider
# ---------------------------------------------------------------------------

class _OpenAIProvider:
    """Embed via OpenAI text-embedding-3-small (1536 dims)."""

    MODEL = "text-embedding-3-small"
    BATCH_SIZE = 100

    def __init__(self):
        from openai import OpenAI
        self._client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

    def embed_texts(self, texts: list, task_type: str = "retrieval_document") -> list:
        # task_type is ignored for OpenAI — included for interface compatibility
        embeddings = []
        for i in range(0, len(texts), self.BATCH_SIZE):
            batch = texts[i:i + self.BATCH_SIZE]
            try:
                response = self._client.embeddings.create(model=self.MODEL, input=batch)
                embeddings.extend([e.embedding for e in response.data])
            except Exception:
                embeddings.extend([[] for _ in batch])
        return embeddings


# ---------------------------------------------------------------------------
# sentence-transformers provider (fully offline)
# ---------------------------------------------------------------------------

class _SentenceTransformersProvider:
    """Embed via sentence-transformers (all-MiniLM-L6-v2 by default, 384 dims).

    Runs entirely offline — no API key required. Install with:
        pip install sentence-transformers
    Override the model with JDOCMUNCH_ST_MODEL env var.
    """

    DEFAULT_MODEL = "all-MiniLM-L6-v2"
    BATCH_SIZE = 64

    def __init__(self):
        from sentence_transformers import SentenceTransformer
        model_name = os.environ.get("JDOCMUNCH_ST_MODEL", self.DEFAULT_MODEL)
        self._model = SentenceTransformer(model_name)

    def embed_texts(self, texts: list, task_type: str = "retrieval_document") -> list:
        # task_type is ignored — sentence-transformers handles asymmetric search
        # via separate query/passage models when needed; for MiniLM it's symmetric.
        try:
            embeddings = self._model.encode(texts, batch_size=self.BATCH_SIZE, show_progress_bar=False)
            return [emb.tolist() for emb in embeddings]
        except Exception:
            return [[] for _ in texts]


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def _get_provider():
    name = get_provider_name()
    if name == "gemini":
        try:
            return _GeminiProvider()
        except Exception:
            return None
    if name == "openai":
        try:
            return _OpenAIProvider()
        except Exception:
            return None
    if name == "sentence-transformers":
        try:
            return _SentenceTransformersProvider()
        except Exception:
            return None
    return None


def embed_sections(sections: list) -> list:
    """Generate and attach embeddings to sections in-place.

    Mirrors the summarize_sections() interface — modifies and returns the list.
    Silently degrades (leaves embedding=None) if no provider is configured.
    """
    provider = _get_provider()
    if not provider:
        return sections

    texts = [_section_embed_text(s) for s in sections]
    try:
        embeddings = provider.embed_texts(texts, task_type="retrieval_document")
        for sec, emb in zip(sections, embeddings):
            if emb:
                sec.embedding = emb
    except Exception:
        pass  # lexical search still works

    return sections


def embed_query(query: str) -> Optional[list]:
    """Embed a search query. Returns None if no provider is configured."""
    provider = _get_provider()
    if not provider:
        return None
    try:
        results = provider.embed_texts([query], task_type="retrieval_query")
        return results[0] if results and results[0] else None
    except Exception:
        return None
