from typing import Any, Dict, Optional, Type
from langchain_core.tools import BaseTool, ToolException
from pydantic import BaseModel, ConfigDict, Field
from langchain_igpt._igpt_search_api_wrapper import IgptSearchAPIWrapper
import warnings

class IgptSearchInput(BaseModel):
    model_config = ConfigDict(extra="allow")

    query: Optional[str] = Field(
        default=None,
        description=("The search query to execute.")
    )

    date_from: Optional[str] = Field(
        default=None,
        description="""Filter by start date. 

        Format must be YYYY-MM-DD (e.g., "2024-01-15" for January 15, 2024).

        Examples:
        - "2025-01-02" - Results from January 2, 2025 onwards
        - "2024-12-21" - Results from December 21, 2024 onwards

        Default is None.
        """,
    )

    date_to: Optional[str] = Field(
        default=None,
        description="""Filter by end date.

        Format must be YYYY-MM-DD (e.g., "2024-03-31" for March 31, 2024).

        Examples:
        - "2024-03-31" - Results up to and including March 31, 2024
        - "2023-12-31" - Results up to and including December 31, 2023

        Default is None.
        """,
    )

class IgptSearch(BaseTool):
    """

    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"
        ```

    Instantiate:

        ```python
        from langchain_igpt import IgptSearch

        tool = IgptSearch(max_results=5)
        ```

    Invoke directly with args:

        ```python
        tool.invoke({"query": "invoices", "date_from":"2024-03-31", "date_to":"2025-01-02"})
        ```

        ```json
       {
            "id": "...",
            "context": {
                "indexed": 1,
                "datasources": [{...}]
            },
            "results": [{
                "id": "...",
                "timestamp": "2024-05-13T11:48:34Z",
                "type": "message",
                "from": "...",
                "subject": "...",
                "content": "..."
                },{...}],
            "durations": {
                "retrieve": 47,
                "resolve": 10
            }
        }
        ```

    """
    name: str = "igpt_search"
    description: str = (
        "Search in connected datasources. "
        "The Search tool provides direct, retrieval-only access to indexed context across your connected datasources, "
        "including emails, messages, file attachments, and related metadata. "
        "It returns raw, context-aware results without applying any reasoning, summarization, "
        "or generation - so you can inspect and use the underlying evidence as-is. "
        "Each result includes stable, structured metadata, making it well-suited for building citations, assembling lists, "
        "and implementing custom filtering or downstream processing pipelines."
    )
    args_schema: Type[BaseModel] = IgptSearchInput
    """The schema that is passed to the model when performing tool calling."""

    params: Optional[Dict[str, Any]] = None

    api_wrapper: IgptSearchAPIWrapper = Field(default_factory=IgptSearchAPIWrapper)

    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 _run(self, user: Optional[str] = None, **kwargs: Any) -> Dict[str, Any]:
        params ={}
        params.update(self.params)
        if kwargs:
            params.update(kwargs)

        try:
            res = self.api_wrapper.search(user=user, **params)

            if res.get("context") and not res["context"].get("datasources"):
                warnings.warn(
                    "No datasources found. Please authorize and connect a new datasource at "
                    "https://igpt.ai/hub/playground/connectors-authorize/"
                )

            return res
        except ToolException:
            # Re-raise tool exceptions
            raise
        except Exception as e:
            return {"error": e}

    async def _arun(self, user: Optional[str] = None, **kwargs: Any) -> Dict[str, Any]:
        params ={}
        params.update(self.params)
        if kwargs:
            params.update(kwargs)
        try:
            res = await self.api_wrapper.search_async(user=user, **params)
            
            if res.get("context") and not res["context"].get("datasources"):
                warnings.warn(
                    "No datasources found. Please authorize and connect a new datasource at " 
                    "https://igpt.ai/hub/playground/connectors-authorize/"
                )

            return res
        except ToolException:
            # Re-raise tool exceptions
            raise
        except Exception as e:
            return {"error": e}
