from typing import Dict, List, Optional, Type
import json, logging
from pathlib import Path
from sys import platform

from bs4 import BeautifulSoup
from selenium.common.exceptions import WebDriverException
from selenium.webdriver.chrome.options import Options as ChromeOptions
from selenium.webdriver.chrome.service import Service as ChromeDriverService
from selenium.webdriver.chrome.webdriver import WebDriver as ChromeDriver
from selenium.webdriver.common.by import By
from selenium.webdriver.edge.options import Options as EdgeOptions
from selenium.webdriver.edge.service import Service as EdgeDriverService
from selenium.webdriver.edge.webdriver import WebDriver as EdgeDriver
from selenium.webdriver.firefox.options import Options as FirefoxOptions
from selenium.webdriver.firefox.service import Service as GeckoDriverService
from selenium.webdriver.firefox.webdriver import WebDriver as FirefoxDriver
from selenium.webdriver.remote.webdriver import WebDriver
from selenium.webdriver.safari.options import Options as SafariOptions
from selenium.webdriver.safari.webdriver import WebDriver as SafariDriver
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.wait import WebDriverWait
from webdriver_manager.chrome import ChromeDriverManager
from webdriver_manager.firefox import GeckoDriverManager
from webdriver_manager.microsoft import EdgeChromiumDriverManager as EdgeDriverManager

from autoxx.config.config import GlobalConfig
from autoxx.utils.llm import get_chat_completion
from autoxx.utils.processing_text import split_text
from autoxx.utils.token_counter import count_message_tokens
from autoxx.utils.processing_html import extract_hyperlinks, format_hyperlinks
from autoxx.utils.base import Message

BrowserOptions = ChromeOptions | EdgeOptions | FirefoxOptions | SafariOptions
BrowserSummaryModel = "gpt-3.5-turbo-16k"
MaxSummaryTokenSize = 2048
summary_response_format = {
    "relevant_content": "Extract all information in detail relevant to question. if there is no relevant information, summarize the text briefly.",
    "relevance": "Yes/No # Whether there is a relevance",
}
formatted_summary_response_format = json.dumps(summary_response_format, indent=4)


FILE_DIR = Path(__file__).parent

def create_summary_message(chunk: str, question: str) -> Message:
    """Create a message for the chat completion

    Args:
        chunk (str): The chunk of text to summarize
        question (str): The question to answer

    Returns:
        Dict[str, str]: The message to send to the chat completion
    """
    return Message(
        role="user", 
        content=(
            f'"""{chunk}""" Analyze the above text and extract all information in detail relevant to '
            f'question: "{question}" -- if there is no relevant information, summarize the text.'
            f"\nResponse Format: \n{formatted_summary_response_format} "
            f"\nEnsure the response can be parsed by Python json.loads"
        )
    )

def no_relevance(answer: str) -> bool:
    return answer.strip().replace(".", "").lower() == "no"

def scroll_to_percentage(driver: WebDriver, ratio: float) -> None:
    """Scroll to a percentage of the page

    Args:
        driver (WebDriver): The webdriver to use
        ratio (float): The percentage to scroll to

    Raises:
        ValueError: If the ratio is not between 0 and 1
    """
    if ratio < 0 or ratio > 1:
        raise ValueError("Percentage should be between 0 and 1")
    driver.execute_script(f"window.scrollTo(0, document.body.scrollHeight * {ratio});")

def scrape_links_with_selenium(driver: WebDriver, url: str) -> list[str]:
    """Scrape links from a website using selenium

    Args:
        driver (WebDriver): The webdriver to use to scrape the links

    Returns:
        List[str]: The links scraped from the website
    """
    page_source = driver.page_source
    soup = BeautifulSoup(page_source, "html.parser")

    for script in soup(["script", "style"]):
        script.extract()

    hyperlinks = extract_hyperlinks(soup, url)

    return format_hyperlinks(hyperlinks)

