import streamlit as st
from langchain.agents import initialize_agent, Tool, AgentType
from langchain_groq import ChatGroq
from typing import TypedDict, Dict, List
from langchain_community.tools import TavilySearchResults
from dotenv import load_dotenv
import os

load_dotenv()

GROQ_API_KEY = os.getenv("GROQ_API_KEY")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")



llm = ChatGroq(
    name="chat_groq",
    model="llama3-groq-70b-8192-tool-use-preview",
    max_tokens=8000,
    temperature=0.5,
    api_key=GROQ_API_KEY,
)



def web_search_tool(query: str) -> str:
    """This tool performs web search to gather relevant data."""
    # client = TavilyClient(api_key="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")
    websearch = TavilySearchResults(max_results=5, search_depth="advanced",api_key=TAVILY_API_KEY )
    response = websearch.invoke({"query": query})
    return response



tools = [
    Tool(
        name="WebSearchTool",
        func=web_search_tool,
        description="Use this tool to perform a web search for the given query."
    )
]


agent = initialize_agent(
    tools,
    llm,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)



class ResearchState(TypedDict):
    company: str
    company_keywords: str
    exclude_keywords: str
    documents: Dict[str, str]
    RAG_docs: List[str]



def research_model(state: ResearchState) -> ResearchState:
    query = state['company_keywords']
    search_result = web_search_tool(query)
    state['documents'][query] = search_result
    return state



def process_documents(state: ResearchState) -> ResearchState:
    relevant_docs = []
    for query, doc in state['documents'].items():
        if state['exclude_keywords'] not in doc:
            relevant_docs.append(doc)
    state['RAG_docs'] = relevant_docs
    return state



def write_report(state: ResearchState) -> str:

    top_companies = f"""Top Companies in {state['company']} Domain: list down all the companies that are being a top player in this domain\n"""


    ecosystems = f"Related Ecosystems or Domains:\n{state['company_keywords']} list down other subdomains are industries that are hugely impacted because of this\n"


    future_trends = f"""list down the interesting future trends of {state['company']} adoption across industries mentioned in the NEWS articles"""


    prompt = f"""Based on the documents gathered for the domain '{state['company']}', summarize the key companies, trends, and ecosystem.
    - Key Companies:
    {top_companies}
    - Ecosystems:
    {ecosystems}
    - Future Trends:
    {future_trends}
    Documents considered: {str(state['RAG_docs'])}.
    Ensure that the output is structured into sections like Top Companies, Related Ecosystems, and Future Trends.
    """

    messages = [
        {"role": "system", "content": "You are an assistant that helps with research."},
        {"role": "user", "content": prompt}
    ]

    try:
        # Pass the messages to the LLM and get the result
        result = llm.invoke(messages)
        st.header("Research Report:")
        st.write(result.content)
        st.write(result.response_metadata)
        st.subheader("References:")
        st.write(state)
        return result
    except Exception as e:
        return f"An error occurred: {str(e)}"


# Streamlit UI layout
st.set_page_config(page_title="Domain Research", layout="wide")
st.title("Research on Domains and Ecosystems")

# Left panel for user input
with st.sidebar:
    st.header("Input Section")
    company = st.text_input("Company Name:", "Electric Vehicles")
    company_keywords = st.text_area("Company Keywords:",
                                    "electric vehicle industry, top EV brands, EV trends, EV ecosystem")
    exclude_keywords = st.text_area("Exclude Keywords:", "obsolete, outdated")

# Run the research workflow when the user submits the form
if st.button("Generate Report"):
    state = ResearchState(
        company=company,  # User input company
        company_keywords=company_keywords,  # User input keywords
        exclude_keywords=exclude_keywords,  # User input exclude keywords
        documents={},
        RAG_docs=[]
    )

    state = research_model(state)
    state = process_documents(state)
    report = write_report(state)