#for internet agent
from langchain.agents import create_agent # The modern agent constructor
from langchain_ollama import ChatOllama # for configuring the LLM
from langchain_community.tools import DuckDuckGoSearchRun #for using an online search tool
from langchain_core.messages import HumanMessage # for processing the human prompt

#packages installation command - pip install langgraph langchain-ollama duckduckgo-search ddgs
#ollama 3.2 installation command  - ollama pull llama3.2

def run_internet_agent(user_message):
    SYSTEM_INSTRUCTION = """
    You are a highly flexible and helpful AI research assistant. Your primary goal is to answer the user's question, using the available tools only when necessary.

    **FLEXIBLE FORMATTING RULES:**
    1.  **PRIORITY:** Adhere strictly to any formatting requested by the user (e.g., "return as bullet points," "write a short paragraph," or "output as JSON").
    2.  **LINK MANDATE (Hard Rule):** For every piece of factual information, movie, book, or specific data point you mention that came from the web search tool, you MUST embed the source URL from the search results directly into the text using **Markdown link syntax: [Relevant Text](Source URL)**.
    """
    search_tool = DuckDuckGoSearchRun()
    tools = [search_tool]
    llm = ChatOllama(
    model="llama3.2",
    temperature=0.1, # Keep it low for factual/tool-based tasks
)
# 3. Create the Agent Executor
# This bundles the LLM and the tools into a decision-making loop
    agent_executor = create_agent(llm, tools,system_prompt=SYSTEM_INSTRUCTION)
    
    # Use streaming to show the agent's step-by-step process
    # The agent_executor.stream() output is what we need to safely parse
    for chunk in agent_executor.stream(
        {"messages": [HumanMessage(content=user_message)]},
        stream_mode="values"
    ):
        latest_message = chunk["messages"][-1]
        
        # 1. Check for Tool Call (AI's decision to search)
        # Use hasattr() to safely check for the 'tool_calls' attribute
        if hasattr(latest_message, 'tool_calls') and latest_message.tool_calls:
            
            tool_call_args = [tc['args'] for tc in latest_message.tool_calls]
            print(f"--- 🧠 Thinking: I need to search for: {tool_call_args}")
        
        # 2. Check if the AI has received tool results
        elif latest_message.type == "tool":
            print(f"--- 🔍 Found: {latest_message.content[:100]}... (truncated results)")
        
        # 3. Check for the final answer
        # Check if it's an AI message AND it has no tool calls (or the attribute is missing/empty)
        elif latest_message.type == "ai" and (not hasattr(latest_message, 'tool_calls') or not latest_message.tool_calls):
          return  f"{latest_message.content}"

def submit():
    user_message = st.session_state['user_input']
    if user_message:
        
        if 'active_profile' not in st.session_state or st.session_state.active_profile != username:
            st.session_state.active_profile = username
    # load saved chat for this username and populate session history
        
        if username:
                CHROMA_DIR = f".chroma_db_{username}"
                client, collection = init_chromadb(CHROMA_DIR)
        else:
                CHROMA_DIR = ".chroma_db"
                client, collection = init_chromadb(CHROMA_DIR)

        
        st.session_state.history.append(("user", user_message))
        save_chat_message(client, collection, role="user", text=user_message, user_id=username) # save user's message
        
        result = run_internet_agent(user_message)
        st.session_state.history.append(("agent", result))
        
        save_chat_message(client, collection, role="agent", text=result, user_id=username) #save agent message
        if 'empty_key' not in st.session_state:
            st.session_state.empty_key=''
        st.session_state.empty_key=st.session_state.user_input
        # if you set user_input to a blank valueto clear the prompt, streamlit throws an error
        st.session_state.user_input = ''