import json
import requests
from typing import List, Optional
from langchain_community.document_transformers import MarkdownifyTransformer
from langchain_core.documents import Document
from langchain_core.pydantic_v1 import BaseModel, Field

class BraveAIWrapper(BaseModel):
    api_key: str = Field(..., description="API key for Brave Search")
    base_search_url: str = Field("https://api.search.brave.com/res/v1/web/search", const=True)
    base_summarize_url: str = Field("https://api.search.brave.com/res/v1/summarizer/search", const=True)
    headers: dict = Field(default_factory=dict, description="HTTP headers for API requests")

    class Config:
        arbitrary_types_allowed = True

    def __init__(self, **data):
        super().__init__(**data)
        self.headers = {
            "X-Subscription-Token": self.api_key,
            "Accept": "application/json",
            "Accept-Encoding": "gzip",
        }

    def get_brave_results(self, query: str, count: int = 3, safe_search: str = 'off') -> Optional[dict]:
        """
        Get search results from Brave Search.

        Args:
            query (str): The search query.
            count (int): Number of results to return.
            safe_search (str): Safe search filter (off, moderate, strict).

        Returns:
            Optional[dict]: JSON response from Brave Search API or None if an error occurs.
        """
        params = {
            "q": query,
            "count": count,
            "summary": True,
            "safe_search": safe_search,
            "extra_snippets": True,
        }
        try:
            response = requests.get(self.base_search_url, headers=self.headers, params=params)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error: {e}")
            return None

    def summarize_results(self, summarizer_key: str) -> Optional[dict]:
        """
        Summarize search results using Brave Summarizer.

        Args:
            summarizer_key (str): The key for the summarizer.

        Returns:
            Optional[dict]: JSON response from Brave Summarizer API or None if an error occurs.
        """
        params = {"key": summarizer_key}
        try:
            response = requests.get(self.base_summarize_url, headers=self.headers, params=params)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error: {e}")
            return None

    def get_and_summarize(self, query: str, count: int = 3, safe_search: str = 'off') -> Optional[str]:
        """
        Get and summarize search results from Brave Search.

        Args:
            query (str): The search query.
            count (int): Number of results to return.
            safe_search (str): Safe search filter (off, moderate, strict).

        Returns:
            Optional[str]: Summarized result or None if an error occurs.
        """
        results = self.get_brave_results(query, count, safe_search)
        if results and 'summarizer' in results:
            summarizer_key = results['summarizer']['key']
            summary = self.summarize_results(summarizer_key)
            if summary and 'summary' in summary and len(summary['summary']) > 0:
                return summary['summary'][0]['data']
        return None