def scrape_text_with_selenium(url: str) -> tuple[WebDriver, str]:
    """Scrape text from a website using selenium

    Args:
        url (str): The url of the website to scrape

    Returns:
        Tuple[WebDriver, str]: The webdriver and the text scraped from the website
    """
    logging.getLogger("selenium").setLevel(logging.CRITICAL)

    options_available: dict[str, Type[BrowserOptions]] = {
        "chrome": ChromeOptions,
        "edge": EdgeOptions,
        "firefox": FirefoxOptions,
        "safari": SafariOptions,
    }

    CFG = GlobalConfig().get()
    options: BrowserOptions = options_available[CFG.selenium_web_browser]()
    options.add_argument(
        "user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/112.0.5615.49 Safari/537.36"
    )

    if CFG.selenium_web_browser == "firefox":
        if CFG.selenium_headless:
            options.headless = True
            options.add_argument("--disable-gpu")
        driver = FirefoxDriver(
            service=GeckoDriverService(GeckoDriverManager().install()), options=options
        )
    elif CFG.selenium_web_browser == "edge":
        driver = EdgeDriver(
            service=EdgeDriverService(EdgeDriverManager().install()), options=options
        )
    elif CFG.selenium_web_browser == "safari":
        # Requires a bit more setup on the users end
        # See https://developer.apple.com/documentation/webkit/testing_with_webdriver_in_safari
        driver = SafariDriver(options=options)
    else:
        if platform == "linux" or platform == "linux2":
            options.add_argument("--disable-dev-shm-usage")
            options.add_argument("--remote-debugging-port=9222")

        options.add_argument("--no-sandbox")
        if CFG.selenium_headless:
            options.add_argument("--headless=new")
            options.add_argument("--disable-gpu")

        chromium_driver_path = Path("/usr/bin/chromedriver")

        driver = ChromeDriver(
            service=ChromeDriverService(str(chromium_driver_path))
            if chromium_driver_path.exists()
            else ChromeDriverService(ChromeDriverManager().install()),
            options=options,
        )
    driver.get(url)

    WebDriverWait(driver, 10).until(
        EC.presence_of_element_located((By.TAG_NAME, "body"))
    )

    # Get the HTML content directly from the browser's DOM
    page_source = driver.execute_script("return document.body.outerHTML;")
    soup = BeautifulSoup(page_source, "html.parser")

    for script in soup(["script", "style"]):
        script.extract()

    text = soup.get_text()
    lines = (line.strip() for line in text.splitlines())
    chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
    text = "\n".join(chunk for chunk in chunks if chunk)
    return driver, text

def summarize_chunks(text:str, question: str, url: str, driver: Optional[WebDriver] = None) -> List[str]:
    """Summarize a list of chunks

    Args:
        chunks (list[str]): The chunks to summarize
        question (str): The question to answer

    Returns:
        str: The summary of the chunks
    """
    logging.debug(f"Memorizing text:\n{'-'*32}\n{text}\n{'-'*32}\n")
    chunks = [
        chunk
        for chunk, _ in (
            split_text(text, for_model=BrowserSummaryModel)
        )
    ]

    summaries = []
    scroll_ratio = 1 / len(chunks)
    for i, chunk in enumerate(chunks):
        if driver:
            scroll_to_percentage(driver, scroll_ratio * i)

        messages = [create_summary_message(chunk, question)]
        tokens_for_chunk = count_message_tokens(messages)

        logging.info(
            f"Summarizing {url} chunk {i + 1} / {len(chunks)} of length {len(chunk)} characters, or {tokens_for_chunk} tokens"
        )

        summary_response_str = get_chat_completion(
            model=BrowserSummaryModel,
            messages=messages,
        )

        summary_response = json.loads(summary_response_str)
        if no_relevance(summary_response.get("relevance", "")):
            logging.info(f"Skipping chunk {i + 1} / {len(chunks)} of length {len(chunk)} characters, or {tokens_for_chunk} tokens, not relevant.")
            continue
        summary = summary_response.get("relevant_content", summary_response_str)

        logging.info(f"Chunk {i + 1} / {len(chunks)} summary.\n{summary}")

        summaries.append(summary)

    return summaries

