from typing import Any, Dict, List, Optional
from pydantic import Field
from langchain_core.retrievers import BaseRetriever
from langchain_core.documents import Document
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_igpt._igpt_search_api_wrapper import IgptSearchAPIWrapper
import warnings

class IgptSearchRetriever(BaseRetriever):
    """
    iGPT Search API retriever.

    Provides retrieval over the user's connected personal datasources,
    including emails, messages, files, attachments, and metadata.


    Setup:
        Install `langchain-igpt` and set environment variables `IGPT_API_KEY`, `IGPT_API_USER`.

        ```bash
        pip install -U langchain-igpt
        export IGPT_API_KEY="your-api-key"
        export IGPT_API_USER="your-api-user"
        ```

    Key init args:
        date_from:(string, optional)
            Start date filter ('YYYY-MM-DD')
        date_to (string, optional)
            End date filter ('YYYY-MM-DD').
        max_results: (number, optional)
            Limit number of results (e.g., 50)

    Instantiate:
        ```python
        from langchain_igpt import IgptSearchRetriever
        retriever = IgptSearchRetriever(max_results=50)
        ```

    Example:
        ```python
        docs = retriever.invoke("contract renewal date")
        ```

    """

    api_wrapper: IgptSearchAPIWrapper = Field(
        default_factory=IgptSearchAPIWrapper
    )

    # Search params (max_results, date filters, etc.)
    params: Dict[str, Any] = Field(default_factory=dict)

    def __init__(self, user: Optional[str] = None, api_key: Optional[str] = None, base_url: Optional[str] = None, max_retries: int = 3, backoff_factor: int = 2, backoff_base: int = 100, **kwargs: Any) -> None:
        data = {}
        if api_key or base_url or user:
            wrapper_kwargs = {}
            if api_key:
                wrapper_kwargs["api_key"] = api_key
            if base_url:
                wrapper_kwargs["base_url"] = base_url.rstrip("/")
            if user:
                wrapper_kwargs["user"] = user
                
            wrapper_kwargs["max_retries"] = max(1, max_retries) if isinstance(max_retries, int) else 3
            wrapper_kwargs["backoff_factor"] = max(2, backoff_factor) if isinstance(backoff_factor, int) else 2

            backoff_base_ms = max(100, backoff_base) if isinstance(backoff_base, int) else 100
            wrapper_kwargs["backoff_base"] = backoff_base_ms / 1000.0 # 0.1 seconds
            data["api_wrapper"] = IgptSearchAPIWrapper(**wrapper_kwargs)

        data["params"] = dict(kwargs or {})

        super().__init__(**data)

    def _get_relevant_documents(
        self,
        query: str,
        *, # This forces all following arguments to be keyword-only
        run_manager=CallbackManagerForRetrieverRun, # This is LangChain’s callback tracing system
        user: Optional[str] = None,
        **kwargs: Any,
    ) -> List[Document]:
        params = dict(self.params)
        params.update(kwargs)
        params["query"] = query

        try:
            response = self.api_wrapper.search(user=user, **params)
            if response.get("context") and not response["context"].get("datasources"):
                warnings.warn(
                    "No datasources found. Please authorize and connect a new datasource at "
                    "https://igpt.ai/hub/playground/connectors-authorize/"
                )

            return self._parse_response(response)
        except Exception as e:
            return {"error": e}

    # Async retrieval
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        user: Optional[str] = None,
        **kwargs: Any,
    ) -> List[Document]:
        params = dict(self.params)
        params.update(kwargs)
        params["query"] = query

        try:
            response = await self.api_wrapper.search_async(user=user, **params)
            if response.get("context") and not response["context"].get("datasources"):
                warnings.warn(
                    "No datasources found. Please authorize and connect a new datasource at "
                    "https://igpt.ai/hub/playground/connectors-authorize/"
                )
            return self._parse_response(response)
        except Exception as e:
            return {"error": e}


    # Documents mapping
    def _parse_response(self, response: Dict[str, Any]) -> List[Document]:
        """
        Convert iGPT search response into LangChain Documents.

        Preserves:
        - id
        - timestamp
        - type
        - from or name
        - subject
        - source (normalized)
        """

        results = response.get("results", [])
        documents: List[Document] = []

        for item in results:
            content = item.get("content") or ""
            metadata = {k: v for k, v in item.items() if k != "content"} 
            if not metadata.get("source"):
                metadata["source"] = self._build_source(item)

            documents.append(
                Document(
                    page_content=content,
                    metadata=metadata,
                )
            )

        return documents
    
    def _build_source(self, item: Dict[str, Any]) -> str:
        """
        Build a stable, human-readable source identifier.
        """
        item_type = str(item.get("type") or "unknown").strip()
        item_id = str(item.get("id") or "").strip()
        return f"{item_type}:{item_id}" if item_id else item_type