import operator
import os
from typing import Annotated, TypedDict

import pandas as pd
import streamlit as st
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 langgraph.graph import END, StateGraph

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


load_dotenv()

st.set_page_config(page_title="NVIDIA AI Assistant", page_icon="📈", layout="wide")

CSV_PATH = "nvda_2014_to_2026.csv"
MODEL_PATH = "nvidia_price_model.pkl"
CHROMA_PATH = "./chroma_db_v2"
COLLECTION_NAME = "nvidia_annual_reports_2014_2025"

if not os.getenv("OPENAI_API_KEY"):
    st.error("OPENAI_API_KEY not found in .env")
    st.stop()

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


@st.cache_resource
def load_rag():
    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_PATH,
            settings=Settings(allow_reset=True),
        )
        return Chroma(
            client=client,
            collection_name=COLLECTION_NAME,
            embedding_function=embeddings,
        )
    except Exception as exc:
        st.sidebar.error(f"RAG load failed: {exc}")
        return None


vectorstore = load_rag()


class AgentState(TypedDict):
    query: str
    response: str
    debug_log: Annotated[str, operator.add]
    next_node: str


def router_node(state: AgentState) -> AgentState:
    notebook_route = route_query(state["query"])
    route_map = {
        "ML_agent": "trader",
        "results_strategy_agent": "researcher",
        "outlook_agent": "researcher",
        "correlation_agent": "correlation",
        "general_agent": "general",
    }

    if notebook_route == "ML_agent" and any(
        term in state["query"].lower()
        for term in ["12 months", "one year", "next year", "2027", "long term"]
    ):
        state["next_node"] = "long_term_analyst"
    else:
        state["next_node"] = route_map.get(notebook_route, "general")

    state["debug_log"] = f"Router Decision: {state['next_node']} (notebook route: {notebook_route})\n"
    return state


def researcher_node(state: AgentState) -> AgentState:
    query = state["query"]
    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")[:1000]
    response = llm.invoke(
        [
            SystemMessage(content="You are NVIDIA expert researcher. Use retrieved report context and current news."),
            HumanMessage(content=f"Query: {query}\n\nAnnual-report context:\n{context}\n\nNews:\n{news}"),
        ]
    ).content

    state["response"] = response
    state["debug_log"] += "Researcher Agent completed.\n"
    return state


def trader_node(state: AgentState) -> AgentState:
    try:
        payload = predict_nvidia_stock_payload(periods=7)
        state["response"] = forecast_markdown(payload)
        state["debug_log"] += (
            "Trader Agent -> called shared Prophet v3 residual ensemble; "
            f"model={payload.get('model_version')}; residual={payload.get('uses_residual_model')}\n"
        )
    except Exception as exc:
        state["response"] = f"Forecast error: {exc}"
        state["debug_log"] += f"Trader Agent error: {exc}\n"
    return state


