from typing import Any, Dict, Optional, Type, Union, Literal, AsyncGenerator, Generator
from langchain_core.tools import BaseTool, ToolException
from pydantic import BaseModel, ConfigDict, Field
from langchain_igpt._igpt_ask_api_wrapper import IgptAskAPIWrapper
import warnings

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

    input: str = Field(
        description=("The question to ask")
    )

class IgptAsk(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 IgptAsk

        tool = IgptAsk(quality="cef-1-normal")
        ```

    Invoke directly with args:

        ```python
        tool.invoke({"input": "Summarize my last meeting in 5 bullet points.", "output_format":"text"})
        ```

        ```json
       {
            "id": "...",
            "output": "...",
            "context": {
                "quality": "cef-1-normal",
                "indexed": 1,
                "datasources": [{...}]
            },
            "metadata": {
                "sources": [{...},{...}]
            },
            "usage": {
                "input_tokens": 1272,
                "output_tokens": 281,
                "total_tokens": 1553
            }
        }
        ```

    """
    name: str = "igpt_ask"
    description: str = (
        "Generate a response based on the input and related context. "
        "Use the Ask tool when you need to ask a question and get an answer grounded in the user’s complete personal information archive, "
        "including emails, messages, file attachments, and related metadata. "
        "The Ask tool is ideal for finding relevant content — "
        "such as specific documents, conversation history, key decisions, dates, or background details "
        "the user previously received or discussed."
    )
    args_schema: Type[BaseModel] = IgptAskInput
    """The schema that is passed to the model when performing tool calling."""

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

    api_wrapper: IgptAskAPIWrapper = Field(default_factory=IgptAskAPIWrapper)

    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"] = IgptAskAPIWrapper(**wrapper_kwargs)

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

        super().__init__(**data)

    def _run(self, user: Optional[str] = None, **kwargs: Any) -> Union[Dict[str, Any], Generator[dict, None, None]]:
        params ={}
        params.update(self.params)
        if kwargs:
            params.update(kwargs)
        try:
            res = self.api_wrapper.ask(user=user, **params)

            if isinstance(res, dict):
                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) -> Union[Dict[str, Any], AsyncGenerator[dict, None]]:
        params ={}
        params.update(self.params)
        if kwargs:
            params.update(kwargs)
        try:
            res = await self.api_wrapper.ask_async(user=user, **params)

            if isinstance(res, dict):
                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}

