# FINAL PROJECT FOR GAIA TEST
from llama_index.tools.tavily_research import TavilyToolSpec
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.llms.groq import Groq
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
import os
from llama_index.tools.yahoo_finance import YahooFinanceToolSpec
from llama_index.core.agent.workflow import ReActAgent
from dotenv import load_dotenv
import asyncio
from llama_index.core.tools import FunctionTool, QueryEngineTool, ToolMetadata
import gradio as gr
import datetime
from llama_index.llms.google_genai import GoogleGenAI

load_dotenv()



tavily_key = os.getenv("TAVILY_API_KEY")
groq_key = os.getenv("GROQ_API_KEY")
current_date = datetime.datetime.now().strftime("%A, %B %d, %Y")
my_llm = llm = GoogleGenAI(model="models/gemma-4-26b-a4b-it")
tavily_tool = TavilyToolSpec(
    api_key=tavily_key,
)


# load_dotenv()
# # Set API Key
# groq_api_key = os.environ["GROQ_API_KEY"]

# # 1. Define LLM (Groq)
Settings.llm = my_llm

# # 2. Define Embedding Model (Example: HuggingFace)
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
# Settings.embed_model = HuggingFaceInferenceAPIEmbedding(
#     model_name="BAAI/bge-small-en-v1.5"
# )
Settings.embed_model = embed_model

# 3. Load Data and Build Index
# This looks at the specific FILE
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents)

# 4. Query Engine
query_engine = index.as_query_engine()


# WEB SEARCH AGENT

web_search_tools = tavily_tool.to_tool_list() + YahooFinanceToolSpec().to_tool_list()

web_search_agent = ReActAgent(
    tools=web_search_tools,
    llm=my_llm,
    verbose=True,
    system_prompt="""You are a financial assistant. Use tavily tool to search for a company's ticker and then use it YahooFinanceToolSpec to find current stock price. You can also use the tools to research on a companies market news
           When providing a final Answer:
            1. Do NOT say "Based on the data..." or "I found that...".
            2. Provide ONLY the requested value and the timestamp.
            3. If multiple tools are used, synthesize them into a single-sentence fact.
            4. Eliminate all conversational filler. 
            Today's date is {current_date}. use that to calculate for relative dates""",
    timeout=120,
)


summarize_agent = ReActAgent(
    tools=[],
    llm=my_llm,
    verbose=True,
    system_prompt="""You are a helpful assistant that can summarize information. Use this agent to summarize the information retrieved by the web search agent or to synthesize information from multiple sources. Provide only the final synthesized answer without any explanation or filler.""",
    timeout=120,
)
async def call_researcher(query: str) -> str:
    """Delegates a research task to the web search Agent."""
    response = await web_search_agent.run(user_msg=query)
    return str(response)


def search_documents(query: str) -> str:
    """
    Search through the local documentation for specific facts.
    Use this tool whenever you need information from the uploaded files or when a lookup is needed..
    """
    # query_engine should be defined earlier in your code
    response = query_engine.query(query)
    return str(response)

def summarize(information: str) -> str:
    """
    Summarize the given information using the summarize agent.
    Use this tool to synthesize information from multiple sources or to condense lengthy information into a concise summary.
    Also useful for summarizing information received from the other agent and produce the final answer
    """
    response = summarize_agent.run(user_msg=information)
    return str(response)

# Code for the manager agent

manager_tools = [
    FunctionTool.from_defaults(fn=call_researcher, name="Researcher_Specialist"),
    FunctionTool.from_defaults(fn=search_documents, name="Document_Searcher"),
    FunctionTool.from_defaults(fn=summarize, name="Summarize_Information"),
]

manager_agent = ReActAgent(
    tools=manager_tools,
    llm=my_llm,
    verbose=True,
    system_prompt="""You are a manager agent that delegates financial research tasks to a web search agent. When you receive a query, use the Researcher_Specialist tool to get the information needed to answer the question. Provide only the final synthesized answer without any explanation or filler.""",
    timeout=120,
)

async def main():
    answer = await manager_agent.run(user_msg="What is the current stock price of Apple Inc AAPL?")
    clean_answer = str(answer.response).strip()
    print(clean_answer)
# asyncio.run(main())


async def predict(message, history):
    """
    Handles the conversation. 'history' is provided by Gradio 
    but we'll let the ReActAgent handle its own internal memory for now.
    """
    try:
        # Run the agent asynchronously
        response = await manager_agent.run(user_msg=message)
        return str(response)
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

# 5. Launch Gradio Interface
demo = gr.ChatInterface(
    fn=predict,
    title="📈 AI Finance & Research Agent",
    description="Ask me about stock prices, market news, or company financial data.",
    examples=["What is the current price of AAPL?", "Compare NVDA and AMD performance this week."],
)

if __name__ == "__main__":
    # CRITICAL: server_name and server_port are required for HF Spaces
    demo.launch(server_name="0.0.0.0", server_port=7860)