import os

import gradio as gr
from dotenv import load_dotenv
from langchain_chroma import Chroma
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI

from news_price_correlation import build_correlation_analysis
from nvidia_agent_core import forecast_markdown, predict_nvidia_stock_payload, route_query


load_dotenv()

if not os.getenv("OPENAI_API_KEY"):
    raise ValueError("OPENAI_API_KEY was not found. Add it to .env or your environment.")

llm = ChatOpenAI(model="gpt-5.4", temperature=0.3, max_tokens=1024)
ddg_search = DuckDuckGoSearchRun()

vectorstore = None
try:
    import chromadb
    from chromadb.config import Settings

    embeddings = HuggingFaceEmbeddings(model_name="all-mpnet-base-v2", model_kwargs={"device": "cpu"})
    client = chromadb.PersistentClient(path="./chroma_db_v2", settings=Settings(allow_reset=True))
    vectorstore = Chroma(
        client=client,
        collection_name="nvidia_annual_reports_2014_2025",
        embedding_function=embeddings,
    )
except Exception as exc:
    print(f"RAG unavailable; Gradio app will continue without annual-report retrieval: {exc}")


def forecast_answer() -> str:
    payload = predict_nvidia_stock_payload(periods=7)
    return forecast_markdown(payload)


def correlation_answer(query: str) -> str:
    return build_correlation_analysis(
        query=query,
        search=ddg_search.run,
        llm=llm,
        csv_path="nvda_2014_to_2026.csv",
    )


def researcher_answer(query: str) -> str:
    context = ""
    if vectorstore:
        docs = vectorstore.similarity_search(query, k=5)
        context = "\n\n".join(doc.page_content[:700] for doc in docs)

    news = ddg_search.run(f"NVIDIA {query} latest news OR earnings OR Blackwell OR Huawei")[:1000]
    response = llm.invoke(
        [
            SystemMessage(content="You are NVIDIA's expert financial and strategy analyst."),
            HumanMessage(content=f"Question: {query}\n\nAnnual-report context:\n{context}\n\nNews:\n{news}"),
        ]
    )
    return response.content


def get_response(query: str) -> str:
    route = route_query(query)
    if route == "ML_agent":
        return forecast_answer()
    if route == "correlation_agent":
        return correlation_answer(query)
    if route in {"results_strategy_agent", "outlook_agent"}:
        return researcher_answer(query)

    response = llm.invoke(
        [
            SystemMessage(content="You are a helpful NVIDIA assistant."),
            HumanMessage(content=query),
        ]
    )
    return response.content


def chat(message, history):
    try:
        return get_response(message)
    except Exception as exc:
        return f"Error: {exc}"


demo = gr.ChatInterface(
    fn=chat,
    title="NVIDIA AI Assistant",
    description="Notebook-aligned multi-agent assistant with RAG, news correlation, and Prophet v3 residual ML forecasting.",
    examples=[
        "Predict Nvidia stock price for the next 7 business days",
        "What did Nvidia report about gross profit in the 2025 annual report?",
        "Why did NVDA sell off on China/Huawei export-control news?",
        "If CCP China invades Taiwan, predict the scale of NVDA selloff and share price drop",
    ],
)


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)