from datetime import datetime
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage
from tradingagents.graph.trading_graph import TradingAgentsGraph
from tradingagents.default_config import DEFAULT_CONFIG

# Load environment variables
load_dotenv()

# Patch failing tools
from langchain_core.tools import tool
import tradingagents.agents.utils.agent_utils as agent_utils

@tool
def mock_get_insider_transactions(ticker: str, curr_date: str = None):
    """Mock insider transactions."""
    return "Insider transactions: None significant."

@tool
def mock_get_indicators(ticker: str, curr_date: str = None):
    """Mock indicators."""
    return "Indicators: RSI 50, MACD positive."

@tool
def mock_get_market_movers():
    """Mock market movers."""
    return "Market Movers: NVDA +5%, TSLA -2%."

@tool
def mock_get_earnings_calendar(curr_date: str = None):
    """Mock earnings calendar."""
    return "Earnings: NVDA reporting soon."

@tool
def mock_get_trending_social():
    """Mock trending social."""
    return "Trending: NVDA, TSLA."

agent_utils.get_insider_transactions = mock_get_insider_transactions
agent_utils.get_indicators = mock_get_indicators
agent_utils.get_market_movers = mock_get_market_movers
agent_utils.get_earnings_calendar = mock_get_earnings_calendar
agent_utils.get_trending_social = mock_get_trending_social

def run_full_system(ticker: str = None):
    print("--- Starting Full Agent System ---")
    
    # Configure to use screening
    config = DEFAULT_CONFIG.copy()
    config["deep_think_llm"] = "gpt-4o-mini"
    config["quick_think_llm"] = "gpt-4o-mini"
    
    # Determine if we should include screening
    include_screening = (ticker is None)
    
    # Initialize graph with screening enabled
    print(f"Initializing TradingAgentsGraph with screening={include_screening}...")
    ta = TradingAgentsGraph(
        include_screening=include_screening, 
        config=config,
        debug=True # Enable debug to see the trace
    )
    
    # Create initial state
    # We use a placeholder ticker since screening will find the real one
    trade_date = datetime.now().strftime("%Y-%m-%d")
    
    # If ticker is provided, use it; otherwise use placeholder for screening
    initial_ticker = ticker if ticker else "PENDING"
    
    initial_state = ta.propagator.create_initial_state(initial_ticker, trade_date)
    
    # Override the initial message to trigger screening ONLY if no ticker provided
    if include_screening:
        initial_state["messages"] = [HumanMessage(content="Find a promising stock to analyze based on today's market movers.")]
    else:
        print(f"Bypassing screening, analyzing ticker: {ticker}")
        # Ensure it is a HumanMessage object for consistency and printing
        initial_state["messages"] = [HumanMessage(content=f"Analyze {ticker}")]
    
    print(f"Invoking graph with initial prompt: {initial_state['messages'][0].content}")
    
    # Run the graph
    # We use stream to see progress
    try:
        for chunk in ta.graph.stream(initial_state, config={"recursion_limit": 50}):
            for node, values in chunk.items():
                print(f"--- Node: {node} ---")
                if "messages" in values:
                    last_msg = values["messages"][-1]
                    if last_msg.content:
                        print(f"Output: {last_msg.content}")
                    if hasattr(last_msg, 'tool_calls') and last_msg.tool_calls:
                        for tc in last_msg.tool_calls:
                            print(f"Tool Call: {tc['name']} (Args: {tc['args']})")
                    
                if "company_of_interest" in values:
                    print(f"Current Ticker: {values['company_of_interest']}")
                    
        print("\n--- Execution Completed ---")
        
    except Exception as e:
        print(f"Execution failed: {e}")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Run the Trading Agent System')
    parser.add_argument('--ticker', type=str, help='Stock ticker to analyze (bypasses screening)')
    args = parser.parse_args()
    
    run_full_system(ticker=args.ticker)