from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph, END
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_openai import ChatOpenAI
import requests
import base64
import json
import io
import contextlib
import os
import subprocess
import tempfile

OPENAI = os.getenv("OPENAI")

def replace(_, new_value):
    return new_value

search_tool = DuckDuckGoSearchRun()

tools = [search_tool]


llm = ChatOpenAI(
    model="gpt-4o",
    openai_api_key=OPENAI,
)
llm_with_tools = llm.bind_tools(tools)


# Generate the AgentState and Agent graph
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    question: Annotated[str, replace]
    task_id: Annotated[str, replace]
    file_name: Annotated[str, replace]

def get_question(state: AgentState):
    print("get_question")
    question_dict = json.loads(state['messages'][-1].content)
    print(question_dict)
    state['question'] = question_dict["question"]
    state['task_id'] = question_dict["task_id"]
    state['file_name'] = question_dict["file_name"]
    return state


def reformulate_question(state: AgentState):
    print("reformulate_question")
    sys_msg = SystemMessage(content=f"You are a detective and you have to answer a hard question. The first step is to understand the question. Can you reformulate it to make it clear and concise ? The question is : {state['question']}")
    
    state['messages'] = state['messages'] + [llm_with_tools.invoke([sys_msg])]
    return state

def assistant(state: AgentState):
    print("assistant")
    # System message
    sys_msg="""
You are a helpful assistant that can answer questions about the world. You can use the internet through the search tool to find information.
"""
    sys_msg = SystemMessage(content=sys_msg)
    state['messages'] = state['messages'] + [llm_with_tools.invoke([sys_msg] + state["messages"])]
    return state

def get_final_answer(state: AgentState):
    print("get_final_answer")
    sys_msg = SystemMessage(content=f"""
Reply the answer and only the answer of the question. Do not make a sentence, just the answer. If the answer is a number, return it in numeric form.
If the question specifies a format, please return it in the specified format. Do not add unnecessary uppercasing or punctuation. Here is the question again: {state['question']} and here is the answer: {state['messages'][-1].content}
""")
    final_answer = state["messages"][-1].content  # The last assistant message
    print(f"final answer: {final_answer}")
    state['messages'] = state['messages'] + [llm.invoke([sys_msg])]
    return state

    
def python_interpreter(state: AgentState):
    print("python interpreter")

    question_id = state['task_id']
    url = f"https://agents-course-unit4-scoring.hf.space/files/{question_id}"
    
    # Step 1: Download the code
    response = requests.get(url)
    python_code = response.content.decode("utf-8")

    print("Running user code with subprocess...")

    # Step 2: Write code to temp file
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as tmp_file:
        tmp_file.write(python_code)
        tmp_file_path = tmp_file.name

    try:
        # Step 3: Run using subprocess (set timeout if needed)
        result = subprocess.run(
            ["python3", tmp_file_path],
            capture_output=True,
            text=True,
            timeout=30  # seconds
        )
        if result.returncode != 0:
            output = f"⚠️ Error:\n{result.stderr.strip()}"
        else:
            output = result.stdout.strip() or "Code ran but produced no output."
    except subprocess.TimeoutExpired:
        output = "Execution timed out."
    except Exception as e:
        output = f"Unexpected error: {e}"

    # Step 4: Update agent state
    state['messages'] = state['messages'] + [
        AIMessage(content=f"The output of the Python code is:\n```\n{output}\n```")
    ]
    
    return state



# def python_interpreter(state: AgentState):
#     question_id = state['task_id']
#     url = f"https://agents-course-unit4-scoring.hf.space/files/{question_id}"

#     # Step 1: Fetch Python code as bytes and decode
#     response = requests.get(url)
#     python_code = response.content.decode("utf-8")
#     # print(python_code)

#     # Step 2: Capture stdout
#     stdout = io.StringIO()
#     with contextlib.redirect_stdout(stdout):
#         try:
#             exec(python_code)  # Use empty dict for isolated global context
#         except Exception as e:
#             output = f"Error during execution: {e}"
#         else:
#             output = stdout.getvalue().strip()
#     print("Code output : ", output)
#     # Step 3: Save output as message
#     state['messages'] = state['messages'] + [
#         AIMessage(content=f"The output of the python code is : {output}")
#     ]
#     return state

def image_interpreter(state: AgentState):
    print("image interpreter")
    """Answers a question about an image. Provide raw image bytes and a question."""
    question_id = state['task_id']
    url = f"https://agents-course-unit4-scoring.hf.space/files/{question_id}"
    image = requests.get(url).content
    image_base64 = base64.b64encode(image).decode("utf-8")

    messages=[
        {"role": "system", "content": "You are a helpful assistant that can answer questions based on images."},
        {
            "role": "human",
            "content": [
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}},
                {"type": "text", "text": state['question']}
            ],
        },
    ]
    resp = llm.invoke(messages)
    print("Image response : ", resp)
    state['messages'] = state['messages'] + [resp]

    return state

def not_handled(state: AgentState):
    print("not_handled")
    state['messages'] = state['messages'] + [AIMessage(content="I cannot handle this question.")]
    return state


def tools_condition(state: AgentState) -> str:
    print("tools_condition")
    last_msg = state["messages"][-1]

    if "tool_calls" in last_msg.additional_kwargs:
        # LLM is requesting a tool
        return "tools"

    # No tool requested → go to final summary
    return "get_final_answer"

def files_condition(state: AgentState) -> str:
    print("files_condition")
    if state["file_name"].endswith(".py"):
        return "python_interpreter"
    elif state["file_name"].endswith(".png"):
        return "image_interpreter"
    elif state["file_name"]=="":
        return "reformulate_question"
    else:
        return "not_handled"


def build_agent():
    ## The graph
    builder = StateGraph(AgentState)
    
    # Define nodes: these do the work
    builder.add_node("get_question", get_question)
    builder.add_node("reformulate_question", reformulate_question)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))
    builder.add_node("get_final_answer", get_final_answer)
    builder.add_node("python_interpreter", python_interpreter)
    builder.add_node("image_interpreter", image_interpreter)
    builder.add_node("not_handled", not_handled)
    
    # Define edges: these determine how the control flow moves
    builder.add_edge(START, "get_question")
    builder.add_conditional_edges(
        "get_question",
        files_condition,
        {
            "python_interpreter": "python_interpreter",
            "image_interpreter": "image_interpreter",
            "reformulate_question": "reformulate_question",
            "not_handled": "not_handled"
        }
    )
    builder.add_edge("python_interpreter", "get_final_answer")
    builder.add_edge("image_interpreter", "get_final_answer")
    builder.add_edge("reformulate_question", "assistant")
    builder.add_conditional_edges(
        "assistant",
        tools_condition,
        {
            "tools": "tools",
            "get_final_answer": "get_final_answer"
        }
    )
    builder.add_edge("tools", "assistant")
    builder.add_edge("get_final_answer", END)
    builder.add_edge("not_handled", END)
    mySuperAgent = builder.compile()
    return mySuperAgent