"""LangChain-compatible retriever for RagSaaS."""
from typing import Optional, List, Any, Dict
from ragsaas.config import get_config, RagSaaSConfig
from ragsaas.client import RagSaaSClient

# Try to import LangChain, but don't fail if not installed
try:
    from langchain_core.retrievers import BaseRetriever
    from langchain_core.documents import Document as LCDocument
    from langchain_core.callbacks import CallbackManagerForRetrieverRun
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    BaseRetriever = object
    LCDocument = None
    CallbackManagerForRetrieverRun = None


class RagSaaSRetriever(BaseRetriever if LANGCHAIN_AVAILABLE else object):
    """
    LangChain-compatible retriever for RagSaaS.

    Usage:
        import os
        os.environ["RAGSAAS_API_KEY"] = "sk_xxx"

        from ragsaas import RagSaaSRetriever

        # tenant_id is your customer's identifier (required)
        retriever = RagSaaSRetriever(tenant_id="customer-123")
        docs = retriever.invoke("What is machine learning?")

    With options:
        retriever = RagSaaSRetriever(
            tenant_id="customer-123",
            top_k=10,
            use_hybrid=True,
            collection="support-docs",
        )
    """

    # Pydantic fields for LangChain
    api_key: Optional[str] = None
    tenant_id: Optional[str] = None
    base_url: str = "https://api.ragsaas.com"
    collection: Optional[str] = None
    top_k: int = 5
    threshold: Optional[float] = None
    use_hybrid: bool = False

    _client: Optional[RagSaaSClient] = None

    class Config:
        arbitrary_types_allowed = True
        extra = "allow"

    def __init__(
        self,
        tenant_id: str,  # Required - this is your customer's ID
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        collection: Optional[str] = None,
        top_k: Optional[int] = None,
        threshold: Optional[float] = None,
        use_hybrid: Optional[bool] = None,
        **kwargs,
    ):
        """Initialize retriever.

        Args:
            tenant_id: Required - your customer's identifier
            api_key: Optional - defaults to RAGSAAS_API_KEY env var
            base_url: Optional - defaults to RAGSAAS_BASE_URL or https://api.ragsaas.com
            collection: Optional - filter by collection name
            top_k: Optional - number of results (default: 5)
            threshold: Optional - minimum score threshold
            use_hybrid: Optional - use hybrid search (default: False)
        """
        # Get config from environment, with overrides
        config = get_config(
            tenant_id=tenant_id,
            api_key=api_key,
            base_url=base_url,
            collection=collection,
            top_k=top_k,
            threshold=threshold,
            use_hybrid=use_hybrid,
        )

        if LANGCHAIN_AVAILABLE:
            super().__init__(
                api_key=config.api_key,
                tenant_id=config.tenant_id,
                base_url=config.base_url,
                collection=config.collection,
                top_k=config.top_k,
                threshold=config.threshold,
                use_hybrid=config.use_hybrid,
                **kwargs,
            )
        else:
            self.api_key = config.api_key
            self.tenant_id = config.tenant_id
            self.base_url = config.base_url
            self.collection = config.collection
            self.top_k = config.top_k
            self.threshold = config.threshold
            self.use_hybrid = config.use_hybrid

        self._client = RagSaaSClient(
            api_key=self.api_key,
            tenant_id=self.tenant_id,
            base_url=self.base_url,
        )

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: Any = None,
    ) -> List:
        """Get documents relevant to query (LangChain interface)."""
        if self.use_hybrid:
            results = self._client.hybrid_search(
                query=query,
                collection=self.collection,
                top_k=self.top_k,
                threshold=self.threshold,
            )
        else:
            results = self._client.search(
                query=query,
                collection=self.collection,
                top_k=self.top_k,
                threshold=self.threshold,
            )

        if not LANGCHAIN_AVAILABLE:
            return results

        return [
            LCDocument(
                page_content=doc.content,
                metadata={
                    "score": doc.score,
                    "document_id": doc.document_id,
                    "document_name": doc.document_name,
                    "chunk_index": doc.chunk_index,
                    "page_number": doc.page_number,
                    **doc.metadata,
                },
            )
            for doc in results
        ]

    # Alias for non-LangChain usage
    def retrieve(self, query: str) -> List:
        """Retrieve documents (simple interface)."""
        return self._get_relevant_documents(query)

    def invoke(self, query: str, **kwargs) -> List:
        """Invoke retriever (LangChain interface)."""
        return self._get_relevant_documents(query)

    def get_answer(
        self,
        query: str,
        model: str = "gpt-4o-mini",
        system_prompt: Optional[str] = None,
    ):
        """Get a generated answer with sources."""
        return self._client.answer(
            query=query,
            collection=self.collection,
            top_k=self.top_k,
            model=model,
            system_prompt=system_prompt,
        )


# Convenience function for quick setup
def create_retriever(**kwargs) -> RagSaaSRetriever:
    """Create a configured retriever."""
    return RagSaaSRetriever(**kwargs)
