from typing import Any, Dict, Optional, AsyncGenerator, Generator, Union
import time
from pydantic import BaseModel, ConfigDict, SecretStr, model_validator
from langchain_core.utils import get_from_dict_or_env
import requests
import aiohttp
import asyncio
import json
import warnings
from importlib.metadata import version

IGPT_BASE_URL: str = "https://api.igpt.ai/v1"

class IgptAskAPIWrapper(BaseModel):
    """Wrapper for Igpt Ask API."""

    api_key: SecretStr
    base_url: Optional[str] = None
    user: Optional[str] = None
    max_retries: int = 3
    backoff_factor: int = 2
    backoff_base: float = 0.1 # 100 milliseconds
    _version = version("langchain-igpt")

    model_config = ConfigDict(
        extra="forbid",
    )

    @model_validator(mode="before")
    def validate_environment(cls, values: Dict) -> Any:
        api_key = get_from_dict_or_env(
            values, "api_key", "IGPT_API_KEY"
        )
        values["api_key"] = api_key
        user = get_from_dict_or_env(
            values, "user", "IGPT_API_USER"
        )
        values["user"] = user
        return values

    def ask(self, user: Optional[str] = None, **kwargs: Any) -> Union[Dict[str, Any], Generator[dict, None, None]]:
        payload = {}
        payload["user"] = user or self.user or None
        if kwargs:
            payload.update(kwargs)

        base_url = self.base_url or IGPT_BASE_URL

        url = f"{base_url}/recall/ask/"
        return self._stream_request(url, payload) if payload.get("stream") else self._post(url, payload)

    def _post(self, url: str, payload: dict):
        headers = {
            "Authorization": f"Bearer {self.api_key.get_secret_value()}",
            "Content-Type": "application/json",
            "X-Client": f"langchain-igpt/{self._version}"
        }
        delay = self.backoff_base
        for attempt in range(self.max_retries):
            try:
                resp = requests.post(url, json=payload, headers=headers, timeout=(60, 600)) # connect_timeout 60 seconds, read_timeout 10 minutes
                if resp.status_code == 200:
                    try:
                        return resp.json()
                    except Exception:
                        return {"error": "invalid_json_response"}
                elif resp.status_code == 429 or 500 <= resp.status_code < 600: # Server-side error (429 - Too Many Requests) - retry
                    pass
                else:
                    return {"error": "client_error"} # Client-side error - don't retry
            except requests.exceptions.RequestException: # Network error or timeout - retry
                pass

            if attempt < self.max_retries - 1: # Sleep before next retry, except after the last attempt
                time.sleep(delay)
                delay *= self.backoff_factor  # Exponential backoff

        return {"error": "network_error"}
    
    def _stream_request(self, url: str, payload: dict) -> Generator[dict, None, None]:
        headers = {
            "Authorization": f"Bearer {self.api_key.get_secret_value()}",
            "Content-Type": "application/json",
            "X-Client": f"langchain-igpt/{self._version}"
        }
        delay = self.backoff_base
        for attempt in range(self.max_retries):
            try:
                with requests.post(url, json=payload, stream=True, headers=headers, timeout=(60, 600)) as resp: 
                    if resp.status_code == 200:
                        for chunk in resp.iter_lines():
                            if chunk:
                                chunk_str = chunk.decode()
                                if chunk_str.startswith("data:"):
                                    json_str = chunk_str[len("data:"):].strip()
                                else:
                                    json_str = chunk_str

                                if not json_str:
                                    continue
                                
                                try:
                                    data = json.loads(json_str)
                                    if data.get("context") and not data["context"].get("datasources"):
                                        warnings.warn(
                                            "No datasources found. Please authorize and connect a new datasource at " 
                                            "https://igpt.ai/hub/playground/connectors-authorize/"
                                        )
                                    yield data
                                except json.JSONDecodeError:
                                    continue
                        return
                    elif resp.status_code == 429 or 500 <= resp.status_code < 600:
                        pass
                    else:
                        yield {"error": "client_error"}
                        return
            except requests.exceptions.RequestException:
                pass

            if attempt < self.max_retries - 1:
                time.sleep(delay)
                delay *= self.backoff_factor
            else:  
                yield {"error": "network_error"}

    # Async
    async def ask_async(self, user: Optional[str] = None, **kwargs: Any) -> Union[Dict[str, Any], AsyncGenerator[dict, None]]:
        payload = {}
        payload["user"] = user or self.user or None
        if kwargs:
            payload.update(kwargs)

        base_url = self.base_url or IGPT_BASE_URL

        url = f"{base_url}/recall/ask/"
        return self._stream_request_async(url, payload) if payload.get("stream") else await self._post_async(url, payload)

    async def _post_async(self, url: str, payload: dict):
        headers = {
            "Authorization": f"Bearer {self.api_key.get_secret_value()}",
            "Content-Type": "application/json",
            "X-Client": f"langchain-igpt/{self._version}"
        }

        delay = self.backoff_base
        timeout = aiohttp.ClientTimeout(
            total=None,
            connect=60,      # connect timeout (sec)
            sock_read=600,   # read timeout (sec) = 10 min
        )
        async with aiohttp.ClientSession(headers=headers, timeout=timeout) as session:
            for attempt in range(self.max_retries):
                try:
                    async with session.post(url, json=payload) as resp:
                        if resp.status == 200:
                            try:
                                return await resp.json()
                            except Exception:
                                return {"error": "invalid_json_response"}
                        elif resp.status == 429 or 500 <= resp.status < 600: # Server-side error (429 - Too Many Requests) - retry
                            pass
                        else:
                            return {"error": "client_error"} # Client-side error - don't retry
                except (aiohttp.ClientError, asyncio.TimeoutError):
                    pass

                if attempt < self.max_retries - 1: # Sleep before next retry, except after the last attempt
                    await asyncio.sleep(delay)
                    delay *= self.backoff_factor # Exponential backoff

        return {"error": "network_error"}
    
    async def _stream_request_async(self, url: str, payload: dict) -> AsyncGenerator[dict, None]:
        headers = {
            "Authorization": f"Bearer {self.api_key.get_secret_value()}",
            "Content-Type": "application/json",
            "X-Client": f"langchain-igpt/{self._version}"
        }
        delay = self.backoff_base
        timeout = aiohttp.ClientTimeout(
            total=None,
            connect=60,    # connect timeout (sec)
            sock_read=600, # read timeout (sec): max idle time between chunks
        )
        async with aiohttp.ClientSession(headers=headers, timeout=timeout) as session:
            for attempt in range(self.max_retries):
                try:
                    async with session.post(url, json=payload) as resp:
                        if resp.status == 200:
                            async for chunk in resp.content:
                                if chunk:
                                    chunk_str = chunk.decode()
                                    if chunk_str.startswith("data:"):
                                        json_str = chunk_str[len("data:"):].strip()
                                    else:
                                        json_str = chunk_str
                                        
                                    if not json_str:
                                        continue
                                    
                                    try:
                                        data = json.loads(json_str)
                                        if data.get("context") and not data["context"].get("datasources"):
                                            warnings.warn(
                                                "No datasources found. Please authorize and connect a new datasource at " 
                                                "https://igpt.ai/hub/playground/connectors-authorize/"
                                            )
                                        yield data
                                    except json.JSONDecodeError:
                                        continue
                            return
                        elif resp.status == 429 or 500 <= resp.status < 600:
                            pass
                        else:
                            yield {"error": "client_error"}
                            return
                except (aiohttp.ClientError, asyncio.TimeoutError):
                    pass

                if attempt < self.max_retries - 1:
                    await asyncio.sleep(delay)
                    delay *= self.backoff_factor
                
            yield {"error": "network_error"}
