src/core/FinAgentEngine.py

import json
import logging
import uuid
from typing import Any, Dict, List, Optional

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END

# Import the refactored Agent classes
from src.agents.CryptoAgent import CryptoAgent
from src.agents.EducationAgent import EducationAgent
from src.agents.GoalPlanningAgent import GoalPlanningAgent
from src.agents.MarketAgent import MarketAgent
from src.agents.NewsSynthesizerAgent import NewsSynthesizerAgent
from src.agents.PortfolioAgent import PortfolioAgent
from src.agents.TaxAgent import TaxAgent
from src.core.FinanceState import FinanceState
from src.core.errors import add_error
from src.core.logging_config import configure_logging
from src.core.settings import get_settings



logger = logging.getLogger(__name__)


class FinAgentEngine:
    def __init__(self):
        configure_logging()
        load_dotenv()
        settings = get_settings()
        self.router_llm = ChatOpenAI(model=settings.models.router_model)
        self.agent_llm = ChatOpenAI(model=settings.models.agent_model)

        # Build the graph
        self.graph = StateGraph(FinanceState)
        self.graph.add_node("router", self.router_node)
        self.graph.add_node("multi_agent", self.multi_agent_node)
        self.graph.add_node("portfolio_agent", self.portfolio_node)
        self.graph.add_node("market_agent", self.market_node)
        self.graph.add_node("education_agent", self.education_node)
        self.graph.add_node("crypto_agent", self.crypto_node)
        self.graph.add_node("tax_agent", self.tax_node)
        self.graph.add_node("goal_planning_agent", self.goal_planning_node)
        self.graph.add_node("news_synthesizer_agent", self.news_synthesizer_node)

        self.graph.set_entry_point("router")
        self.graph.add_conditional_edges(
            "router",
            self.route,
            {
                "multi_agent": "multi_agent",
                "portfolio_agent": "portfolio_agent",
                "market_agent": "market_agent",
                "education_agent": "education_agent",
                "crypto_agent": "crypto_agent",
                "tax_agent": "tax_agent",
                "goal_planning_agent": "goal_planning_agent",
                "news_synthesizer_agent": "news_synthesizer_agent",
                "default": END,
            },
        )
        self.graph.add_edge("multi_agent", END)
        self.app = self.graph.compile()

    def llm_router(
        self,
        query: str,
        user_profile: Optional[Dict[str, Any]] = None,
        conversation_history: Optional[List[Dict[str, str]]] = None,
    ):
        prompt_content = f"""
You are an intelligent routing and entity extraction engine for a finance AI system.

Your tasks:
1. Classify the user query into one of:
   - education
   - portfolio
   - market
   - tax
   - crypto
   - goal_planning
   - news
   - none

2. Select relevant agents:
   - education_agent
   - portfolio_agent
   - market_agent
   - tax_agent
   - crypto_agent
   - goal_planning_agent
   - news_synthesizer_agent
   - none

3. Extract stock symbol (if applicable):
   - If a company or stock is mentioned, return its correct ticker symbol
   - Examples:
     Apple → AAPL
     Tesla → TSLA
     Nvidia → NVDA
   - If no stock/company is mentioned, return null

4. Extract cryto symbol (if applicable):
   - If a cryptocurrency is mentioned, return its correct ticker symbol
   - Examples:
     Bitcoin → BTC
     Ethereum → ETH
     Litecoin → LTC
   - If no cryptocurrency is mentioned, return null

5. If the intent is 'portfolio' and the query mentions specific holdings, extract them as a list of dictionaries with 'symbol' and 'quantity'.
   - Example: "My portfolio contains: 100 Apple shares, 1000 Nvidia and 250 Tesla shares" → [{{"symbol": "AAPL", "quantity": 100}}, {{"symbol": "NVDA", "quantity": 1000}}, {{"symbol": "TSLA", "quantity": 250}}]
   - If no portfolio details are mentioned, return null.

Rules:
- Market queries → include market_agent
- Investment decisions → include portfolio_agent
- Learning queries → education_agent
- Infomartion queries → education_agent
- Tax queries → tax_agent
- Crypto SPOT PRICE / QUOTE queries (e.g., "price of bitcoin", "BTC price", "1 ETH price today") → crypto_agent
- Crypto PREDICTION / FORECAST queries (e.g., "bitcoin price predictions", "BTC forecast", "will bitcoin go up") → news_synthesizer_agent (use web search)
- Financial goal planning queries (saving for X, retirement plan, house down payment) → goal_planning_agent
- News queries (today's news, latest headlines, why did stock move, earnings headlines) → news_synthesizer_agent
- If multiple intents → multi
- If no intent → none

Return ONLY valid JSON (no explanation):

{{
  "intent": "...",
  "agents": ["..."],
  "symbol": "AAPL" | null,
  "crypto_symbol": "BTC" | null,
  "portfolio": [ {{ "symbol": "AAPL", "quantity": 100 }}, {{ "symbol": "NVDA", "quantity": 1000 }} ] | null,
  "query": "{query}"
}}

User Query: "{query}"
"""
        # Provide additional context to improve routing stability.
        # Keep it short; the router should still rely on the user query primarily.
        if user_profile:
            prompt_content += (
                f"\nUser Profile (context): {json.dumps(user_profile)[:1200]}\n"
            )
        if conversation_history:
            recent = conversation_history[-6:]
            prompt_content += (
                f"\nRecent Conversation (context): {json.dumps(recent)[:1200]}\n"
            )

        messages_for_llm = [
            {
                "role": "system",
                "content": "You are a strict JSON generator. Do not return anything except valid JSON.",
            },
            {"role": "user", "content": prompt_content},
        ]

        response_message = self.router_llm.invoke(messages_for_llm, temperature=0)
        content = response_message.content

        try:
            # Handle markdown-wrapped JSON
            if content.startswith("```"):
                content = content.replace("```json", "").replace("```", "").strip()

            return json.loads(content)

        except Exception as e:
            logger.error(f"Error parsing JSON: {e}")
            # Best-effort state telemetry (router is called from router_node; this is a fallback).
            return {
                "intent": "education",
                "agents": ["education_agent"],
                "symbol": None,
                "crypto_symbol": None,
                "portfolio": None,
                "query": query,
            }

    def quote_vs_web_classifier(
        self,
        *,
        query: str,
        symbol: Optional[str] = None,
        crypto_symbol: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Lightweight disambiguation for asset questions:

        Decides whether the user wants:
        - a spot quote ("quote") or
        - a web-based synthesis ("web") or
        - general education ("education")

        This is intentionally smaller than `llm_router()` and should only be used when
        heuristic signals are ambiguous.
        """
        q = (query or "").strip()
        prompt = f"""
You are a strict JSON classifier for financial user queries.

Decide the user's intent for the query as one of:
- quote: user wants the current/spot price/quote or latest numeric snapshot
- web: user wants predictions/forecast/outlook/news/why-it-moved analysis (needs web search)
- education: user wants a general explanation/definition (no need for live quotes)

Asset hints (may be null):
- stock symbol: {symbol or None}
- crypto symbol: {crypto_symbol or None}

Rules:
- If the query asks for "prediction/forecast/outlook/target/will it go up" => web
- If the query asks "price/quote/how much is X" or a numeric snapshot => quote
- If the query asks "what is" / definitions / concepts => education

Return ONLY valid JSON:
{{
  "mode": "quote" | "web" | "education",
  "confidence": 0.0-1.0,
  "reason": "short"
}}

Query: {json.dumps(q)}
"""
        try:
            msg = self.router_llm.invoke(
                [
                    {"role": "system", "content": "Return only valid JSON."},
                    {"role": "user", "content": prompt},
                ],
                temperature=0,
            )
            content = str(msg.content or "").strip()
            if content.startswith("```"):
                content = content.replace("```json", "").replace("```", "").strip()
            out = json.loads(content)
            if not isinstance(out, dict):
                return {"mode": "web", "confidence": 0.0, "reason": "invalid_json"}
            mode = out.get("mode")
            if mode not in ("quote", "web", "education"):
                out["mode"] = "web"
            return out
        except Exception as e:
            logger.exception("quote_vs_web_classifier failed: %s", e)
            return {"mode": "web", "confidence": 0.0, "reason": f"error:{e}"}

    def router_node(self, state: FinanceState):
        query = state.get("user_query", "")
        trace_id = str(state.get("trace_id") or "")
        logger.info("[trace=%s] router_node.start query_len=%d", trace_id, len(query or ""))

        decision = self.llm_router(
            query,
            user_profile=state.get("user_profile"),
            conversation_history=state.get("conversation_history"),
        )

        def _has_any(text: str, tokens: set[str]) -> bool:
            return any(t in text for t in tokens)

        qn = (query or "").strip().lower()
        # Prefer the LLM router's entity extraction for deciding whether a query is crypto-related.
        # Keep these heuristic sets small and high-signal.
        quote_tokens = {"price", "quote", "spot", "how much", "current price"}
        web_tokens = {"predict", "prediction", "forecast", "outlook", "price target", "will it"}

        crypto_symbol = decision.get("crypto_symbol")
        has_crypto_symbol = isinstance(crypto_symbol, str) and crypto_symbol.strip() != ""
        mentions_crypto = has_crypto_symbol or decision.get("intent") == "crypto" or "crypto_agent" in (
            decision.get("agents") or []
        )
        is_spot_price = _has_any(qn, quote_tokens)
        is_prediction = _has_any(qn, web_tokens)

        if mentions_crypto and is_prediction and not is_spot_price:
            decision = dict(decision or {})
            decision["intent"] = "news"
            decision["agents"] = ["news_synthesizer_agent"]
            # Keep entity extraction if present; otherwise best-effort map common coins.
            logger.info(
                "[trace=%s] router_node.override crypto_prediction -> news_synthesizer_agent",
                trace_id,
            )

        # Similar override for stocks: if a ticker is present and the user asked for
        # predictions/forecast/outlook (not a spot quote), route to web-search synthesis.
        symbol = decision.get("symbol")
        has_symbol = isinstance(symbol, str) and symbol.strip() != ""
        if has_symbol and is_prediction and not is_spot_price:
            decision = dict(decision or {})
            decision["intent"] = "news"
            decision["agents"] = ["news_synthesizer_agent"]
            logger.info(
                "[trace=%s] router_node.override stock_prediction -> news_synthesizer_agent symbol=%s",
                trace_id,
                symbol,
            )

        # Ambiguous cases: use a small LLM classifier to decide quote vs web vs education.
        # Only run when we have a clear asset mention but no strong heuristic signal.
        asset_mentioned = mentions_crypto or has_symbol or has_crypto_symbol
        ambiguous = asset_mentioned and (not is_spot_price) and (not is_prediction)

        if ambiguous:
            classification = self.quote_vs_web_classifier(
                query=query,
                symbol=str(symbol) if has_symbol else None,
                crypto_symbol=str(crypto_symbol) if has_crypto_symbol else None,
            )
            mode = classification.get("mode")
            conf = classification.get("confidence")
            logger.info(
                "[trace=%s] quote_vs_web_classifier mode=%s confidence=%s reason=%s",
                trace_id,
                mode,
                conf,
                classification.get("reason"),
            )
            if mode == "web":
                decision = dict(decision or {})
                decision["intent"] = "news"
                decision["agents"] = ["news_synthesizer_agent"]
            elif mode == "quote":
                # Route quote requests to the appropriate quote agent.
                decision = dict(decision or {})
                if has_crypto_symbol or mentions_crypto:
                    decision["intent"] = "crypto"
                    decision["agents"] = ["crypto_agent"]
                elif has_symbol:
                    decision["intent"] = "market"
                    decision["agents"] = ["market_agent"]
            elif mode == "education":
                decision = dict(decision or {})
                decision["intent"] = "education"
                decision["agents"] = ["education_agent"]

        logger.info(
            "[trace=%s] router_node.decision intent=%s agents=%s symbol=%s crypto=%s",
            trace_id,
            decision.get("intent"),
            decision.get("agents"),
            decision.get("symbol"),
            decision.get("crypto_symbol"),
        )
        state["intent"] = decision.get("intent", "education")
        state["agents"] = decision.get("agents", ["education_agent"])

        # If the router says "none" (unsupported), stop cleanly with a user-friendly response.
        # This avoids LangGraph trying to route to a non-existent node named "none".
        if state.get("intent") == "none" or "none" in (state.get("agents") or []):
            state["intent"] = "none"
            state["agents"] = []
            state["response"] = "Not supported by the Assistant."
            return state
        if "symbol" in decision:
            state["symbol"] = decision.get("symbol")
        if "crypto_symbol" in decision:
            state["crypto_symbol"] = decision.get("crypto_symbol")
        if decision.get("portfolio") is not None:
            state["portfolio"] = (state.get("portfolio") or []) + (decision.get("portfolio") or [])
        state["user_query"] = decision.get("query", state.get("user_query", ""))
        return state

    def route(self, state: FinanceState):
        """
        This function determines where the graph should go next.
        It looks at the 'agents' list we populated from the LLM.
        """
        agents = state.get("agents", [])

        # If the LLM picked exactly one agent, travel to that agent's node
        if len(agents) == 1:
            if agents[0] == "none":
                return "default"
            return agents[0]
        # Run multiple agents sequentially and combine their outputs.
        if len(agents) > 1:
            return "multi_agent"
        # Otherwise, redirect to the default end node
        else:
            return "default"

    def multi_agent_node(self, state: FinanceState):
        agents = state.get("agents", []) or []
        agent_outputs: List[Dict[str, str]] = []
        trace_id = str(state.get("trace_id") or "")
        logger.info("[trace=%s] multi_agent.start agents=%s", trace_id, agents)
        # Execute agents in a deterministic order to keep output stable.
        preferred_order = [
            "education_agent",
            "tax_agent",
            "market_agent",
            "portfolio_agent",
            "goal_planning_agent",
            "news_synthesizer_agent",
            "crypto_agent",
        ]
        ordered = [a for a in preferred_order if a in agents] + [
            a for a in agents if a not in preferred_order
        ]

        for agent_name in ordered:
            try:
                if agent_name == "education_agent":
                    state = self.education_node(state)
                elif agent_name == "tax_agent":
                    state = self.tax_node(state)
                elif agent_name == "market_agent":
                    state = self.market_node(state)
                elif agent_name == "portfolio_agent":
                    state = self.portfolio_node(state)
                elif agent_name == "goal_planning_agent":
                    state = self.goal_planning_node(state)
                elif agent_name == "news_synthesizer_agent":
                    state = self.news_synthesizer_node(state)
                elif agent_name == "crypto_agent":
                    state = self.crypto_node(state)
                else:
                    continue
                if state.get("response"):
                    agent_outputs.append(
                        {
                            "agent": agent_name,
                            "output": str(state.get("response") or ""),
                        }
                    )
            except Exception as e:
                logger.exception("multi_agent failed for %s", agent_name)
                add_error(
                    state,
                    code="agent_error",
                    message=str(e),
                    agent=agent_name,
                )
                agent_outputs.append({"agent": agent_name, "output": f"Error: {e}"})

        if not agent_outputs:
            state["response"] = "No response was generated by the selected agents."
            return state

        # Only modify the output format for multi-agent intent.
        # If the router marked this as multi (or we have >1 agent), synthesize into a single response.
        if state.get("intent") == "multi" or len(agents) > 1:
            logger.info(
                "[trace=%s] multi_agent: combining_outputs intent=%s agents=%s",
                trace_id,
                state.get("intent"),
                agents,
            )
            combined_inputs = "\n\n".join(
                [f"[{a['agent']}]\n{a['output']}".strip() for a in agent_outputs]
            )
            system = (
                "You are a financial education assistant. Combine multiple specialist agent outputs into ONE coherent response. "
                "Do not add new facts beyond what the agent outputs already contain. "
                "If agent outputs conflict, mention the uncertainty rather than choosing one. "
                "Do not provide trade instructions. Include a single short disclaimer once. "
                "If sources/URLs are present, consolidate them at the end under 'Sources'."
            )
            prompt = f"""
User query:
{state.get("user_query", "")}

User profile (context, may be incomplete):
{state.get("user_profile", {})}

Agent outputs:
{combined_inputs}

Write a single combined answer with this structure:
1) Disclaimer (one line)
2) Direct answer (1-3 short paragraphs)
3) Key points (bulleted)
4) Suggested next questions for the user (3 bullets)
5) Sources (only if present in agent outputs)
"""
            try:
                msg = self.agent_llm.invoke(
                    [
                        {"role": "system", "content": system},
                        {"role": "user", "content": prompt},
                    ],
                    temperature=0.2,
                )
                state["response"] = str(msg.content or "").strip()
                return state
            except Exception:
                logger.exception("multi_agent: combine failed")
                add_error(
                    state,
                    code="combine_error",
                    message="Failed to combine multi-agent outputs",
                    agent="multi_agent",
                )
                # Fall back to a stable, readable concatenation.
                state["response"] = "\n\n".join(
                    [f"## {a['agent']}\n{a['output']}".strip() for a in agent_outputs]
                )
                return state

        # If we somehow reached multi_agent_node without multi intent, keep the previous sectioned output.
        state["response"] = "\n\n".join(
            [f"## {a['agent']}\n{a['output']}".strip() for a in agent_outputs]
        )
        return state

    # ------------------------------------------------------------------------
    # Node Wrappers
    #
    # LangGraph expects simple functions for its nodes, but we use Python
    # Classes for our agents to keep our code organized. These functions
    # act like a "bridge" between the Graph and our Classes.
    # ------------------------------------------------------------------------

    def _run_agent(self, agent_class, state: FinanceState):
        """Generic helper to run any AgentCommand."""
        agent = agent_class(state)
        agent.process()
        return agent.state

    def education_node(self, state: FinanceState):
        return self._run_agent(EducationAgent, state)

    def market_node(self, state: FinanceState):
        return self._run_agent(MarketAgent, state)

    def portfolio_node(self, state: FinanceState):
        return self._run_agent(PortfolioAgent, state)

    def tax_node(self, state: FinanceState):
        return self._run_agent(TaxAgent, state)

    def crypto_node(self, state: FinanceState):
        return self._run_agent(CryptoAgent, state)

    def goal_planning_node(self, state: FinanceState):
        return self._run_agent(GoalPlanningAgent, state)

    def news_synthesizer_node(self, state: FinanceState):
        return self._run_agent(NewsSynthesizerAgent, state)

    def invoke(self, query: str, initial_state: Optional[Dict[str, Any]] = None):
        trace_id = str(uuid.uuid4())
        state: Dict[str, Any] = {
            "user_query": query,
            "intent": "",
            "agents": [],
            "portfolio": [],
            "response": "",
            "symbol": None,
            "crypto_symbol": None,
            "conversation_history": [],
            "user_profile": {},
            "retrieved_sources": [],
            "errors": [],
            "trace_id": trace_id,
        }
        if initial_state:
            # Do not allow callers to replace the user query string.
            merged = dict(state)
            merged.update({k: v for k, v in initial_state.items() if k != "user_query"})
            # Preserve trace_id if caller set one, otherwise keep generated.
            if not merged.get("trace_id"):
                merged["trace_id"] = trace_id
            state = merged
        final_state = self.app.invoke(state)

        return final_state

    def routeAgent(self, query: str):
        return self.invoke(query)