from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from webscout import WEBS, transcriber
from typing import Optional
from fastapi.encoders import jsonable_encoder
from bs4 import BeautifulSoup
import requests
import urllib.parse
import os

app = FastAPI()

@app.get("/")
async def root():
    return {"message": "API documentation can be found at /docs"}

@app.get("/health")
async def health_check():
    return {"status": "OK"}

@app.get("/api/search")
async def search(
    q: str,
    max_results: int = 10,
    timelimit: Optional[str] = None,
    safesearch: str = "moderate",
    region: str = "wt-wt",
    backend: str = "api"
):
    """Perform a text search."""
    try:
        with WEBS() as webs:
            results = webs.text(keywords=q, region=region, safesearch=safesearch, timelimit=timelimit, backend=backend, max_results=max_results)
            return JSONResponse(content=jsonable_encoder(results))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during search: {e}")

@app.get("/api/images")
async def images(
    q: str,
    max_results: int = 10,
    safesearch: str = "moderate",
    region: str = "wt-wt",
    timelimit: Optional[str] = None,
    size: Optional[str] = None,
    color: Optional[str] = None,
    type_image: Optional[str] = None,
    layout: Optional[str] = None,
    license_image: Optional[str] = None
):
    """Perform an image search."""
    try:
        with WEBS() as webs:
            results = webs.images(keywords=q, region=region, safesearch=safesearch, timelimit=timelimit, size=size, color=color, type_image=type_image, layout=layout, license_image=license_image, max_results=max_results)
            return JSONResponse(content=jsonable_encoder(results))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during image search: {e}")

@app.get("/api/videos")
async def videos(
    q: str,
    max_results: int = 10,
    safesearch: str = "moderate",
    region: str = "wt-wt",
    timelimit: Optional[str] = None,
    resolution: Optional[str] = None,
    duration: Optional[str] = None,
    license_videos: Optional[str] = None
):
    """Perform a video search."""
    try:
        with WEBS() as webs:
            results = webs.videos(keywords=q, region=region, safesearch=safesearch, timelimit=timelimit, resolution=resolution, duration=duration, license_videos=license_videos, max_results=max_results)
            return JSONResponse(content=jsonable_encoder(results))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during video search: {e}")

@app.get("/api/news")
async def news(
    q: str,
    max_results: int = 10,
    safesearch: str = "moderate",
    region: str = "wt-wt",
    timelimit: Optional[str] = None
):
    """Perform a news search."""
    try:
        with WEBS() as webs:
            results = webs.news(keywords=q, region=region, safesearch=safesearch, timelimit=timelimit, max_results=max_results)
            return JSONResponse(content=jsonable_encoder(results))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during news search: {e}")

@app.get("/api/answers")
async def answers(q: str):
    """Get instant answers for a query."""
    try:
        with WEBS() as webs:
            results = webs.answers(keywords=q)
            return JSONResponse(content=jsonable_encoder(results))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error getting instant answers: {e}")

@app.get("/api/suggestions")
async def suggestions(q: str, region: str = "wt-wt"):
    """Get search suggestions for a query."""
    try:
        with WEBS() as webs:
            results = webs.suggestions(keywords=q, region=region)
            return JSONResponse(content=jsonable_encoder(results))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error getting search suggestions: {e}")

@app.get("/api/chat")
async def chat(
    q: str,
    model: str = "gpt-3.5"
):
    """Perform a text search."""
    try:
        with WEBS() as webs:
            results = webs.chat(keywords=q, model=model)
            return JSONResponse(content=jsonable_encoder(results))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error getting chat results: {e}")

def extract_text_from_webpage(html_content):
    """Extracts visible text from HTML content using BeautifulSoup."""
    soup = BeautifulSoup(html_content, "html.parser")
    # Remove unwanted tags
    for tag in soup(["script", "style", "header", "footer", "nav"]):
        tag.extract()
    # Get the remaining visible text
    visible_text = soup.get_text(strip=True)
    return visible_text

