from typing import Annotated, Any, Mapping, Optional
from typing_extensions import NotRequired, TypedDict
from langgraph.graph import MessagesState


RESEARCH_PROVENANCE_FIELDS = (
    "research_status",
    "research_mode",
    "timed_out_nodes",
    "degraded_reason",
    "covered_dimensions",
    "manager_confidence",
)


def extract_research_provenance(
    debate_state: Mapping[str, Any] | None,
) -> dict[str, Any] | None:
    if not isinstance(debate_state, Mapping):
        return None
    metadata = {
        key: debate_state.get(key)
        for key in RESEARCH_PROVENANCE_FIELDS
        if key in debate_state
    }
    return metadata or None


# Researcher team state
class InvestDebateState(TypedDict, total=False):
    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
    research_status: NotRequired[Annotated[str, "Research stage status: full/degraded/failed"]]
    research_mode: NotRequired[Annotated[str, "Research mode: debate/degraded_synthesis"]]
    timed_out_nodes: NotRequired[Annotated[list[str], "Research nodes that timed out"]]
    degraded_reason: NotRequired[Annotated[Optional[str], "Research degradation reason"]]
    covered_dimensions: NotRequired[Annotated[list[str], "Research dimensions covered so far"]]
    manager_confidence: NotRequired[Annotated[Optional[float], "Research manager confidence"]]


# Risk management team state
class RiskDebateState(TypedDict):
    aggressive_history: Annotated[
        str, "Aggressive Agent's Conversation history"
    ]  # Conversation history
    conservative_history: Annotated[
        str, "Conservative 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_aggressive_response: Annotated[
        str, "Latest response by the aggressive analyst"
    ]  # Last response
    current_conservative_response: Annotated[
        str, "Latest response by the conservative 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"]
    portfolio_context: Annotated[str, "Optional portfolio/account context for this analysis"]
    peer_context: Annotated[str, "Optional same-theme or peer ranking context for this analysis"]
    peer_context_mode: Annotated[str, "Mode describing whether peer_context is same-theme normalized or only a book snapshot"]

    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"]
    investment_plan_structured: Annotated[
        Mapping[str, Any], "Structured metadata extracted from the research-manager decision"
    ]

    trader_investment_plan: Annotated[str, "Plan generated by the Trader"]
    trader_investment_plan_structured: Annotated[
        Mapping[str, Any], "Structured metadata extracted from the trader decision"
    ]

    # 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"]
    final_trade_decision_report: Annotated[str, "Human-readable final decision report"]
    final_trade_decision_structured: Annotated[
        Mapping[str, Any], "Structured metadata extracted from the portfolio-manager decision"
    ]