"""HTTP client for RagSaaS API."""
import httpx
from typing import Optional, List, Dict, Any
from dataclasses import dataclass


@dataclass
class Document:
    """Retrieved document chunk."""
    content: str
    score: float
    document_id: str
    document_name: str
    chunk_index: int
    page_number: Optional[int]
    metadata: Dict[str, Any]


@dataclass
class Answer:
    """Generated answer with sources."""
    answer: str
    sources: List[Dict[str, Any]]
    tokens_used: int
    latency_ms: int


class RagSaaSClient:
    """Client for RagSaaS API."""

    def __init__(
        self,
        api_key: str,
        tenant_id: str,
        base_url: str = "https://api.ragsaas.com",
        timeout: float = 30.0,
    ):
        self.api_key = api_key
        self.tenant_id = tenant_id
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._client = httpx.Client(
            base_url=self.base_url,
            headers={"X-API-Key": api_key},
            timeout=timeout,
        )

    def search(
        self,
        query: str,
        collection: Optional[str] = None,
        top_k: int = 5,
        threshold: Optional[float] = None,
    ) -> List[Document]:
        """Semantic search for relevant documents."""
        payload = {
            "tenant_id": self.tenant_id,
            "query": query,
            "top_k": top_k,
        }
        if collection:
            payload["collection"] = collection
        if threshold is not None:
            payload["threshold"] = threshold

        response = self._client.post("/v1/query", json=payload)
        response.raise_for_status()
        data = response.json()

        return [
            Document(
                content=r["content"],
                score=r["score"],
                document_id=r["document_id"],
                document_name=r["document_name"],
                chunk_index=r["chunk_index"],
                page_number=r.get("page_number"),
                metadata=r.get("metadata", {}),
            )
            for r in data["results"]
        ]

    def hybrid_search(
        self,
        query: str,
        collection: Optional[str] = None,
        top_k: int = 5,
        semantic_weight: float = 0.7,
        threshold: Optional[float] = None,
    ) -> List[Document]:
        """Hybrid search combining semantic and keyword search."""
        payload = {
            "tenant_id": self.tenant_id,
            "query": query,
            "top_k": top_k,
            "semantic_weight": semantic_weight,
        }
        if collection:
            payload["collection"] = collection
        if threshold is not None:
            payload["threshold"] = threshold

        response = self._client.post("/v1/query/hybrid", json=payload)
        response.raise_for_status()
        data = response.json()

        return [
            Document(
                content=r["content"],
                score=r["score"],
                document_id=r["document_id"],
                document_name=r["document_name"],
                chunk_index=r["chunk_index"],
                page_number=r.get("page_number"),
                metadata=r.get("metadata", {}),
            )
            for r in data["results"]
        ]

    def answer(
        self,
        query: str,
        collection: Optional[str] = None,
        top_k: int = 5,
        model: str = "gpt-4o-mini",
        system_prompt: Optional[str] = None,
    ) -> Answer:
        """Generate an answer with citations."""
        payload = {
            "tenant_id": self.tenant_id,
            "query": query,
            "top_k": top_k,
            "model": model,
            "include_sources": True,
        }
        if collection:
            payload["collection"] = collection
        if system_prompt:
            payload["system_prompt"] = system_prompt

        response = self._client.post("/v1/answer", json=payload)
        response.raise_for_status()
        data = response.json()

        return Answer(
            answer=data["answer"],
            sources=data["sources"],
            tokens_used=data["tokens_used"],
            latency_ms=data["latency_ms"],
        )

    def close(self):
        """Close the HTTP client."""
        self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()
