import base64
from typing import List, TypedDict, Annotated, Optional
from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
from langgraph.graph.message import add_messages
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
#from IPython.display import Image, display
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.tools.tavily_search import TavilySearchResults
#from langchain_community.tools.youtube_search import YoutubeSearchResults
from langchain_core.tools import tool

from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# Validate API key
if not OPENAI_API_KEY:
    print("Warning: OPENAI_API_KEY environment variable not found or empty!")

class AgentState(TypedDict):
    # The document provided
    question_text: Optional[str]
    messages: Annotated[list[AnyMessage], add_messages]

### TOOLS ###
@tool
def search_wikipedia(query: str) -> str:
    """Searches Wikipedia for a given query and returns maximum 1 results."""
    try:
        # Ensure query is properly encoded
        if isinstance(query, bytes):
            query = query.decode('utf-8')
        
        # Use WikipediaLoader with the encoded query
        search_docs = WikipediaLoader(
            query=query, 
            load_max_docs=1, 
            doc_content_chars_max=10000
        ).load()
        
        # Format the results with proper encoding
        formatted_docs = []
        for doc in search_docs:
            source = doc.metadata.get("source", "").encode('utf-8', errors='replace').decode('utf-8')
            page = doc.metadata.get("page", "").encode('utf-8', errors='replace').decode('utf-8')
            content = doc.page_content.encode('utf-8', errors='replace').decode('utf-8')
            formatted_docs.append(
                f'<Document source="{source}" page="{page}"/>\n{content}\n</Document>'
            )
        
        return "\n\n---\n\n".join(formatted_docs)
    except Exception as e:
        print(f"Error in search_wikipedia: {e}")
        return f"Error searching Wikipedia: {str(e)}"

