import json
from dotenv import load_dotenv
from pathlib import Path
from datetime import datetime

from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import ToolNode
from langgraph.prebuilt import tools_condition
import requests

from agent_state import AgentState
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI
from logging_config import setup_logging, get_logger
from prompts import ASSISTANT_PROMPT

load_dotenv()

# Initialize logging FIRST (before importing agent_tools)
# Create logs folder and generate timestamped log file
logs_dir = Path("logs")
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
log_file = logs_dir / f"agent_run_{timestamp}.log"

setup_logging(log_level="INFO", log_file=str(log_file))
logger = get_logger(__name__)

# Now import agent_tools after logging is configured
from agent_tools import download_file, query_spreadsheet, query_media_file, web_search

logger.info("🚀 Starting agent...")
logger.info(f"📁 Logs will be saved to: {log_file}")

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
questions_url = f"{DEFAULT_API_URL}/random-question"

llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
tools = [web_search, download_file, query_spreadsheet, query_media_file]
llm_with_tools = llm.bind_tools(tools, parallel_tool_calls=False)


sys_msg = SystemMessage(content=ASSISTANT_PROMPT)


def assistant(state: AgentState) -> AgentState:
    logger.info("🤖 Assistant function called")
    logger.info(f"🤖 Last message type: {type(state['messages'][-1])}")
    logger.info(f"🤖 Last message: {state['messages'][-1]}")
    
    result = llm_with_tools.invoke([sys_msg] + state["messages"])

    logger.info(f"🤖 Assistant response: {result.content[:100]}...")
    return {"messages": [result]}


graph = StateGraph(AgentState)

graph.add_node("assistant", assistant)
graph.add_node("tools", ToolNode(tools))

graph.add_edge(START, "assistant")
graph.add_conditional_edges(
    "assistant",
    tools_condition,
)
graph.add_edge("tools", "assistant")

gaia_hf_agent = graph.compile()