import logging
from typing import Optional

import httpx

from config import settings

logger = logging.getLogger("rag-service.embedding")

_TIMEOUT = httpx.Timeout(timeout=120.0, connect=10.0)


class EmbeddingClient:
    """HTTP client for the embedding-service (port 8087)."""

    def __init__(self) -> None:
        self._base_url: str = settings.EMBEDDING_SERVICE_URL.rstrip("/")

    def _url(self, path: str) -> str:
        return f"{self._base_url}{path}"

    # ------------------------------------------------------------------
    # Embeddings
    # ------------------------------------------------------------------

    async def generate_embeddings(self, texts: list[str]) -> list[list[float]]:
        """
        Send a batch of texts to the embedding service and return a list of
        embedding vectors.
        """
        async with httpx.AsyncClient(timeout=_TIMEOUT) as client:
            response = await client.post(
                self._url("/api/v1/embeddings"),
                json={"texts": texts},
            )
            response.raise_for_status()
            data = response.json()
            return data.get("embeddings", [])

    async def generate_single_embedding(self, text: str) -> list[float]:
        """Convenience wrapper for a single text."""
        results = await self.generate_embeddings([text])
        if not results:
            raise ValueError("Embedding service returned empty result")
        return results[0]

    # ------------------------------------------------------------------
    # Reranking
    # ------------------------------------------------------------------

    async def rerank_documents(
        self,
        query: str,
        documents: list[str],
        top_k: int = 10,
    ) -> list[dict]:
        """
        Ask the embedding service to re-rank documents for a given query.
        Returns a list of {index, score, text}.
        """
        async with httpx.AsyncClient(timeout=_TIMEOUT) as client:
            response = await client.post(
                self._url("/api/v1/rerank"),
                json={
                    "query": query,
                    "documents": documents,
                    "top_k": top_k,
                },
            )
            response.raise_for_status()
            data = response.json()
            return data.get("results", [])

    # ------------------------------------------------------------------
    # Chunking
    # ------------------------------------------------------------------

    async def chunk_text(
        self,
        text: str,
        strategy: str = "recursive",
        chunk_size: int = 512,
        overlap: int = 50,
    ) -> list[str]:
        """
        Ask the embedding service to chunk a long text.
        Returns a list of chunk strings.
        """
        async with httpx.AsyncClient(timeout=_TIMEOUT) as client:
            response = await client.post(
                self._url("/api/v1/chunk"),
                json={
                    "text": text,
                    "strategy": strategy,
                    "chunk_size": chunk_size,
                    "overlap": overlap,
                },
            )
            response.raise_for_status()
            data = response.json()
            return data.get("chunks", [])

    # ------------------------------------------------------------------
    # PDF extraction
    # ------------------------------------------------------------------

    async def extract_pdf(self, pdf_bytes: bytes) -> str:
        """
        Send raw PDF bytes to the embedding service for text extraction.
        Returns the extracted text as a string.
        """
        async with httpx.AsyncClient(timeout=_TIMEOUT) as client:
            response = await client.post(
                self._url("/api/v1/extract-pdf"),
                files={"file": ("document.pdf", pdf_bytes, "application/pdf")},
            )
            response.raise_for_status()
            data = response.json()
            return data.get("text", "")


# Singleton
embedding_client = EmbeddingClient()