def long_term_analyst_node(state: AgentState) -> AgentState:
    try:
        checkpoint, forecast, _ = run_ensemble_forecast(periods=30, model_path=MODEL_PATH)
        last_close = checkpoint["last_close"]
        final_30d = float(forecast["yhat"].iloc[-1])
        upside_30d = (final_30d / last_close - 1) * 100
        pred = forecast.tail(7)[["ds", "yhat"]].copy()
        pred["ds"] = pd.to_datetime(pred["ds"]).dt.strftime("%Y-%m-%d")

        context = ""
        if vectorstore:
            docs = vectorstore.similarity_search(state["query"], k=6)
            context = "\n\n".join(doc.page_content[:600] for doc in docs)
        news = ddg_search.run("NVIDIA price target 2026 2027 analyst forecast Blackwell revenue")[:1000]

        llm_response = llm.invoke(
            [
                SystemMessage(content="You are NVIDIA senior equity analyst. Combine quantitative forecast with fundamentals."),
                HumanMessage(
                    content=f"""Query: {state['query']}
Short-term Prophet ensemble: expected {upside_30d:+.1f}% in 30 business days to ${final_30d:.2f}
Annual-report context: {context}
Latest analyst/news: {news}

Give a balanced 12-month outlook with key drivers, risks, a realistic range, and a trade view."""
                ),
            ]
        ).content

        forecast_lines = "\n".join(f"{row['ds']}: **${row['yhat']:.2f}**" for _, row in pred.iterrows())
        state["response"] = f"""**NVIDIA 12-Month Outlook (Hybrid ML + LLM)**

**Short-term Prophet v3 residual ensemble, next 7 business days:**
{forecast_lines}

**30-business-day expected move:** **{upside_30d:+.1f}%** -> ~${final_30d:.2f}

{llm_response}

*Model: {checkpoint.get('model_version', 'unknown')} | Residual ML: {checkpoint.get('residual_model') is not None} | Backtested MAPE: {checkpoint.get('backtest_mape', 0):.2f}% | Directional Accuracy: {checkpoint.get('directional_accuracy', 0):.1f}%*
"""
        state["debug_log"] += "Long-term Analyst completed with shared ensemble forecast.\n"
    except Exception as exc:
        state["response"] = f"Long-term analysis error: {exc}"
        state["debug_log"] += f"Long-term Analyst error: {exc}\n"
    return state


def correlation_node(state: AgentState) -> AgentState:
    try:
        state["response"] = build_correlation_analysis(
            query=state["query"],
            search=ddg_search.run,
            llm=llm,
            csv_path=CSV_PATH,
        )
        state["debug_log"] += "Correlation Agent -> shared news/price correlation analysis.\n"
    except Exception as exc:
        state["response"] = f"Correlation analysis error: {exc}"
        state["debug_log"] += f"Correlation Agent error: {exc}\n"
    return state


def general_node(state: AgentState) -> AgentState:
    response = llm.invoke(
        [
            SystemMessage(content="You are a helpful NVIDIA assistant."),
            HumanMessage(content=state["query"]),
        ]
    ).content
    state["response"] = response
    state["debug_log"] += "General Agent completed.\n"
    return state


workflow = StateGraph(AgentState)
workflow.add_node("router", router_node)
workflow.add_node("researcher", researcher_node)
workflow.add_node("trader", trader_node)
workflow.add_node("long_term_analyst", long_term_analyst_node)
workflow.add_node("correlation", correlation_node)
workflow.add_node("general", general_node)
workflow.set_entry_point("router")
workflow.add_conditional_edges(
    "router",
    lambda state: state["next_node"],
    {
        "researcher": "researcher",
        "trader": "trader",
        "long_term_analyst": "long_term_analyst",
        "correlation": "correlation",
        "general": "general",
    },
)
for node in ["researcher", "trader", "long_term_analyst", "correlation", "general"]:
    workflow.add_edge(node, END)

app = workflow.compile()


st.title("NVIDIA AI Assistant - Multi-Agent System")
st.caption("NTU DSAI Capstone | RAG + DuckDuckGo + Prophet v3 Residual ML Ensemble")

if "messages" not in st.session_state:
    st.session_state.messages = []

for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

if prompt := st.chat_input("Ask anything about NVIDIA..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            try:
                result = app.invoke({"query": prompt, "response": "", "debug_log": "", "next_node": ""})
                answer = result.get("response", "Sorry, I couldn't generate an answer.")
                log = result.get("debug_log", "No debug info")
            except Exception as exc:
                answer = f"Agent error: {exc}"
                log = f"Error: {exc}"

            st.markdown(answer)
            with st.expander("Debug Trace"):
                st.code(log)

    st.session_state.messages.append({"role": "assistant", "content": answer})

with st.sidebar:
    st.success("Prophet v3 residual ensemble ready")
    st.success("News/price correlation agent ready")
    st.success("RAG ready" if vectorstore else "RAG unavailable")
    st.success("DuckDuckGo ready")
    if st.button("Clear Chat"):
        st.session_state.messages = []
        st.rerun()