@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers.
    Args:
        a: first int
        b: second int
    """
    return a * b

@tool
def add(a: int, b: int) -> int:
    """Add two numbers.
    
    Args:
        a: first int
        b: second int
    """
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """Subtract two numbers.
    
    Args:
        a: first int
        b: second int
    """
    return a - b

@tool
def divide(a: int, b: int) -> int:
    """Divide two numbers.
    
    Args:
        a: first int
        b: second int
    """
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b

@tool
def modulus(a: int, b: int) -> int:
    """Get the modulus of two numbers.
    
    Args:
        a: first int
        b: second int
    """
    return a % b

@tool
def web_search(query: str) -> str:
    """Search the web for a query and return maximum 3 results.
    
    Args:
        query: The search query."""
    search_tool = TavilySearchResults(max_results=3)
    search_results = search_tool.invoke(query)
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{result["url"]}" title="{result.get("title", "")}"/>\n{result["content"]}\n</Document>'
            for result in search_results
        ])
    return formatted_search_docs

# Equip the butler with tools
tools = [
    multiply,
    subtract,
    add,
    divide,
    modulus,
    search_wikipedia,
    web_search
]

# Initialize the OpenAI model with explicit encoding parameters
try:
    # Use encoding='utf-8' to ensure proper handling of special characters
    llm = ChatOpenAI(
        model="gpt-4o-mini",
        openai_api_key=OPENAI_API_KEY,
        default_headers={"Content-Type": "application/json; charset=utf-8"}
    )
    print("Binding tools...")
    llm_with_tools = llm.bind_tools(tools, parallel_tool_calls=False)
except Exception as e:
    print(f"Error initializing OpenAI model: {e}")
    raise

#### THE NODES #####
def assistant(state: AgentState):
    try:
        # Get the question text from state and ensure it's properly encoded
        question = state.get("question_text", "")
        
        # Build the system prompt piece by piece
        parts = [
            "You are a helpful assistant tasked with answering questions using a set of tools.",
            "Report your thoughts step by step.",
            "Once you have determined the answer, provide it in the following XML format:",
            "<FINAL_ANSWER>Your concise final answer here</FINAL_ANSWER>",
            "Your final answer should be as concise as possible - a number OR as few words as possible OR a comma separated list of numbers and/or strings.",
            "If asked for a number, don't use commas or units such as $ or percent sign unless specified otherwise.",
            "If asked for a string, don't use articles or abbreviations (e.g. for cities), and write digits in plain text unless specified otherwise.",
            "The answer inside <FINAL_ANSWER> tags must be complete, clear, and not truncated.",
            
            # Tool descriptions
            "You can use these provided tools:",
            "search_wikipedia(query: str) -> str:",
            "    Searches Wikipedia for a given query and returns maximum 1 results.",
            "web_search(query: str) -> str:",
            "    Searches the web for a query and returns maximum 3 results.",
            "multiply, add, subtract, divide, modulus: Basic math operations between two numbers."
        ]
        
        # Add the question
        parts.append(f"The question you have to answer is: {question}")
        
        # Join all parts with newlines and ensure proper encoding
        sys_prompt = "\n".join(parts).encode('utf-8', errors='replace').decode('utf-8')
        
        # Create system message with explicit encoding
        sys_msg = SystemMessage(content=sys_prompt)
        
        # Invoke the model with proper error handling
        messages = llm_with_tools.invoke([sys_msg] + state["messages"])
        return {
            "messages": messages,
            "question_text": question
        }
    except Exception as e:
        error_msg = f"Error in assistant node: {str(e)}"
        print(error_msg)
        return {
            "messages": [SystemMessage(content=f"An error occurred: {error_msg}")],
            "question_text": state.get("question_text", "")
        }

def build_graph():
    # The graph
    builder = StateGraph(AgentState)

    # Define nodes: these do the work
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))

    # Define edges: these determine how the control flow moves
    builder.add_edge(START, "assistant")
    builder.add_conditional_edges(
        "assistant",
        # If the latest message requires a tool, route to tools
        # Otherwise, provide a direct response
        tools_condition,
    )
    builder.add_edge("tools", "assistant")
    react_graph = builder.compile()
    return react_graph



class WalidzAgent:
    """A langgraph agent."""
    def __init__(self):
        print("BasicAgent initialized.")
        self.graph = build_graph()

    def __call__(self, question: str) -> str:
        if isinstance(question, dict) and "messages" in question and "question_text" in question:
            # We got a properly formatted input
            input_state = question
        else:
            # Create a proper input state from the question string
            human_msg = HumanMessage(content=question)
            input_state = {
                "messages": [human_msg],
                "question_text": question
            }

        # Stream the execution and capture final step
        final_state = None
        for step in self.graph.stream(input_state):
            print("--- STEP ---")
            print(step)
            final_state = step
        
        # Extract the answer from the final state
        if final_state:
            # Extract message content from the appropriate structure
            message_content = None
            
            # Check if the result has the nested 'assistant' structure
            if 'assistant' in final_state and 'messages' in final_state['assistant']:
                # Get message from assistant's messages field
                message = final_state['assistant']['messages']
                message_content = message.content
            # Fallback to direct messages list if available
            elif 'messages' in final_state:
                messages_list = final_state['messages']
                if messages_list and len(messages_list) > 0:
                    message_content = messages_list[-1].content
            
            if not message_content:
                print("Warning: Could not extract message content from response")
                return "Failed to extract message content from agent response"
            
            # First try XML tag extraction
            import re
            xml_pattern = re.compile(r'<FINAL_ANSWER>(.*?)</FINAL_ANSWER>', re.DOTALL)
            match = xml_pattern.search(message_content)
            
            if match:
                # Found XML tags
                answer = match.group(1).strip()
            else:
                # Fall back to "FINAL ANSWER:" format
                if "FINAL ANSWER:" in message_content:
                    # Split by "FINAL ANSWER:" and take the last part
                    final_answer_parts = message_content.split("FINAL ANSWER:")
                    answer = final_answer_parts[-1].strip()
                else:
                    # If no explicit final answer format is found, return the whole message
                    print("Warning: No final answer format found in response")
                    answer = message_content
            
            print(f"Extracted final answer: {answer}")
            return answer
            
        # Fallback if no valid answer found
        print("Warning: Could not extract answer from agent response")
        return "No valid answer could be extracted from agent response"
    
if __name__ == "__main__":

    question_text = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
    question_text = "What country had the least number of athletes at the 1928 Summer Olympics? If there's a tie for a number of athletes, return the first in alphabetical order. Give the IOC country code as your answer."

    agent = WalidzAgent()
    # Prepare input state
    human_msg = HumanMessage(content=question_text)
    input_state = {
        "messages": [human_msg],
        "question_text": question_text
    }

    agent(input_state)


    # Save the graph visualization to disk
    #graph_image = react_graph.get_graph(xray=True).draw_mermaid_png()
    #with open("graph_visualization.png", "wb") as f:
    #    f.write(graph_image)
    #print("Graph visualization saved to 'graph_visualization.png'")

    # Also display it if in an environment that supports display
    #display(Image(graph_image))


    #question_sample = question_sample.encode('utf-8', errors='replace').decode('utf-8')


    """print("Running with sample question:", question_sample)

    try:
        # Create message with proper encoding
        human_msg = HumanMessage(content=question_sample)
        
        # Prepare input state
        input_state = {
            "messages": [human_msg],
            "question_text": question_sample
        }
        
        # Invoke graph with proper encoding handling
        print("Invoking graph...")
        messages = react_graph.invoke(input_state)
        
        # Show the messages
        for m in messages['messages']:
            m.pretty_print()
    except Exception as e:
        print(f"Error in Agent execution: {str(e)}")
        import traceback
        traceback.print_exc()"""