def add_header(driver: WebDriver) -> None:
    """Add a header to the website

    Args:
        driver (WebDriver): The webdriver to use to add the header

    Returns:
        None
    """
    try:
        with open(f"{FILE_DIR}/js/overlay.js", "r") as overlay_file:
            overlay_script = overlay_file.read()
        driver.execute_script(overlay_script)
    except Exception as e:
        print(f"Error executing overlay.js: {e}")

def close_browser(driver: WebDriver) -> None:
    """Close the browser

    Args:
        driver (WebDriver): The webdriver to close

    Returns:
        None
    """
    driver.quit()


def browse_website(url: str, question: str) -> str | None:
    """Browse a website and return the answer and links to the user

    Args:
        url (str): The url of the website to browse
        question (str): The question asked by the user

    Returns:
        Tuple[str, WebDriver]: The answer and links to the user and the webdriver
    """
    try:
        driver, text = scrape_text_with_selenium(url=url)
    except WebDriverException as e:
        # These errors are often quite long and include lots of context.
        # Just grab the first line.
        msg = e.msg.split("\n")[0]
        return f"Error: {msg}"
    add_header(driver)

    summaries = summarize_chunks(text=text, question=question, url=url, driver=driver)
    text_gen = "\n".join(summaries)
    while count_message_tokens([create_summary_message(text_gen, question)]) > MaxSummaryTokenSize:
       summaries = summarize_chunks(text=text_gen, question=question, url=url, driver=driver)
       text_gen = "\n".join(summaries)

    if len(text_gen) == 0:
        close_browser(driver)
        return None

    links = scrape_links_with_selenium(driver, url)
    # Limit links to 5
    if len(links) > 5:
        links = links[:5]
    close_browser(driver)
    return f"Answer gathered from website: {text_gen}\n\nLinks: {links}"

def web_search(question:str, search_num:Optional[int]=3) -> List[Dict]:
    search_result = google_official_search(query=question, num_results=search_num+3)
    search_summaries = []

    index = 0
    for res in search_result:
        print(f"{res['title']} x {res['url']}\n")
        try:
            summary = browse_website(url=res['url'], question=question)
        except Exception as e:
            print(f"Error browsing websit {res['title']} x {res['url']}: {e}")
            continue

        if summary is None:
            continue

        search_summaries.append({
            "relevant_content": summary,
            "title": res['title'],
            "url": res['url']
        })

        print(f"{res['title']} x {res['url']} summary: {summary}\n")
        index += 1
        if index >= search_num:
            break

    return search_summaries

def google_official_search(query: str, num_results: int = 8) -> List[Dict]:
    """Return the results of a Google search using the official Google API

    Args:
        query (str): The search query.
        num_results (int): The number of results to return.

    Returns:
        str: The results of the search.
    """

    from googleapiclient.discovery import build
    from googleapiclient.errors import HttpError

    CFG = GlobalConfig().get()
    search_summaries = []
    try:
        # Get the Google API key and Custom Search Engine ID from the config file
        api_key = CFG.google_api_key
        custom_search_engine_id = CFG.custom_search_engine_id

        # Initialize the Custom Search API service
        service = build("customsearch", "v1", developerKey=api_key)

        # Send the search query and retrieve the results
        result = (
            service.cse()
            .list(q=query, cx=custom_search_engine_id, num=num_results)
            .execute()
        )

        # Extract the search result items from the response
        search_results = result.get("items", [])
    except HttpError as e:
        # Handle errors in the API call
        error_details = json.loads(e.content.decode())

        # Check if the error is related to an invalid or missing API key
        if error_details.get("error", {}).get(
            "code"
        ) == 403 and "invalid API key" in error_details.get("error", {}).get(
            "message", ""
        ):
            raise "Error: The provided Google API key is invalid or missing."
        else:
            raise f"Error: {e}"

    # Return the list of search result URLs
    for search_result in search_results:
        search_summaries.append({
            "relevant_content": search_result['snippet'],
            "title": search_result['title'],
            "url": search_result['link']
        })
    return search_summaries
