from typing import Annotated, Any, Sequence
from datetime import date, timedelta, datetime
from typing_extensions import TypedDict, Optional
from langchain_openai import ChatOpenAI
from tradingagents.agents import *
from langgraph.prebuilt import ToolNode
from langgraph.graph import END, StateGraph, START, MessagesState


# 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):
    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 FairValueRange(TypedDict):
    low: Optional[float]
    high: Optional[float]


class ValuationData(TypedDict):
    fair_value_range: FairValueRange
    expected_return_pct: Optional[float]
    primary_method: str
    thesis: str


class SegmentData(TypedDict):
    ticker: str
    analysis_date: str
    business_unit_decomposition: list[dict[str, Any]]
    segment_economics: dict[str, Any]
    value_driver_map: list[dict[str, Any]]


class ScenarioCaseData(TypedDict):
    probability: Optional[float]
    price_target: Optional[float]
    thesis: str


class ScenarioCatalystData(TypedDict):
    ticker: str
    analysis_date: str
    scenario_map: list[dict[str, Any]]
    dated_catalyst_map: list[dict[str, Any]]
    invalidation_triggers: list[dict[str, Any]]


class PositionSizingData(TypedDict):
    ticker: str
    analysis_date: str
    conviction: str
    target_weight_pct: Optional[float]
    initial_weight_pct: Optional[float]
    max_loss_pct: Optional[float]
    sizing_rationale: str


class ChiefAnalystData(TypedDict):
    action: str
    summary: str
    thesis: str
    confidence: str


def make_default_valuation_data() -> ValuationData:
    return {
        "fair_value_range": {"low": None, "high": None},
        "expected_return_pct": None,
        "primary_method": "",
        "thesis": "",
    }


def make_default_segment_data() -> SegmentData:
    return {
        "ticker": "",
        "analysis_date": "",
        "business_unit_decomposition": [],
        "segment_economics": {},
        "value_driver_map": [],
    }


def make_default_scenario_case_data() -> ScenarioCaseData:
    return {
        "probability": None,
        "price_target": None,
        "thesis": "",
    }


def make_default_scenario_catalyst_data() -> ScenarioCatalystData:
    return {
        "ticker": "",
        "analysis_date": "",
        "scenario_map": [],
        "dated_catalyst_map": [],
        "invalidation_triggers": [],
    }


def make_default_position_sizing_data() -> PositionSizingData:
    return {
        "ticker": "",
        "analysis_date": "",
        "conviction": "",
        "target_weight_pct": None,
        "initial_weight_pct": None,
        "max_loss_pct": None,
        "sizing_rationale": "",
    }


def make_default_chief_analyst_data() -> ChiefAnalystData:
    return {
        "action": "",
        "summary": "",
        "thesis": "",
        "confidence": "",
    }


def make_default_structured_stock_underwriting_state() -> dict[str, Any]:
    return {
        "valuation_data": make_default_valuation_data(),
        "segment_data": make_default_segment_data(),
        "scenario_catalyst_data": make_default_scenario_catalyst_data(),
        "position_sizing_data": make_default_position_sizing_data(),
        "chief_analyst_data": make_default_chief_analyst_data(),
    }


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"
    ]
    macro_report: Annotated[str, "Report from the Macro Analyst"]
    fundamentals_report: Annotated[str, "Report from the Fundamentals Researcher"]
    factor_rules_report: Annotated[
        str, "Summary from the optional factor rule analyst"
    ]
    segment_report: Annotated[str, "Report from the Segment Analyst"]
    scenario_catalyst_report: Annotated[
        str,
        "Report from the Scenario and Catalyst Analyst",
    ]
    position_sizing_report: Annotated[
        str,
        "Report from the Position Sizing Analyst",
    ]
    valuation_data: Annotated[
        ValuationData, "Structured valuation underwriting output"
    ]
    segment_data: Annotated[
        SegmentData, "Structured segment underwriting output"
    ]
    scenario_catalyst_data: Annotated[
        ScenarioCatalystData, "Structured scenario and catalyst underwriting output"
    ]
    position_sizing_data: Annotated[
        PositionSizingData, "Structured position sizing underwriting output"
    ]
    chief_analyst_data: Annotated[
        ChiefAnalystData, "Structured chief analyst summary output"
    ]

    # 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"]