import os
import streamlit as st
from langchain.agents import create_react_agent, AgentExecutor
from langchain_community.tools import DuckDuckGoSearchRun
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_google_genai import ChatGoogleGenerativeAI

# Get API Key from Hugging Face secrets
key = os.environ.get("GOOGLE_API_KEY")
if not key:
    st.error("❌ GOOGLE_API_KEY not found. Please set it in your Hugging Face Secrets.")
    st.stop()

# Set page config
st.set_page_config(page_title="Company Info Agent", layout="wide")

st.title("📈 Business Intelligence Assistant")
st.markdown(
    "Ask anything about companies: *stock, revenue, growth, acquisitions, news*, etc."
)

# Custom Prompt
prompt = PromptTemplate(
    input_variables=["input", "agent_scratchpad", "tool_names", "tools"],
    template="""You are a smart and reliable business assistant.

You help users by answering questions about *any company* (public or private), including their:
- Stock performance
- Revenue
- Funding
- Acquisitions
- Market valuation
- Recent news
- Business activities

You can search online to gather accurate and current data. You must cite sources (like URLs or page titles) in your final answer.

You have access to the following tools:

{tools}

Use the following format:

Question: the input question you must answer  
Thought: think step-by-step about what to do  
Action: the action to take, should be one of [{tool_names}]  
Action Input: the input to the action  
Observation: the result of the action  

(Repeat the above Thought/Action/Action Input/Observation format only *2 times*)

Then:

Thought: I now know the final answer  
Final Answer: provide a detailed answer, *include helpful insights and links to your sources*.  
If no specific info is available, say so politely and try to offer related context.

---

Begin!

Question: {input}  
Thought:{agent_scratchpad}
"""
)

# Tool and LLM
search_tool = DuckDuckGoSearchRun()
tools = [search_tool]

llm = ChatGoogleGenerativeAI(
    api_key=key,
    model="gemini-2.0-flash"
)

# User Input
user_input = st.text_input("Enter your question:")

# Create Agent and Executor
agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    handle_parsing_errors=True,
)

# Handle User Input
if user_input:
    with st.spinner("🔎 Thinking..."):
        response = agent_executor.invoke({"input": user_input})
    st.markdown("### 🧠 Answer:")
    st.write(response["output"])