import requests
from geopy.geocoders import Nominatim
from langchain import OpenAI, LLMMathChain, LLMChain, PromptTemplate, Wikipedia
from langchain.agents import Tool
from langchain.agents.react.base import DocstoreExplorer
from langchain_community.document_loaders import TextLoader
from langchain.indexes import VectorstoreIndexCreator
from langchain_community.utilities import SerpAPIWrapper, DuckDuckGoSearchAPIWrapper
from langchain_community.utilities import WolframAlphaAPIWrapper

from nodes.Node import Node


class GoogleWorker(Node):
    def __init__(self, name="Google"):
        super().__init__(name, input_type=str, output_type=str)
        self.isLLMBased = False
        self.description = "Worker that searches results from Google. Useful when you need to find short " \
                           "and succinct answers about a specific topic. Input should be a search query."

    def run(self, input, log=False):
        assert isinstance(input, self.input_type)
        try:
            # Detect if input is in Vietnamese
            import re
            is_vietnamese = bool(re.search(r'[àáạảãâầấậẩẫăằắặẳẵèéẹẻẽêềếệểễìíịỉĩòóọỏõôồốộổỗơờớợởỡùúụủũưừứựửữỳýỵỷỹđ]', input.lower()))
            
            tool = SerpAPIWrapper()
            # Add language hint for better results
            if is_vietnamese:
                input = input + " lang:vi"
            evidence = tool.run(input)
            
            if evidence is None:
                evidence = "Không tìm thấy kết quả." if is_vietnamese else "No results found from Google search."

        except Exception as e:
            evidence = f"Lỗi tìm kiếm Google: {str(e)}" if is_vietnamese else f"Error performing Google search: {str(e)}"

        evidence = str(evidence).strip()
        if log:
            return {"input": input, "output": evidence}
        return evidence

class DuckduckgoWorker(Node):
    def __init__(self, name="Duckduckgo"):
        super().__init__(name, input_type=str, output_type=str)
        self.isLLMBased = False
        self.description = "Worker that searches results from the web. Useful when you need to find short " \
                           "and succinct answers about a specific topic. Input should be a search query."

    def run(self, input, log=False):
        assert isinstance(input, self.input_type)
        try:
            # Detect if input is in Vietnamese
            import re
            is_vietnamese = bool(
                re.search(r'[àáạảãâầấậẩẫăằắặẳẵèéẹẻẽêềếệểễìíịỉĩòóọỏõôồốộổỗơờớợởỡùúụủũưừứựửữỳýỵỷỹđ]', input.lower()))

            tool = DuckDuckGoSearchAPIWrapper(
                backend="html",
                region="vn-vi" if is_vietnamese else "wt-wt",
                time=None,
                max_results=5,
            )
            evidence = tool.run(input)

            if evidence is None:
                evidence = "Không tìm thấy kết quả." if is_vietnamese else "No results found from DDG search."

        except Exception as e:
            evidence = f"Lỗi tìm kiếm DDG: {str(e)}" if is_vietnamese else f"Error performing DDG search: {str(e)}"

        evidence = str(evidence).strip()
        if log:
            return {"input": input, "output": evidence}
        return evidence

class WikipediaWorker(Node):
    def __init__(self, name="Wikipedia", docstore=None):
        super().__init__(name, input_type=str, output_type=str)
        self.isLLMBased = False
        self.description = "Worker that search for similar page contents from Wikipedia. Useful when you need to " \
                           "get holistic knowledge about people, places, companies, historical events, " \
                           "or other subjects. The response are long and might contain some irrelevant information. " \
                           "Input should be a search query."
        self.docstore = docstore

    def run(self, input, log=False):
        if not self.docstore:
            self.docstore = DocstoreExplorer(Wikipedia())
        assert isinstance(input, self.input_type)
        tool = Tool(
            name="Search",
            func=self.docstore.search,
            description="useful for when you need to ask with search"
        )
        evidence = tool.run(input)
        assert isinstance(evidence, self.output_type)
        if log:
            print(f"Running {self.name} with input {input}\nOutput: {evidence}\n")
        return evidence


