from typing import Annotated

from langgraph.graph import MessagesState
from typing_extensions import TypedDict

from tradingagents.agents import *


# Researcher team state
class InvestDebateState(TypedDict):
    bull_history: Annotated[str, "Bullish Conversation history"]  # Bullish Conversation history
    bear_history: Annotated[str, "Bearish Conversation history"]  # Bullish Conversation history
    history: Annotated[str, "Conversation history"]  # Conversation history
    current_response: Annotated[str, "Latest response"]  # Last response
    judge_decision: Annotated[str, "Final judge decision"]  # Last response
    count: Annotated[int, "Length of the current conversation"]  # Conversation length


# Risk management team state
class RiskDebateState(TypedDict):
    risky_history: Annotated[str, "Risky Agent's Conversation history"]  # Conversation history
    safe_history: Annotated[str, "Safe Agent's Conversation history"]  # Conversation history
    neutral_history: Annotated[str, "Neutral Agent's Conversation history"]  # Conversation history
    history: Annotated[str, "Conversation history"]  # Conversation history
    latest_speaker: Annotated[str, "Analyst that spoke last"]
    current_risky_response: Annotated[str, "Latest response by the risky analyst"]  # Last response
    current_safe_response: Annotated[str, "Latest response by the safe analyst"]  # Last response
    current_neutral_response: Annotated[
        str, "Latest response by the neutral analyst"
    ]  # Last response
    judge_decision: Annotated[str, "Judge's decision"]
    count: Annotated[int, "Length of the current conversation"]  # Conversation length


class AgentState(MessagesState):
    company_of_interest: Annotated[str, "Company that we are interested in trading"]
    trade_date: Annotated[str, "What date we are trading at"]

    sender: Annotated[str, "Agent that sent this message"]

    # research step
    market_report: Annotated[str, "Report from the Market Analyst"]
    sentiment_report: Annotated[str, "Report from the Social Media Analyst"]
    news_report: Annotated[str, "Report from the News Researcher of current world affairs"]
    fundamentals_report: Annotated[str, "Report from the Fundamentals Researcher"]

    # researcher team discussion step
    investment_debate_state: Annotated[
        InvestDebateState, "Current state of the debate on if to invest or not"
    ]
    investment_plan: Annotated[str, "Plan generated by the Analyst"]

    trader_investment_plan: Annotated[str, "Plan generated by the Trader"]

    # risk management team discussion step
    risk_debate_state: Annotated[RiskDebateState, "Current state of the debate on evaluating risk"]
    final_trade_decision: Annotated[str, "Final decision made by the Risk Analysts"]


class DiscoveryState(TypedDict):
    trade_date: Annotated[str, "Current trading date (yyyy-mm-dd format)"]
    tickers: Annotated[list[str], "List of tickers found"]
    candidate_metadata: Annotated[list[dict], "Metadata for candidates (source, strategy)"]
    filtered_tickers: Annotated[list[str], "List of tickers after filtering"]
    opportunities: Annotated[list[dict], "List of final opportunities with rationale"]
    final_ranking: Annotated[str, "Final ranking from LLM"]
    status: Annotated[str, "Current status of discovery"]
    tool_logs: Annotated[
        list[dict], "Detailed logs of all tool calls across all nodes (scanner, filter, deep_dive)"
    ]