from typing import List, TypedDict, Annotated, Optional
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 langchain_groq import ChatGroq
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.tools import tool

from langchain_community.tools import DuckDuckGoSearchRun
from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper


import requests, io, pandas as pd, PyPDF2, ast, pytesseract
from PIL import Image


class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]


@tool
def add(a: float, b: float) -> float:
    """
    Adds two numbers.

    Args:
        a (float): First number.
        b (float): Second number.

    Returns:
        float: Sum of a and b.
    """
    return a + b

@tool
def subtract(a: float, b: float) -> float:
    """
    Subtracts one number from another.

    Args:
        a (float): Minuend.
        b (float): Subtrahend.

    Returns:
        float: Result of a - b.
    """
    return a - b

@tool
def multiply(a: float, b: float) -> float:
    """
    Multiplies two numbers.

    Args:
        a (float): First number.
        b (float): Second number.

    Returns:
        float: Product of a and b.
    """
    return a * b

@tool
def divide(a: float, b: float) -> float:
    """
    Divides one number by another.

    Args:
        a (float): Dividend.
        b (float): Divisor.

    Raises:
        ValueError: If b is zero.

    Returns:
        float: Result of a / b.
    """
    if b == 0:
        raise ValueError("Cannot divide by zero")
    return a / b

@tool
def web_search(query: str) -> str:
    """
    Searches the web for the given query.

    Args:
        query (str): The search query.

    Returns:
        str: The search results
    """
    search = DuckDuckGoSearchRun()
    print("search")
    return search.invoke(query)

@tool
def wikisearch(query: str) -> str:
    """
    Searches wikipedia for the given query.

    Args:
        query (str): The search query.

    Returns:
        str: The wikipedia results.
    """
    wikipedia = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
    print("wiki")
    return wikipedia.run(query)


@tool
def gaia_retriever_tool(task_id: str) -> str:
    """
    Downloads a file from a Gaia task using the provided task_id and returns its raw text content.
    Does NOT analyze or interpret the content.
    Supports PDFs, Excel/CSV spreadsheets, Python scripts, and images.
    """
    print("retriever")
    url = f'https://agents-course-unit4-scoring.hf.space/files/{task_id}'
    resp = requests.get(url)
    if resp.status_code != 200:
        return f"Failed to download file {task_id}"

    file_bytes = io.BytesIO(resp.content)
    text = ""

    try:
        df = pd.read_excel(file_bytes)
        text = df.to_csv(index=False)
    except Exception:
        try:
            file_bytes.seek(0)
            df = pd.read_csv(file_bytes)
            text = df.to_csv(index=False)
        except Exception:
            try:
                file_bytes.seek(0)
                reader = PyPDF2.PdfReader(file_bytes)
                text = "\n".join([p.extract_text() or "" for p in reader.pages])
            except Exception:
                try:
                    file_bytes.seek(0)
                    img = Image.open(file_bytes)
                    text = pytesseract.image_to_string(img)
                except Exception:
                    file_bytes.seek(0)
                    text = file_bytes.read().decode("utf-8", errors="ignore")
    return text

tools = [add, subtract, multiply, divide, web_search, wikisearch, gaia_retriever_tool]

llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    api_key="...",
    temperature=0.3,
)

llm_with_tools = llm.bind_tools(tools)

system_prompt = SystemMessage(
    content="""
You are a concise AI assistant. Use tools only when strictly necessary.

**Math:** use add(a,b), subtract(a,b), multiply(a,b), divide(a,b).  
**Search:** 
- Always try wikisearch(query) first.
- If wiki fails or is insufficient, use web_search(query).
- Only fallback to general knowledge if both fail.

**Files:** use gaia_retriever_tool(task_id) **whenever the question mentions a file, attachment, picture, Excel, CSV, PDF, or Python file.**  
- Variations like "attached Excel file", "following CSV file", or "image attached" count as explicit mentions.  
- **The agent should pass the task_id from the question to gaia_retriever_tool.**  
- **The tool itself handles downloading and extracting content.**  
- **Do NOT attempt to download or read the file outside the tool.**  
- **Do NOT call gaia_retriever_tool if no file is mentioned.**

**Answer rules:**
1. Give **exact answer only**, no explanations.  
2. Numbers → output only the number.
3. Lists → comma-separated values, no words.  
4. Text → 1–5 words.  
5. Conceptual questions (opposites, synonyms, meanings) → answer correctly in 1 word.  
6. Never include extra context, tool names, or sentences.
7. Your answer must strictly be short

**Example 1:**  
Question: Who is the president of France?  
Action: wikisearch("president of France") → returns "Emmanuel Macron"  
Answer: Emmanuel Macron

**Example 2:**  
Question: Some obscure topic, task_id: 123  
Action: wikisearch("Some obscure topic") → returns ""  
Since wiki returned nothing, call web_search("Some obscure topic")  
Answer: [result from web_search]  
**Do NOT call gaia_retriever_tool because the question does not mention a file.**

**Example 3:**  
Question: The attached Excel file contains sales data, task_id: 7bd855d8-463d-4ed5-93ca-5fe35145f733  
Action: gaia_retriever_tool("7bd855d8-463d-4ed5-93ca-5fe35145f733") → returns file content  
Answer: 30420.00
"""
)


def assistant(state: AgentState):
    return {
        "messages": [llm_with_tools.invoke([system_prompt] + state["messages"])],
    }


builder = StateGraph(AgentState)
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
builder.add_edge(START, "assistant")
builder.add_conditional_edges("assistant", tools_condition)
builder.add_edge("tools", "assistant")
graph = builder.compile()

def run(prompt, task_id='1'):
    state: AgentState = {"messages": [HumanMessage(prompt)]}
    output = graph.invoke(state)
    return output["messages"][-1].content