class DocStoreLookUpWorker(Node):
    def __init__(self, name="LookUp", docstore=None):
        super().__init__(name, input_type=str, output_type=str)
        self.isLLMBased = False
        self.description = "Worker that search the direct sentence in current Wikipedia result page. Useful when you " \
                           "need to find information about a specific keyword from a existing Wikipedia search " \
                           "result. Input should be a search keyword."
        self.docstore = docstore

    def run(self, input, log=False):
        if not self.docstore:
            raise ValueError("Docstore must be provided for lookup")
        assert isinstance(input, self.input_type)
        tool = Tool(
            name="Lookup",
            func=self.docstore.lookup,
            description="useful for when you need to ask with lookup"
        )
        evidence = tool.run(input)
        assert isinstance(evidence, self.output_type)
        if log:
            print(f"Running {self.name} with input {input}\nOutput: {evidence}\n")
        return evidence


class CustomWolframAlphaAPITool(WolframAlphaAPIWrapper):
    def __init__(self):
        super().__init__()

    def run(self, query: str) -> str:
        """Run query through WolframAlpha and parse result."""
        res = self.wolfram_client.query(query)

        try:
            answer = next(res.results).text
        except StopIteration:
            return "Wolfram Alpha wasn't able to answer it"

        if answer is None or answer == "":
            return "No good Wolfram Alpha Result was found"
        else:
            return f"Answer: {answer}"


class WolframAlphaWorker(Node):
    def __init__(self, name="WolframAlpha"):
        super().__init__(name, input_type=str, output_type=str)
        self.isLLMBased = False
        self.description = "A WolframAlpha search engine. Useful when you need to solve a complicated Mathematical or " \
                           "Algebraic equation. Input should be an equation or function."

    def run(self, input, log=False):
        assert isinstance(input, self.input_type)
        tool = CustomWolframAlphaAPITool()
        evidence = tool.run(input).replace("Answer:", "").strip()
        assert isinstance(evidence, self.output_type)
        if log:
            print(f"Running {self.name} with input {input}\nOutput: {evidence}\n")
        return evidence


class CalculatorWorker(Node):
    def __init__(self, name="Calculator"):
        super().__init__(name, input_type=str, output_type=str)
        self.isLLMBased = True
        self.description = "A calculator that can compute arithmetic expressions. Useful when you need to perform " \
                           "math calculations. Input should be a mathematical expression"

    def run(self, input, log=False):
        assert isinstance(input, self.input_type)
        try:
            from langchain_openai import OpenAI
            llm = OpenAI(temperature=0)
            
            # Create math chain with latest components
            from langchain_core.prompts import PromptTemplate
            template = """Given the math problem:
            {question}
            
            Let's solve this step by step:
            1) First, let's understand what we're calculating
            2) Then, perform the calculation
            3) Finally, provide the numeric answer
            
            The answer is:"""
            
            prompt = PromptTemplate(template=template, input_variables=["question"])
            chain = prompt | llm
            
            response = chain.invoke({"question": input})
            evidence = str(response).strip()
            
            # Try to extract just the numeric answer if possible
            import re
            numeric_match = re.search(r'[-+]?\d*\.?\d+', evidence)
            if numeric_match:
                evidence = numeric_match.group()
                
        except Exception as e:
            evidence = f"Error in calculation: {str(e)}"
            
        assert isinstance(evidence, self.output_type)
        if log:
            return {"input": input, "output": evidence}
        return evidence


