import streamlit as st
import os
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_google_community import GoogleSearchAPIWrapper
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain.tools import DuckDuckGoSearchRun, Tool
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.agents import create_openai_tools_agent, AgentExecutor
from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage
from typing_extensions import TypedDict
from typing import Annotated, Sequence
import functools
import operator



# Initialize tools
llm = ChatOpenAI()

tavily_tool = TavilySearchResults(max_results=5)
search_google_tool = Tool(
    name="GoogleSearch",
    func=GoogleSearchAPIWrapper().run,
    description="Search information using Google Search API."
)

duckduck_search_tool = Tool(
    name="DuckDuckGoSearch",
    func=DuckDuckGoSearchRun().run,
    description="Search information using DuckDuckGo."
)

serper_tool = Tool(
    name="GoogleSerperSearch",
    func=GoogleSerperAPIWrapper(max_results=5).run,
    description="Perform searches using Google Serper API."
)

tavily_tool_wrapped = Tool(
    name="TavilySearch",
    func=tavily_tool.run,
    description="Retrieve search results from Tavily API."
)

# Define reusable function for agent creation
def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            MessagesPlaceholder(variable_name="messages"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
    )
    agent = create_openai_tools_agent(llm, tools, prompt)
    executor = AgentExecutor(agent=agent, tools=tools)
    return executor


# Define agents
def get_agents():
    cto_agent = create_agent(
        llm,
        [duckduck_search_tool],
        "You are a CTO name finder. Extract the CTO's name from the provided company data."
    )

    glassdoor_agent = create_agent(
        llm,
        [tavily_tool_wrapped, serper_tool],
        "You are a Glassdoor review scraper. Retrieve reviews about the given company. "
        "Consider points like Overall Rating, Compensation, Senior Management, Career Opportunities."
        "Provide me number of stars against each point."
        "Always scrap the same data"
    )

    competitor_agent = create_agent(
        llm,
        [tavily_tool_wrapped, serper_tool],
        "You are a competitor finder. Provide details such as a description of competitors and their primary differences."
        "Output the results in a table format."
    )

    information_agent = create_agent(
        llm,
        [search_google_tool, tavily_tool_wrapped, serper_tool],
        "You are an information collector. Retrieve details such as Website, Sector, Industry, Location, Employees, Founding Year, and LinkedIn URL. Provide me all these detail in a tabular format."
        "Linkedin URL will be always like this https://www.linkedin.com/company/company_name"
    )

    return cto_agent, glassdoor_agent, competitor_agent, information_agent


# Streamlit App
def main():
    st.title("Company Insights API")
    st.write("Enter a company name to fetch details about its CTO, competitors, Glassdoor reviews, and general information.")

    # Input for company name
    company_name = st.text_input("Enter company name")
    run_queries = st.button("Run Queries")

    if run_queries:
        # Prepare agents
        cto_agent, glassdoor_agent, competitor_agent, information_agent = get_agents()

        # Queries
        queries = {
            "CTO": f"Who is the CTO of {company_name}?",
            "Glassdoor Reviews": f"What are the Glassdoor reviews of {company_name}?",
            "Competitors": f"What are the competitors of {company_name}?",
            "Information": f"Give me all information about {company_name}.",
        }

        results = {}
        for query_name, query in queries.items():
            agent = {
                "CTO": cto_agent,
                "Glassdoor Reviews": glassdoor_agent,
                "Competitors": competitor_agent,
                "Information": information_agent,
            }[query_name]

            state = {
                "messages": [HumanMessage(content=query)]
            }

            try:
                response = agent.invoke(state)
                results[query_name] = response.get("output", "No response")
            except Exception as e:
                results[query_name] = f"Error: {e}"

        # Display results
        for query_name, result in results.items():
            st.subheader(query_name)
            st.write(result)


if __name__ == "__main__":
    main()