"""Embedding provider for docsctx GraphRAG pipeline."""

from __future__ import annotations

import math
import struct

import httpx

from docsctx.config import DocsSettings


def encode_vector(vec: list[float]) -> bytes:
    """Pack a float vector into a compact bytes blob."""
    return struct.pack(f"<{len(vec)}f", *vec)


def decode_vector(data: bytes) -> list[float]:
    """Unpack a bytes blob into a float list."""
    count = len(data) // 4
    return list(struct.unpack(f"<{count}f", data))


def cosine_similarity(a: list[float], b: list[float]) -> float:
    """Compute cosine similarity between two vectors."""
    if len(a) != len(b) or not a:
        return 0.0
    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 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


from langchain_core.embeddings import Embeddings
from langchain_openai import OpenAIEmbeddings, AzureOpenAIEmbeddings
from langchain_ollama import OllamaEmbeddings


# Export BaseEmbedder as an alias to Embeddings for backward compatibility in type hints
BaseEmbedder = Embeddings


class HuggingFaceEmbedder(Embeddings):
    """Fallback manual embedding via HuggingFace Inference API."""

    def __init__(self, api_key: str, model: str = "sentence-transformers/all-MiniLM-L6-v2", timeout: float = 30.0) -> None:
        self.api_key = api_key
        self.model = model
        self.timeout = timeout
        self._base_url = f"https://api-inference.huggingface.co/pipeline/feature-extraction/{model}"

    def embed_query(self, text: str) -> list[float]:
        return self.embed_documents([text])[0]

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        response = httpx.post(
            self._base_url,
            headers={"Authorization": f"Bearer {self.api_key}"},
            json={"inputs": texts, "options": {"wait_for_model": True}},
            timeout=self.timeout,
        )
        response.raise_for_status()
        result = response.json()
        if isinstance(result, list) and result and isinstance(result[0], list):
            return result
        if isinstance(result, list) and result and isinstance(result[0], (int, float)):
            return [result]
        return []


def create_embedder(settings: DocsSettings) -> Embeddings | None:
    """Create an embedder from settings; returns None if not configured."""
    provider = settings.embedding.provider.lower().strip()
    model = settings.embedding.model
    if provider in {"", "none", "mock"}:
        return None
    if provider == "ollama":
        base_url = settings.embedding.base_url or settings.chat.base_url
        client_kwargs = {"timeout": settings.chat.timeout_seconds}
        api_key = settings.embedding.api_key or settings.chat.api_key
        if api_key:
            client_kwargs["headers"] = {"Authorization": f"Bearer {api_key}"}
            
        return OllamaEmbeddings(
            base_url=base_url,
            model=model,
            client_kwargs=client_kwargs,
        )
    if provider in {"openai", "azure"}:
        base_url = settings.embedding.base_url or settings.chat.base_url
        api_key = settings.embedding.api_key or settings.chat.api_key
        if not api_key:
            return None
        if provider == "azure":
            return AzureOpenAIEmbeddings(
                azure_endpoint=base_url,
                api_key=api_key,
                deployment=model,
                api_version="2024-02-15-preview",
            )
        else:
            return OpenAIEmbeddings(
                base_url=base_url,
                api_key=api_key,
                model=model,
            )
    if provider == "huggingface":
        api_key = settings.embedding.api_key or settings.chat.api_key
        if not api_key:
            return None
        return HuggingFaceEmbedder(api_key, model)
    return None


def validate_embedder_access(settings: DocsSettings, embedder: Embeddings | None = None) -> Embeddings:
    """Ensure embedding configuration exists and the backend is reachable."""
    provider = settings.embedding.provider.lower().strip()
    if provider in {"", "none"}:
        raise ValueError("DOCSCTX_EMBEDDING__PROVIDER must be set to a real embedding provider")

    # Mock provider is allowed for testing
    if provider == "mock":
        return None  # type: ignore

    resolved_embedder = embedder or create_embedder(settings)
    if resolved_embedder is None:
        raise ValueError(
            "Embedding configuration is incomplete. Set DOCSCTX_EMBEDDING__MODEL and credentials."
        )

    base_url = (settings.embedding.base_url or settings.chat.base_url or "").rstrip("/")
    if base_url:
        try:
            import httpx
            if provider == "ollama":
                headers = {}
                api_key = settings.embedding.api_key or settings.chat.api_key
                if api_key:
                    headers["Authorization"] = f"Bearer {api_key}"
                response = httpx.get(f"{base_url}/api/tags", headers=headers, timeout=5.0)
                response.raise_for_status()
            elif provider in {"openai", "azure"}:
                api_key = settings.embedding.api_key or settings.chat.api_key
                headers = {"Authorization": f"Bearer {api_key}"}
                if provider == "azure":
                    headers["api-key"] = api_key
                    url = f"{base_url}/models?api-version=2024-02-15-preview"
                else:
                    url = f"{base_url}/models"
                response = httpx.get(url, headers=headers, timeout=5.0)
                response.raise_for_status()
            elif provider == "huggingface":
                api_key = settings.embedding.api_key or settings.chat.api_key
                headers = {"Authorization": f"Bearer {api_key}"}
                response = httpx.get("https://huggingface.co/api/models", headers=headers, timeout=5.0)
                response.raise_for_status()
            else:
                response = httpx.get(base_url, timeout=5.0)
        except Exception as exc:
            raise RuntimeError(f"Embedding provider host {base_url} is not accessible: {exc}")

    return resolved_embedder


def embed_texts(
    embedder: Embeddings,
    texts: list[str],
    batch_size: int = 20,
) -> list[list[float]]:
    """Embed a list of texts in batches."""
    results: list[list[float]] = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        results.extend(embedder.embed_documents(batch))
    return results