class LLMWorker(Node):
    def __init__(self, name="LLM"):
        super().__init__(name, input_type=str, output_type=str)
        self.isLLMBased = True
        self.description = "A pretrained LLM like yourself. Useful when you need to act with general world " \
                           "knowledge and common sense. Prioritize it when you are confident in solving the problem " \
                           "yourself. Input can be any instruction."

    def run(self, input, log=False):
        assert isinstance(input, self.input_type)
        try:
            from langchain_openai import OpenAI
            llm = OpenAI(temperature=0)
            import re
            is_vietnamese = bool(re.search(r'[àáạảãâầấậẩẫăằắặẳẵèéẹẻẽêềếệểễìíịỉĩòóọỏõôồốộổỗơờớợởỡùúụủũưừứựửữỳýỵỷỹđ]', input.lower()))
            # Create a more detailed prompt template
            if is_vietnamese:
                template = """Dựa trên yêu cầu sau đây, hãy cung cấp câu trả lời rõ ràng và trực tiếp. 
                Nếu sử dụng thông tin có sẵn, hãy phân tích cẩn thận.
                Nếu không chắc chắn, hãy giải thích những gì bạn biết về chủ đề này.

                Yêu cầu: {request}

                Trả lời: """
            else:
                template = """Based on the following request, provide a clear and direct answer. 
                If using context, analyze it carefully.
                If you don't know or aren't sure, explain what you do know about the topic.

                Request: {request}

                Answer: """

            prompt = PromptTemplate(template=template, input_variables=["request"])
            chain = prompt | llm
            
            # Run the chain and get response
            response = chain.invoke({"request": input})
            evidence = str(response).strip()
            
            # If response is too short or says "I don't know", try to provide more detail
            if len(evidence) < 10 or "don't know" in evidence.lower() or "không thể" in evidence.lower():
                if is_vietnamese:
                    fallback_prompt = f"Hãy phân tích và cung cấp thông tin chi tiết về: {input}"
                else:
                    fallback_prompt = f"Please provide detailed analysis and information about: {input}"
                fallback_response = llm.invoke(fallback_prompt)
                evidence = str(fallback_response).strip()
            
        except Exception as e:
            if is_vietnamese:
                evidence = f"Lỗi xử lý LLM: {str(e)}"
            else:
                evidence = f"Error in LLM processing: {str(e)}"
        
        assert isinstance(evidence, self.output_type)
        if log:
            return {"input": input, "output": evidence}
        return evidence


class ZipCodeRetriever(Node):

    def __init__(self, name="ZipCodeRetriever"):
        super().__init__(name, input_type=str, output_type=str)
        self.isLLMBased = False
        self.description = "A zip code retriever. Useful when you need to get users' current zip code. Input can be " \
                           "left blank."

    def get_ip_address(self):
        response = requests.get("https://ipinfo.io/json")
        data = response.json()
        return data["ip"]

    def get_location_data(sefl, ip_address):
        url = f"https://ipinfo.io/{ip_address}/json"
        response = requests.get(url)
        data = response.json()
        return data

    def get_zipcode_from_lat_long(self, lat, long):
        geolocator = Nominatim(user_agent="zipcode_locator")
        location = geolocator.reverse((lat, long))
        return location.raw["address"]["postcode"]

    def get_current_zipcode(self):
        ip_address = self.get_ip_address()
        location_data = self.get_location_data(ip_address)
        lat, long = location_data["loc"].split(",")
        zipcode = self.get_zipcode_from_lat_long(float(lat), float(long))
        return zipcode

    def run(self, input):
        assert isinstance(input, self.input_type)
        evidence = self.get_current_zipcode()
        assert isinstance(evidence, self.output_type)


class SearchDocWorker(Node):

    def __init__(self, doc_name, doc_path, name="SearchDoc"):
        super().__init__(name, input_type=str, output_type=str)
        self.isLLMBased = True
        self.doc_path = doc_path
        self.description = f"A vector store that searches for similar and related content in document: {doc_name}. " \
                           f"The result is a huge chunk of text related to your search but can also " \
                           f"contain irrelevant info. Input should be a search query."

    def run(self, input, log=False):
        assert isinstance(input, self.input_type)
        loader = TextLoader(self.doc_path)
        vectorstore = VectorstoreIndexCreator().from_loaders([loader]).vectorstore
        evidence = vectorstore.similarity_search(input, k=1)[0].page_content
        assert isinstance(evidence, self.output_type)
        if log:
            print(f"Running {self.name} with input {input}\nOutput: {evidence}\n")
        return evidence


class SearchSOTUWorker(SearchDocWorker):
    def __init__(self, name="SearchSOTU"):
        super().__init__(name=name, doc_name="state_of_the_union", doc_path="data/docs/state_of_the_union.txt")


WORKER_REGISTRY = {"Google": GoogleWorker(),
                   "Wikipedia": WikipediaWorker(),
                   "LookUp": DocStoreLookUpWorker(),
                   "WolframAlpha": WolframAlphaWorker(),
                   "Calculator": CalculatorWorker(),
                   "LLM": LLMWorker(),
                   "SearchSOTU": SearchSOTUWorker(),
                   "Duckduckgo": DuckduckgoWorker()
                   }