"""
LangChain integration for UnSearch - drop-in Tavily replacement.

Usage:
    # Instead of:
    from langchain_community.tools import TavilySearchResults
    tool = TavilySearchResults(api_key="tvly-...")
    
    # Use:
    from unsearch.langchain import UnSearchResults
    tool = UnSearchResults(api_key="uns-...")
    
    # Or migrate existing code with one line:
    from unsearch.langchain import UnSearchResults as TavilySearchResults
"""
import os
from typing import Any, Dict, List, Optional, Type, Literal

try:
    from langchain_core.callbacks import CallbackManagerForToolRun
    from langchain_core.tools import BaseTool
    from pydantic import BaseModel, Field
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    BaseTool = object
    BaseModel = object


from unsearch.client import UnSearchClient


if LANGCHAIN_AVAILABLE:
    
    class UnSearchInput(BaseModel):
        """Input schema for UnSearch tool."""
        query: str = Field(description="Search query to look up")

    
    class UnSearchResults(BaseTool):
        """
        UnSearch search tool for LangChain - Tavily drop-in replacement.
        
        This tool searches the web using UnSearch API and returns results
        optimized for LLM consumption.
        
        Example:
            >>> from unsearch.langchain import UnSearchResults
            >>> tool = UnSearchResults(api_key="your-api-key")
            >>> results = tool.invoke("What is LangChain?")
        
        Migration from Tavily:
            # Before
            from langchain_community.tools import TavilySearchResults
            tool = TavilySearchResults(api_key="tvly-...")
            
            # After (option 1: direct replacement)
            from unsearch.langchain import UnSearchResults
            tool = UnSearchResults(api_key="uns-...")
            
            # After (option 2: alias for minimal code changes)
            from unsearch.langchain import UnSearchResults as TavilySearchResults
            tool = TavilySearchResults(api_key="uns-...")
        """
        
        name: str = "unsearch_results"
        description: str = (
            "A search engine optimized for comprehensive, accurate, and trusted results. "
            "Useful for answering questions about current events, facts, or any topic "
            "requiring up-to-date information from the web."
        )
        args_schema: Type[BaseModel] = UnSearchInput
        
        api_key: Optional[str] = None
        base_url: Optional[str] = None
        max_results: int = 5
        search_depth: Literal["basic", "advanced"] = "basic"
        include_answer: bool = False
        include_raw_content: bool = False
        include_images: bool = False
        zero_retention: bool = False
        
        _client: Optional[UnSearchClient] = None
        
        def __init__(
            self,
            api_key: Optional[str] = None,
            base_url: Optional[str] = None,
            max_results: int = 5,
            search_depth: Literal["basic", "advanced"] = "basic",
            include_answer: bool = False,
            include_raw_content: bool = False,
            include_images: bool = False,
            zero_retention: bool = False,
            **kwargs
        ):
            """
            Initialize UnSearch tool.
            
            Args:
                api_key: UnSearch API key (or set UNSEARCH_API_KEY env var)
                base_url: API base URL (for self-hosted instances)
                max_results: Maximum search results (1-20)
                search_depth: "basic" or "advanced"
                include_answer: Include AI-generated answer
                include_raw_content: Include full page content
                include_images: Include image results
                zero_retention: Enable zero-retention mode
            """
            super().__init__(**kwargs)
            self.api_key = api_key or os.environ.get("UNSEARCH_API_KEY")
            self.base_url = base_url
            self.max_results = max_results
            self.search_depth = search_depth
            self.include_answer = include_answer
            self.include_raw_content = include_raw_content
            self.include_images = include_images
            self.zero_retention = zero_retention
        
        @property
        def client(self) -> UnSearchClient:
            """Get or create UnSearch client."""
            if self._client is None:
                self._client = UnSearchClient(
                    api_key=self.api_key,
                    base_url=self.base_url,
                    zero_retention=self.zero_retention
                )
            return self._client
        
        def _run(
            self,
            query: str,
            run_manager: Optional[CallbackManagerForToolRun] = None
        ) -> List[Dict[str, Any]]:
            """Execute search and return results."""
            response = self.client.search(
                query=query,
                search_depth=self.search_depth,
                max_results=self.max_results,
                include_answer=self.include_answer,
                include_raw_content=self.include_raw_content,
                include_images=self.include_images
            )
            
            results = []
            for r in response.results:
                result = {
                    "title": r.title,
                    "url": r.url,
                    "content": r.content,
                    "score": r.score
                }
                if r.raw_content:
                    result["raw_content"] = r.raw_content
                results.append(result)
            
            return results
    
    
    class UnSearchQAResults(BaseTool):
        """
        UnSearch Q&A tool - returns AI-generated answer with sources.
        
        Similar to TavilyAnswer but with UnSearch backend.
        """
        
        name: str = "unsearch_qa"
        description: str = (
            "A search engine that provides direct answers to questions "
            "with cited sources. Best for factual questions."
        )
        args_schema: Type[BaseModel] = UnSearchInput
        
        api_key: Optional[str] = None
        base_url: Optional[str] = None
        max_results: int = 5
        search_depth: Literal["basic", "advanced"] = "advanced"
        zero_retention: bool = False
        
        _client: Optional[UnSearchClient] = None
        
        def __init__(
            self,
            api_key: Optional[str] = None,
            base_url: Optional[str] = None,
            max_results: int = 5,
            search_depth: Literal["basic", "advanced"] = "advanced",
            zero_retention: bool = False,
            **kwargs
        ):
            super().__init__(**kwargs)
            self.api_key = api_key or os.environ.get("UNSEARCH_API_KEY")
            self.base_url = base_url
            self.max_results = max_results
            self.search_depth = search_depth
            self.zero_retention = zero_retention
        
        @property
        def client(self) -> UnSearchClient:
            if self._client is None:
                self._client = UnSearchClient(
                    api_key=self.api_key,
                    base_url=self.base_url,
                    zero_retention=self.zero_retention
                )
            return self._client
        
        def _run(
            self,
            query: str,
            run_manager: Optional[CallbackManagerForToolRun] = None
        ) -> str:
            """Execute Q&A search and return answer."""
            return self.client.qna_search(
                query=query,
                search_depth=self.search_depth,
                max_results=self.max_results
            )


    # Aliases for Tavily migration
    TavilySearchResults = UnSearchResults
    TavilyAnswer = UnSearchQAResults

else:
    # Provide helpful error if LangChain not installed
    class UnSearchResults:
        def __init__(self, *args, **kwargs):
            raise ImportError(
                "LangChain is required for this integration. "
                "Install with: pip install langchain-core"
            )
    
    class UnSearchQAResults:
        def __init__(self, *args, **kwargs):
            raise ImportError(
                "LangChain is required for this integration. "
                "Install with: pip install langchain-core"
            )
    
    TavilySearchResults = UnSearchResults
    TavilyAnswer = UnSearchQAResults
