"""LangGraph: agent graph w/ tools"""
import os
from dotenv import load_dotenv
from typing import List, Dict, Any, Optional
import tempfile
import re
import json
import requests
from urllib.parse import urlparse
import pytesseract
from PIL import Image, ImageDraw, ImageFont, ImageEnhance, ImageFilter
import cmath
import pandas as pd
import uuid
import numpy as np


""" Langchain imports"""
from langgraph.graph import START, StateGraph, MessagesState
from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_core.tools import tool
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader
#from langchain_community.vectorstores import SupabaseVectorStore
from langchain_google_genai import ChatGoogleGenerativeAI
#from langchain.tools.retriever import create_retriever_tool
#from supabase.client import Client, create_client
#from code_interpreter import CodeInterpreter
#interpreter_instance = CodeInterpreter()
#from image_processing import *

"""
import getpass
import os

if "GOOGLE_API_KEY" not in os.environ:
    os.environ["GOOGLE_API_KEY"] = getpass.getpass("Enter your Google AI API key: ")
"""

load_dotenv()

@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 power(a: float, b: float) -> float:
    """
    Get the power of two numbers.
    Args:
        a (float): the first number
        b (float): the second number
    """
    return a**b

@tool
def square_root(a: float) -> float | complex:
    """
    Get the square root of a number.
    Args:
        a (float): the number to get the square root of
    """
    if a >= 0:
        return a**0.5
    return cmath.sqrt(a)


@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return maximum 2 results.
    
    Args:
        query: The search query."""
    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in search_docs
        ])
    return {"wiki_results": formatted_search_docs}

@tool
def web_search(query: str) -> str:
    """Search Tavily for a query and return maximum 3 results.
    
    Args:
        query: The search query."""

    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in search_docs
        ])
    return {"web_results": formatted_search_docs}




# load the system prompt from the file
with open("system_prompt.txt", "r", encoding="utf-8") as f:
    system_prompt = f.read()

# System message
sys_msg = SystemMessage(content=system_prompt)


tools = [
    multiply,
    add,
    subtract,
    divide,
    modulus,
    power,
    square_root,
    wiki_search,
    web_search,
]

# Build graph function
def build_graph(provider: str = "huggingface"):
    """Build the graph"""
    # Load environment variables from .env file
    if provider == "huggingface":
        # Huggingface endpoint
        """
        llm = ChatHuggingFace(
            llm=HuggingFaceEndpoint(
                #endpoint_url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
                #endpoint_url="https://api-inference.huggingface.co/models/Qwen/Qwen3-30B-A3B",
                endpoint_url="https://api-inference.huggingface.co/models/Qwen/Qwen2.5-Coder-32B.Instruct",
                #endpoint_url="https://api-inference.huggingface.co/models/Qwen/Qwen3-4B",
                temperature=0,
            ),
        )
        """
        llm = ChatHuggingFace(
            llm=HuggingFaceEndpoint(
                repo_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
                #endpoint_url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
                #endpoint_url="https://api-inference.huggingface.co/models/microsoft/phi-4",
                #endpoint_url="https://api-inference.huggingface.co/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0",
                task="text-generation",  # for chat‐style use “text-generation”
                #max_new_tokens=1024,
                #do_sample=False,
                #repetition_penalty=1.03,
                temperature=0,
            ),
            #verbose=True,
        )

    elif provider == "google":
        # Google Gemini
        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
        #llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)
        
    else:
        raise ValueError("Invalid provider. Choose 'huggingface'.")
        
    # Bind tools to LLM
    llm_with_tools = llm.bind_tools(tools)
    
    # Node
    def assistant(state: MessagesState):
        """Assistant node"""      
        return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}
        
    
    #def retriever(state: MessagesState):
    #    """Retriever node"""
    #    return {"messages": [sys_msg] + state["messages"]}
   

    builder = StateGraph(MessagesState)
    #builder.add_node("retriever", retriever)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))
    #builder.add_edge(START, "retriever")
    builder.add_edge(START, "assistant")
    #builder.add_edge("retriever", "assistant")
    builder.add_conditional_edges(
        "assistant",
        tools_condition,
    )
    #builder.add_edge("tools", "retriever")
    builder.add_edge("tools", "assistant")

    # Compile graph
    return builder.compile()

# test
if __name__ == "__main__":
    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
    # Build the graph
    graph = build_graph(provider="huggingface")
    # Run the graph
    messages = [HumanMessage(content=question)]
    messages = graph.invoke({"messages": messages})
    for m in messages["messages"]:
        m.pretty_print()