@app.get("/api/web_extract")
async def web_extract(
    url: str,
    max_chars: int = 12000,  # Adjust based on token limit
):
    """Extracts text from a given URL."""
    try:
        response = requests.get(url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"})
        response.raise_for_status()
        visible_text = extract_text_from_webpage(response.text)
        if len(visible_text) > max_chars:
            visible_text = visible_text[:max_chars] + "..."
        return {"url": url, "text": visible_text}
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=f"Error fetching or processing URL: {e}")

@app.get("/api/search-and-extract")
async def web_search_and_extract(
    q: str,
    max_results: int = 3,
    timelimit: Optional[str] = None,
    safesearch: str = "moderate",
    region: str = "wt-wt",
    backend: str = "api",
    max_chars: int = 6000,
    extract_only: bool = False
):
    """
    Searches using WEBS, extracts text from the top results, and returns both.
    """
    try:
        with WEBS() as webs:
            # Perform WEBS search
            search_results = webs.text(keywords=q, region=region, safesearch=safesearch,
                                     timelimit=timelimit, backend=backend, max_results=max_results)

            # Extract text from each result's link
            extracted_results = []
            for result in search_results:
                if 'href' in result:
                    link = result['href']
                    try:
                        response = requests.get(link, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"})
                        response.raise_for_status()
                        visible_text = extract_text_from_webpage(response.text)
                        if len(visible_text) > max_chars:
                            visible_text = visible_text[:max_chars] + "..."
                        extracted_results.append({"link": link, "text": visible_text})
                    except requests.exceptions.RequestException as e:
                        print(f"Error fetching or processing {link}: {e}")
                        extracted_results.append({"link": link, "text": None})
                else:
                    extracted_results.append({"link": None, "text": None})
            if extract_only:
                return JSONResponse(content=jsonable_encoder({extracted_results}))
            else:
                return JSONResponse(content=jsonable_encoder({"search_results": search_results, "extracted_results": extracted_results}))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during search and extraction: {e}")

@app.get("/api/website_summarizer")
async def website_summarizer(url: str):
    """Summarizes the content of a given URL using a chat model."""
    try:
        # Extract text from the given URL
        response = requests.get(url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"})
        response.raise_for_status()
        visible_text = extract_text_from_webpage(response.text)
        if len(visible_text) > 7500:  # Adjust max_chars based on your needs
            visible_text = visible_text[:7500] + "..."

        # Use chat model to summarize the extracted text
        with WEBS() as webs:
            summary_prompt = f"Summarize this in detail in Paragraph: {visible_text}"
            summary_result = webs.chat(keywords=summary_prompt, model="gpt-3.5")

        # Return the summary result
        return JSONResponse(content=jsonable_encoder({summary_result}))

    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=f"Error fetching or processing URL: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during summarization: {e}")

@app.get("/api/ask_website")
async def ask_website(url: str, question: str, model: str = "llama-3-70b"):
    """
    Asks a question about the content of a given website.
    """
    try:
        # Extract text from the given URL
        response = requests.get(url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"})
        response.raise_for_status()
        visible_text = extract_text_from_webpage(response.text)
        if len(visible_text) > 7500:  # Adjust max_chars based on your needs
            visible_text = visible_text[:7500] + "..."

        # Construct a prompt for the chat model
        prompt = f"Based on the following text, answer this question in Paragraph: [QUESTION] {question} [TEXT] {visible_text}"

        # Use chat model to get the answer
        with WEBS() as webs:
            answer_result = webs.chat(keywords=prompt, model=model)

        # Return the answer result
        return JSONResponse(content=jsonable_encoder({answer_result}))

    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=f"Error fetching or processing URL: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during question answering: {e}")
        
@app.get("/api/maps")
async def maps(
    q: str,
    place: Optional[str] = None,
    street: Optional[str] = None,
    city: Optional[str] = None,
    county: Optional[str] = None,
    state: Optional[str] = None,
    country: Optional[str] = None,
    postalcode: Optional[str] = None,
    latitude: Optional[str] = None,
    longitude: Optional[str] = None,
    radius: int = 0,
    max_results: int = 10
):
    """Perform a maps search."""
    try:
        with WEBS() as webs:
            results = webs.maps(keywords=q, place=place, street=street, city=city, county=county, state=state, country=country, postalcode=postalcode, latitude=latitude, longitude=longitude, radius=radius, max_results=max_results)
            return JSONResponse(content=jsonable_encoder(results))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during maps search: {e}")

@app.get("/api/translate")
async def translate(
    q: str,
    from_: Optional[str] = None,
    to: str = "en"
):
    """Translate text."""
    try:
        with WEBS() as webs:
            results = webs.translate(keywords=q, from_=from_, to=to)
            return JSONResponse(content=jsonable_encoder(results))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during translation: {e}")

@app.get("/api/youtube/transcript")
async def youtube_transcript(
    video_id: str,
    languages: str = "en",
    preserve_formatting: bool = False
):
    """Get the transcript of a YouTube video."""
    try:
        languages_list = languages.split(",")
        transcript = transcriber.get_transcript(video_id, languages=languages_list, preserve_formatting=preserve_formatting)
        return JSONResponse(content=jsonable_encoder(transcript))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error getting YouTube transcript: {e}")
        
import requests
@app.get("/weather/json/{location}")
def get_weather_json(location: str):
    url = f"https://wttr.in/{location}?format=j1"
    response = requests.get(url)
    if response.status_code == 200:
        return response.json()
    else:
        return {"error": f"Unable to fetch weather data. Status code: {response.status_code}"}

@app.get("/weather/ascii/{location}")
def get_ascii_weather(location: str):
    url = f"https://wttr.in/{location}"
    response = requests.get(url, headers={'User-Agent': 'curl'})
    if response.status_code == 200:
        return response.text
    else:
        return {"error": f"Unable to fetch weather data. Status code: {response.status_code}"}


import os
from typing import Optional
from webscout import WEBS
from webscout.DWEBS import (
    GoogleSearcher,
    QueryResultsExtractor,
    BatchWebpageFetcher,
    BatchWebpageContentExtractor,
    OSEnver,
)
from webscout.Provider import DeepInfra
from rich.console import Console
from rich.markdown import Markdown
from tiktoken import get_encoding

class AISearchEngine:
    def __init__(
        self,
        google_search_result_num: int = 3,
        duckduckgo_search_result_num: int = 3,
        ai_model_1: str = 'microsoft/WizardLM-2-8x22B',  # First AI model
        ai_model_2: str = 'Qwen/Qwen2-72B-Instruct',  # Second AI model for summarization
        ai_max_tokens: int = 8192,
        proxy: Optional[str] = None,
        max_google_content_tokens: int = 60000  # Maximum tokens for Google content
    ):
        self.console = Console()
        self.enver = OSEnver()
        self.enver.set_envs(proxies=proxy)
        self.google_searcher = GoogleSearcher()
        self.query_results_extractor = QueryResultsExtractor()
        self.batch_webpage_fetcher = BatchWebpageFetcher()
        self.batch_webpage_content_extractor = BatchWebpageContentExtractor()
        self.web_search_client = WEBS(proxy=proxy)
        self.google_search_result_num = google_search_result_num
        self.duckduckgo_search_result_num = duckduckgo_search_result_num
        self.max_google_content_tokens = max_google_content_tokens
        self.encoding = get_encoding("cl100k_base")  # Use tiktoken for tokenization

        # Detailed system prompt for the first DeepInfra AI model (WizardLM)
        self.system_prompt_1 = """You are a highly advanced AI assistant, designed to be exceptionally helpful and informative. Your primary function is to provide comprehensive, accurate, and insightful answers to user queries, even if they are complex or open-ended.

        You have access to powerful tools:
        - **Google Search Results:** These provide detailed content extracted from relevant web pages.
        - **DuckDuckGo Search Summaries:** These offer additional perspectives and quick summaries of findings from another search engine.

        Your top priorities are:
        1. **Accuracy:** Ensure your responses are factually correct and grounded in the provided source material. Verify information whenever possible.
        2. **Relevance:**  Focus on directly addressing the user's question. Don't stray into irrelevant tangents. 
        3. **Clarity:**  Present information in a clear, organized, and easy-to-understand manner.  Use bullet points, numbered lists, or concise paragraphs to structure your response.
        4. **Source Citation:** Whenever you use information from the web search results, clearly cite your source using the webpage title and URL. This helps the user understand where the information is coming from.

        Here's a detailed breakdown of how to use the search results effectively:
        - **Google (Extracted Content):** Treat this as your primary source. Analyze the content thoroughly to find the most relevant information. Don't hesitate to quote directly from the source when appropriate.
        - **DuckDuckGo (Summaries):** Use these to broaden your understanding and potentially find additional perspectives. The summaries provide a quick overview of the top results, which can be helpful for identifying key themes or different angles on the topic. 

        Additional Guidelines:
        - **Vague Queries:** If the user's query is unclear or lacks specifics, politely ask clarifying questions to better understand their needs.
        - **Information Gaps:** If you can't find a satisfactory answer in the provided search results, acknowledge this honestly. Let the user know that you need more information to help them or that the specific answer is not available in the current search results. 
        - **Neutrality:** Maintain a neutral tone and avoid expressing personal opinions, beliefs, or emotions.  
        - **No Self-Promotion:**  Do not refer to yourself as an AI or a language model. Focus on providing helpful information to the user.

        Remember, your ultimate goal is to be a reliable and helpful assistant to the user. Use your knowledge and the provided search results to their fullest potential to achieve this goal."""

        self.ai_client_1 = DeepInfra(
            model=ai_model_1, 
            system_prompt=self.system_prompt_1, 
            max_tokens=ai_max_tokens, 
            is_conversation=False,  # Disable conversation history
            timeout=100 # Set timeout to 100 seconds
        )

        # System prompt for the second AI (Qwen) for professional summarization
        self.system_prompt_2 = """You are a highly skilled AI assistant, specifically designed to condense and refine information into professional summaries. Your primary function is to transform the detailed responses of another AI model into clear, concise, and easily digestible reports. 

        Your objective is to craft summaries that are:
        - **Accurate:** Faithfully represent the key points and findings of the original AI response without introducing any new information or interpretations.
        - **Concise:**  Distill the information into its most essential elements, eliminating unnecessary details or redundancy. Aim for brevity without sacrificing clarity.
        - **Informative:**  Ensure the summary provides a comprehensive overview of the main points, insights, and relevant sources.  
        - **Professional:**  Use a formal and objective tone. Avoid casual language, personal opinions, or subjective statements.  

        Here's a breakdown of the key elements to include in your summaries:
        - **Main Points:** Identify and succinctly state the most important arguments, conclusions, or findings presented in the AI's response.
        - **Key Insights:** Highlight any particularly insightful observations, trends, or analyses that emerge from the AI's response. 
        - **Source Attribution:**  If the original AI cited any sources (web pages, articles, etc.), list them clearly in your summary, using proper citation format (e.g., title and URL).

        Formatting Guidelines: 
        - **Structure:** Organize your summary using bullet points or a numbered list for maximum clarity and readability.
        - **Length:**  Keep your summaries concise, aiming for a length that is significantly shorter than the original AI response. 

        Remember, your role is to provide a refined and professional distillation of the AI's output, making it readily accessible and understandable for a professional audience."""

        self.ai_client_2 = DeepInfra(
            model=ai_model_2,
            system_prompt=self.system_prompt_2,
            max_tokens=ai_max_tokens,
            is_conversation=False,
            timeout=100
        )

    def search_google(self, query: str) -> dict:
        """Search Google and extract results."""
        self.console.print(f"[bold blue]Searching Google for:[/] {query}")
        html_path = self.google_searcher.search(query, result_num=self.google_search_result_num)
        search_results = self.query_results_extractor.extract(html_path)
        self.console.print(f"[bold blue]Extracted {len(search_results['query_results'])} Google results.[/]")
        return search_results

    def search_duckduckgo(self, query: str) -> list[dict]:
        """Search DuckDuckGo and extract results."""
        self.console.print(f"[bold blue]Searching DuckDuckGo for:[/] {query}")
        search_results = self.web_search_client.text(keywords=query, max_results=self.duckduckgo_search_result_num)
        self.console.print(f"[bold blue]Extracted {len(search_results)} DuckDuckGo results.[/]")
        return search_results

    def fetch_and_extract_content(self, google_results: list[dict]) -> list[dict]:
        """Fetch and extract content from Google search result URLs only, with truncation."""
        urls = [result['url'] for result in google_results]

        self.console.print(f"[bold blue]Fetching {len(urls)} webpages...[/]")
        url_and_html_path_list = self.batch_webpage_fetcher.fetch(urls)
        html_paths = [item['html_path'] for item in url_and_html_path_list]
        self.console.print(f"[bold blue]Extracting content from {len(html_paths)} webpages...[/]")
        html_path_and_content_list = self.batch_webpage_content_extractor.extract(html_paths)

        # Truncate the combined content to the maximum token limit
        combined_content = ""
        current_token_count = 0
        for item in html_path_and_content_list:
            content = item['extracted_content']
            token_count = len(self.encoding.encode(content))
            if current_token_count + token_count <= self.max_google_content_tokens:
                combined_content += content + "\n\n"
                current_token_count += token_count
            else:
                # Truncate the current content to fit within the limit
                remaining_tokens = self.max_google_content_tokens - current_token_count
                truncated_content = self.encoding.decode(self.encoding.encode(content)[:remaining_tokens])
                combined_content += truncated_content
                break

        # Update the content in the list with the truncated combined content
        html_path_and_content_list = [{'html_path': '', 'extracted_content': combined_content}]

        return html_path_and_content_list

    def ask_ai_1(self, query: str, web_content: str, duckduckgo_results: list[dict]) -> str:
        """Ask the first AI model (WizardLM) a question."""
        self.console.print(f"[bold blue]Asking AI 1:[/] {query}")
        
        duckduckgo_summary = "\n\n".join([
            f"**{result['title']}** ({result['href']})\n{result['body']}" 
            for result in duckduckgo_results
        ])

        prompt = (
            f"Based on the following Google search results (with extracted content):\n\n"
            f"{web_content}\n\n"
            f"And the following DuckDuckGo search summaries:\n\n"
            f"{duckduckgo_summary}\n\n"
            f"Answer the following question: {query}"
        )

        ai_response = self.ai_client_1.chat(prompt)
        # self.console.print(f"[bold blue]AI 1 Response:[/]\n{ai_response}")
        return ai_response

    def ask_ai_2(self, ai_1_response: str) -> str:
        """Ask the second AI model (Qwen) to summarize the first AI's response."""
        self.console.print(f"[bold blue]Asking AI 2 to summarize AI 1's response...[/]")
        prompt = f"Please summarize the following text in a professional format:\n\n{ai_1_response}"
        ai_response = self.ai_client_2.chat(prompt)
        # self.console.print(f"[bold blue]AI 2 Summary:[/]\n{ai_response}")
        return ai_response

    def run(self, query: str):
        """Run the AI search engine."""
        google_results = self.search_google(query)['query_results']
        duckduckgo_results = self.search_duckduckgo(query)
        html_path_and_content_list = self.fetch_and_extract_content(google_results)
        web_content = "\n\n".join([item['extracted_content'] for item in html_path_and_content_list])
        
        # Get response from the first AI
        ai_1_response = self.ask_ai_1(query, web_content, duckduckgo_results)

        # Summarize the first AI's response using the second AI
        ai_2_summary = self.ask_ai_2(ai_1_response)
        self.console.print("[bold green]Full Response[/]:", end="")
        self.console.print(Markdown(ai_1_response))

        self.console.print("[bold red]Summary[/]:", end="")
        self.console.print(Markdown(ai_2_summary))

# Initialize the AI search engine outside of the endpoint
ai_search_engine = AISearchEngine()  

@app.post("/api/ai-search")
async def ai_search(query: str):
    """
    Performs an AI-powered search using Google, DuckDuckGo, and two large language models. 
    """
    try:
        ai_search_engine.run(query)
        return {"message": "AI search completed. Check your console for the results."}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during AI search: {e}")

# Run the API server if this script is executed
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)