import requests
from bs4 import BeautifulSoup
import re
import wikipedia
import os
from typing import Dict, List, Any, Optional
from duckduckgo_search import DDGS

from base_agent import BaseAgent

class ResearchAgent(BaseAgent):
    """
    Agent for handling factual questions by performing web searches, with enhanced structured data scraping.
    """
    def __init__(self, model: str):
        super().__init__(model=model)

    def __call__(self, question: str, file_name: str = None, **kwargs) -> str:
        print(f"Research agent processing question: {question[:100]}...")

        # First extract the search query using LLM
        search_query = self.extract_search_query(question)
        # print(f"Extracted search query: {search_query}")

        # Then try to search Wikipedia
        answer = self.try_structured_wikipedia_scrape(search_query)
        # print(f"Structured answer: {answer}")
        if not answer:
            # If structured answer is not found, perform a web search
            search_results = self.search_web(search_query)
            answer = self.format_search_results(search_results)


        # Now over to LLM again to interpret the result and give final answer

        system_prompt = """
        You are a research expert tasked with answering factual questions based on search results.
        Analyze the search results carefully and extract the information needed to answer the question accurately.
        If the search results don't contain enough information, indicate what's missing.
        Provide a concise, accurate answer based only on the provided information.
        You must provide only the exact answer requested by the user, in the correct type and minimal format, based purely on the provided information.

        YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
        If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
        If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
        If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.

        If you cannot find the answer based on the provided information, reply only with "Unknown".
        """
        
        user_prompt = f"""
        Question: {question}
        
        Search results:
        {answer}
        
        Based on these search results, provide a concise and accurate answer to the question.
        """

        final_answer = self.llm_call(system_prompt, user_prompt)

        print(f"Final answer: {final_answer}")

        if final_answer.lower() == "unknown":
            final_answer = self.ask_james(question)
            
        return final_answer

    def try_structured_wikipedia_scrape(self, query: str) -> Optional[str]:
        try:
            # topic = self.extract_search_query(query)  # <-- Use your LLM-based extractor here
            # if not topic:
            #     print("No topic extracted.")
            #     return None

            print(f"Structured scrape attempt for topic: {query}")

            page_title = wikipedia.search(query, results=1)
            if not page_title:
                print("No Wikipedia page found.")
                return None

            page = wikipedia.page(page_title[0])
            page_url = page.url
            html = requests.get(page_url).text
            soup = BeautifulSoup(html, "html.parser")

            target_table = self.find_relevant_table(soup, query)
            if not target_table:
                print("No relevant table found.")
                return None

            extracted_items = self.parse_table_to_json(target_table)
            if not extracted_items:
                return "No relevant structured data found matching the question."

            answer = "Here is the structured data extracted from Wikipedia:\n\n"
            for item in extracted_items:
                answer += f"- {item}\n"

            return answer

        except Exception as e:
            print(f"Structured scrape failed: {e}")
            return None

    def find_relevant_table(self, soup: BeautifulSoup, question: str) -> Optional[BeautifulSoup]:
        """
        Find the most relevant table based on the question type.
        """
        print("Finding relevant table in " + question)
        target_keywords = []
        if "studio albums" in question.lower():
            target_keywords = ["Studio albums"]
        elif "album" in question.lower():
            target_keywords = ["Discography"]
        elif "book" in question.lower():
            target_keywords = ["Bibliography", "Books"]
        elif "battle" in question.lower():
            target_keywords = ["Battles", "Military history"]
        else:
            # fallback to first wikitable
            tables = soup.find_all("table", {"class": "wikitable"})
            return tables[0] if tables else None

        print(f"Target keywords for table search: {target_keywords}")

        # headings = soup.find_all(['h2', 'h3'])
        headings = soup.find_all(lambda tag: tag.name in ["h2", "h3"] and any(keyword.lower() in tag.get_text(strip=True).lower() for keyword in target_keywords))
        # print(f"Found {len(headings)} headings in the page.")
        for heading in headings:
            heading_text = heading.get_text()
            # print(f"Checking heading: {heading_text}")
            for keyword in target_keywords:
                if keyword.lower() in heading_text.lower():
                    print(f"Found relevant heading: {heading_text}")
                    next_sibling = heading.find_next_sibling()
                    while next_sibling:
                        # Move to the parent div
                        heading_div = heading.find_parent("div")

                        if heading_div:
                            next_sibling = heading_div.find_next_sibling()

                            while next_sibling:
                                if next_sibling.name == "table" and "wikitable" in next_sibling.get("class", []):
                                    print(f"Found relevant table under heading: {heading_text}")
                                    return next_sibling
                                next_sibling = next_sibling.find_next_sibling()

        return None

    def parse_table_to_json(self, table: BeautifulSoup) -> List[Dict[str, Any]]:
        """
        Generic parser for any Wikipedia 'wikitable' to JSON.
        Keys are extracted from <th> headers.
        Values from <td> cells.
        """
        result = []

        # Extract headers
        headers = []
        header_row = table.find("tr")
        if not header_row:
            return result

        for th in header_row.find_all("th"):
            header_text = th.get_text(strip=True)
            headers.append(header_text)

        # Parse data rows
        for row in table.find_all("tr")[1:]:  # skip header
            cols = row.find_all("td")
            if not cols or len(cols) != len(headers):
                continue

            row_data = {}
            for header, col in zip(headers, cols):
                value = col.get_text(separator=" ", strip=True)

                if header.lower() == "album details":
                    # Try to get album title
                    album_link = col.find("a")
                    if album_link:
                        album_title = album_link.get_text(strip=True)
                    else:
                        # Get plain text up to first <ul> if exists
                        ul = col.find("ul")
                        if ul:
                            # Extract text before the <ul>
                            raw_text = "".join(
                                [str(item) for item in col.contents if getattr(item, 'name', None) != "ul"]
                            ).strip()

                            if '<' in raw_text and '>' in raw_text:
                                album_title = BeautifulSoup(raw_text, "html.parser").get_text(strip=True)
                            else:
                                album_title = raw_text.strip()
                        else:
                            album_title = col.get_text(strip=True)

                    # Extract extras from <ul>
                    extras = []
                    ul = col.find("ul")
                    if ul:
                        extras = [li.get_text(strip=True) for li in ul.find_all("li")]

                    # Build structured value
                    if album_title and extras:
                        value = {
                            "title": album_title,
                            "extra": extras
                        }
                    elif album_title:
                        value = {
                            "title": album_title
                        }
                    elif extras:
                        value = extras

                row_data[header] = value

            result.append(row_data)

        return result
    
    def search_web(self, query: str, max_results: int = 5) -> List[Dict[str, str]]:
        """
        Perform a web search using DuckDuckGo.
        
        Args:
            query: The search query
            max_results: Maximum number of results to return
            
        Returns:
            List of search results with title, content, and url
        """
        results = []
        DUCKDUCKGO_ENABLED = os.getenv("DUCKDUCKGO_ENABLED")
        
        try:    
            print("Using DuckDuckGo to conduct a web search")
            # Then use DuckDuckGo if enabled and we need more results
            if DUCKDUCKGO_ENABLED and len(results) < max_results:
                with DDGS() as ddgs:
                    ddgs_results = list(ddgs.text(query, max_results=max_results))
                    for r in ddgs_results:
                        results.append({
                            "title": r.get("title", ""),
                            "content": r.get("body", ""),
                            "url": r.get("href", "")
                        })
        except Exception as e:
            print(f"Web search error: {str(e)}")
            
        return results[:max_results]
    
    def format_search_results(self, results: List[Dict[str, str]]) -> str:
        """
        Format search results for the LLM.
        
        Args:
            results: List of search results
            
        Returns:
            Formatted search results as a string
        """
        if not results:
            return "No search results found."
            
        formatted = ""
        for i, result in enumerate(results, 1):
            formatted += f"Result {i}:\n"
            formatted += f"Title: {result.get('title', 'No title')}\n"
            formatted += f"URL: {result.get('url', 'No URL')}\n"
            formatted += f"Content: {result.get('content', 'No content')}\n\n"
            
        return formatted