discord-stock-bot/bot.py

import discord
from discord.ext import commands
import logging
from datetime import datetime
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger
import pytz

from config import Config
from stock_api import YahooFinanceAPI, FinnhubAPI
from market_hours import MarketHours
from crypto_api import CoinGeckoAPI
from chart_generator import ChartGenerator


# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class StockBot(commands.Bot):
    """Discord bot for stock price tracking and reporting."""

    def __init__(self):
        intents = discord.Intents.default()
        intents.message_content = True
        super().__init__(command_prefix=Config.COMMAND_PREFIX, intents=intents)

        # Initialize stock API based on configuration
        if Config.STOCK_API_PROVIDER == 'finnhub':
            self.stock_api = FinnhubAPI(Config.FINNHUB_API_KEY)
            logger.info("Using Finnhub API for stock data")
        else:
            self.stock_api = YahooFinanceAPI()
            logger.info("Using Yahoo Finance API for stock data")

        # Always initialize Yahoo Finance for chart data (free historical data)
        self.yahoo_api = YahooFinanceAPI()
        logger.info("Using Yahoo Finance API for chart data")

        self.scheduler = AsyncIOScheduler(timezone=pytz.timezone('America/New_York'))
        self.target_channel_ids = [int(id) for id in Config.CHANNEL_IDS]
        self.primary_ticker = Config.PRIMARY_TICKER
        self.startup_announced = False
        self.crypto_api = CoinGeckoAPI()

    async def setup_hook(self):
        """Initialize the bot and set up scheduled tasks."""
        logger.info("Bot setup initiated")

        # Schedule market open update at 9:30 AM ET on weekdays
        self.scheduler.add_job(
            self.send_market_open_update,
            CronTrigger(hour=9, minute=30, day_of_week='mon-fri', timezone=MarketHours.NYSE_TIMEZONE),
            id='market_open_update',
            name='Market Open Update'
        )

        # Schedule noon update at 12:00 PM ET on weekdays
        self.scheduler.add_job(
            self.send_noon_update,
            CronTrigger(hour=12, minute=0, day_of_week='mon-fri', timezone=MarketHours.NYSE_TIMEZONE),
            id='noon_update',
            name='Noon Update'
        )

        # Schedule market close update at 4:00 PM ET on weekdays
        self.scheduler.add_job(
            self.send_market_close_update,
            CronTrigger(hour=16, minute=0, day_of_week='mon-fri', timezone=MarketHours.NYSE_TIMEZONE),
            id='market_close_update',
            name='Market Close Update'
        )

        # Schedule weekly earnings calendar
        self.scheduler.add_job(
            self.send_weekly_earnings_calendar,
            CronTrigger(
                day_of_week=Config.EARNINGS_SCHEDULE_DAY,
                hour=Config.EARNINGS_SCHEDULE_HOUR,
                minute=0,
                timezone=MarketHours.NYSE_TIMEZONE
            ),
            id='weekly_earnings_calendar',
            name='Weekly Earnings Calendar'
        )

        self.scheduler.start()
        logger.info("Scheduler started for market open, noon, market close, and weekly earnings")

    async def on_ready(self):
        """Called when bot successfully connects to Discord."""
        logger.info(f'Bot connected as {self.user.name} (ID: {self.user.id})')
        logger.info(f'Monitoring ticker: {self.primary_ticker}')
        logger.info(f'Target channel IDs: {self.target_channel_ids}')

        # Verify channels exist
        for channel_id in self.target_channel_ids:
            channel = self.get_channel(channel_id)
            if channel:
                logger.info(f'Target channel found: #{channel.name} (ID: {channel_id})')
            else:
                logger.error(f'Could not find channel with ID {channel_id}')

        # Send startup announcement if configured and not already sent
        if Config.STARTUP_ANNOUNCEMENT and not self.startup_announced:
            await self.send_startup_announcement()
            self.startup_announced = True

    async def send_market_open_update(self):
        """Send stock price update at market open on trading days."""
        if not MarketHours.is_trading_day():
            logger.info("Not a trading day, skipping market open update")
            return

        logger.info(f"Sending market open update for {self.primary_ticker}")
        for channel_id in self.target_channel_ids:
            await self.send_stock_update(self.primary_ticker, channel_id)

    async def send_noon_update(self):
        """Send stock price update at noon on trading days."""
        if not MarketHours.is_trading_day():
            logger.info("Not a trading day, skipping noon update")
            return

        logger.info(f"Sending noon update for {self.primary_ticker}")
        for channel_id in self.target_channel_ids:
            await self.send_stock_update(self.primary_ticker, channel_id)

    async def send_market_close_update(self):
        """Send stock price update at market close on trading days."""
        if not MarketHours.is_trading_day():
            logger.info("Not a trading day, skipping market close update")
            return

        logger.info(f"Sending market close update for {self.primary_ticker}")
        for channel_id in self.target_channel_ids:
            await self.send_stock_update(self.primary_ticker, channel_id)

    async def send_weekly_earnings_calendar(self):
        """Send weekly earnings calendar to target channels."""
        from collections import defaultdict
        from datetime import timedelta
        import asyncio

        logger.info("Sending weekly earnings calendar")

        # Calculate upcoming week (Monday-Friday)
        now_et = datetime.now(MarketHours.NYSE_TIMEZONE)

        # Find next Monday
        days_until_monday = (7 - now_et.weekday()) % 7
        if days_until_monday == 0:
            days_until_monday = 7

        start_date = now_et.date() + timedelta(days=days_until_monday)
        end_date = start_date + timedelta(days=4)  # Friday

        from_date_str = start_date.strftime('%Y-%m-%d')
        to_date_str = end_date.strftime('%Y-%m-%d')

        logger.info(f"Fetching earnings from {from_date_str} to {to_date_str}")

        # Check if using Finnhub API
        if Config.STOCK_API_PROVIDER != 'finnhub':
            logger.warning("Earnings calendar requires Finnhub API provider")
            return

        # Fetch earnings data
        min_market_cap = Config.EARNINGS_MIN_MARKET_CAP_BILLIONS * 1_000_000_000
        earnings_list = self.stock_api.get_earnings_calendar(from_date_str, to_date_str, min_market_cap)

        if earnings_list is None:
            logger.error("Failed to retrieve earnings calendar")
            return

        if not earnings_list:
            logger.info(f"No major earnings scheduled for {from_date_str} to {to_date_str}")
            # Optionally send a message saying no earnings this week
            return

        # Group earnings by day
        earnings_by_day = defaultdict(list)
        for earning in earnings_list:
            date = earning['date']
            earnings_by_day[date].append(earning)

        # Create embeds
        embeds = self.create_earnings_embeds(earnings_by_day, Config.EARNINGS_MIN_MARKET_CAP_BILLIONS)

        if not embeds:
            logger.warning("No embeds created for earnings calendar")
            return

        # Send to each target channel
        for channel_id in self.target_channel_ids:
            channel = self.get_channel(channel_id)
            if not channel:
                logger.error(f"Could not find channel {channel_id}")
                continue

            # Send embeds with delay
            for embed in embeds:
                await channel.send(embed=embed)
                await asyncio.sleep(0.5)

            logger.info(f"Sent {len(embeds)} earnings calendar embeds to channel {channel_id}")

    async def send_startup_announcement(self):
        """Send startup announcement to configured announcement channel."""
        embed = discord.Embed(
            title="🤖 Bot Restarted",
            description=Config.STARTUP_ANNOUNCEMENT,
            color=discord.Color.blue(),
            timestamp=datetime.now(pytz.timezone('America/New_York'))
        )

        # Use dedicated announcement channel if configured, otherwise use first target channel
        announcement_channel_id = Config.STARTUP_ANNOUNCEMENT_CHANNEL_ID
        if announcement_channel_id:
            try:
                channel_id = int(announcement_channel_id)
            except ValueError:
                logger.error(f"Invalid STARTUP_ANNOUNCEMENT_CHANNEL_ID: {announcement_channel_id}")
                return
        else:
            channel_id = self.target_channel_ids[0] if self.target_channel_ids else None

        if not channel_id:
            logger.error("No channel configured for startup announcement")
            return

        channel = self.get_channel(channel_id)
        if channel:
            await channel.send(embed=embed)
            logger.info(f"Sent startup announcement to channel {channel_id}")
        else:
            logger.error(f"Could not find announcement channel {channel_id}")

    async def send_stock_update(self, ticker: str, channel_id: int):
        """
        Fetch stock data and send formatted embed to specified channel.

        Args:
            ticker: Stock ticker symbol
            channel_id: Discord channel ID to send message to
        """
        channel = self.get_channel(channel_id)
        if not channel:
            logger.error(f"Could not find channel {channel_id}")
            return

        stock_data = self.stock_api.get_stock_price(ticker)
        if not stock_data:
            await channel.send(f"Unable to retrieve data for ticker: {ticker}")
            return

        embed = self.create_stock_embed(stock_data)

        # Generate chart using Yahoo Finance historical data
        chart_file = None
        candle_data = self.yahoo_api.get_candles(ticker, days=30)
        if candle_data:
            chart_buffer = ChartGenerator.create_price_chart(
                ticker,
                candle_data,
                stock_data.get('company_name')
            )
            if chart_buffer:
                chart_file = discord.File(chart_buffer, filename="chart.png")

        # Send with chart attachment if available
        if chart_file:
            await channel.send(embed=embed, file=chart_file)
        else:
            await channel.send(embed=embed)

        logger.info(f"Sent stock update for {ticker} to channel {channel_id}")

    def create_stock_embed(self, stock_data: dict) -> discord.Embed:
        """
        Create a formatted Discord embed for stock price data.

        Args:
            stock_data: Dictionary containing stock information

        Returns:
            Discord embed object
        """
        ticker = stock_data['ticker']
        company_name = stock_data.get('company_name', ticker)

        # Custom company name for AMD
        if ticker == "AMD":
            company_name = "Advanced Money Destroyer"

        current_price = stock_data['current_price']
        change_dollar = stock_data['change_dollar']
        change_percent = stock_data['change_percent']

        # Determine color based on price movement
        if change_dollar > 0:
            color = discord.Color.green()
            change_emoji = "📈"
        elif change_dollar < 0:
            color = discord.Color.red()
            change_emoji = "📉"
        else:
            color = discord.Color.blue()
            change_emoji = "➡️"

        # Format change string
        change_sign = "+" if change_dollar >= 0 else ""
        change_str = f"{change_sign}${change_dollar} ({change_sign}{change_percent}%)"

        # Custom title for PYPL ticker
        if ticker == "PYPL":
            title = f"{change_emoji} Steaks Stablecoin Value"
        else:
            title = f"{change_emoji} {company_name} ({ticker}) Stock Price"

        # Create embed
        embed = discord.Embed(
            title=title,
            description=f"**${current_price}**",
            color=color,
            timestamp=datetime.now(pytz.timezone('America/New_York'))
        )

        embed.add_field(name="Change", value=change_str, inline=True)
        embed.add_field(name="Previous Close", value=f"${stock_data['previous_close']}", inline=True)

        # Reference attached chart image
        embed.set_image(url="attachment://chart.png")

        market_status = "🟢 Market Open" if stock_data['market_open'] else "🔴 Market Closed"
        embed.set_footer(text=f"{market_status}")

        return embed

    def create_crypto_embed(self, crypto_data: dict) -> discord.Embed:
        """
        Create a formatted Discord embed for cryptocurrency price data.

        Args:
            crypto_data: Dictionary containing crypto information

        Returns:
            Discord embed object
        """
        symbol = crypto_data['symbol']
        current_price = crypto_data['current_price']
        change_dollar = crypto_data['change_dollar']
        change_percent = crypto_data['change_percent']

        # Determine color based on price movement
        if change_dollar > 0:
            color = discord.Color.green()
            change_emoji = "📈"
        elif change_dollar < 0:
            color = discord.Color.red()
            change_emoji = "📉"
        else:
            color = discord.Color.blue()
            change_emoji = "➡️"

        # Format change string
        change_sign = "+" if change_dollar >= 0 else ""
        change_str = f"{change_sign}${change_dollar} ({change_sign}{change_percent}%)"

        # Create embed
        embed = discord.Embed(
            title=f"{change_emoji} {symbol} Price",
            description=f"**${current_price:,.2f}**",
            color=color,
            timestamp=datetime.now(pytz.timezone('America/New_York'))
        )

        embed.add_field(name="24h Change", value=change_str, inline=True)
        embed.add_field(name="Previous Price", value=f"${crypto_data['previous_price']:,.2f}", inline=True)

        # Add volume if available
        volume = crypto_data.get('volume_24h', 0)
        if volume > 0:
            volume_str = f"${volume:,.0f}"
            embed.add_field(name="24h Volume", value=volume_str, inline=False)

        embed.set_footer(text="Data from CoinGecko")

        return embed

    def create_earnings_embeds(self, earnings_by_day: dict, min_market_cap_billions: float) -> list:
        """
        Create Discord embeds for earnings calendar, grouped by day.

        Args:
            earnings_by_day: Dictionary mapping date strings to list of earnings
            min_market_cap_billions: Minimum market cap filter value for footer

        Returns:
            List of discord.Embed objects
        """
        from collections import defaultdict

        embeds = []

        # Process each day
        for date_str in sorted(earnings_by_day.keys()):
            earnings_for_day = earnings_by_day[date_str]

            if not earnings_for_day:
                continue

            # Format date for title
            try:
                date_obj = datetime.strptime(date_str, '%Y-%m-%d')
                day_name = date_obj.strftime('%A, %B %d, %Y')
            except:
                day_name = date_str

            # Build description
            description_lines = []

            for earning in earnings_for_day:
                symbol = earning['symbol']
                company_name = earning['company_name']
                market_cap = earning['market_cap']
                eps_estimate = earning.get('eps_estimate')
                revenue_estimate = earning.get('revenue_estimate')
                hour = earning.get('hour', 'TBD')

                # Format market cap
                if market_cap >= 1_000_000_000_000:  # Trillions
                    market_cap_str = f"${market_cap / 1_000_000_000_000:.1f}T"
                elif market_cap >= 1_000_000_000:  # Billions
                    market_cap_str = f"${market_cap / 1_000_000_000:.1f}B"
                else:
                    market_cap_str = f"${market_cap / 1_000_000:.0f}M"

                # Format EPS estimate
                if eps_estimate is not None:
                    eps_str = f"${eps_estimate:.2f}"
                else:
                    eps_str = "N/A"

                # Format revenue estimate (usually in millions)
                if revenue_estimate is not None:
                    if revenue_estimate >= 1_000:  # Billions
                        rev_str = f"${revenue_estimate / 1_000:.1f}B"
                    else:
                        rev_str = f"${revenue_estimate:.0f}M"
                else:
                    rev_str = "N/A"

                # Determine timing emoji
                if hour and 'bmo' in hour.lower():
                    time_emoji = "🌅"
                    time_str = "Before Market Open"
                elif hour and 'amc' in hour.lower():
                    time_emoji = "🌆"
                    time_str = "After Market Close"
                else:
                    time_emoji = "📅"
                    time_str = "TBD"

                # Build entry
                entry = f"**{symbol}** - {company_name} ({market_cap_str})\n"
                entry += f"├ EPS Est: {eps_str} | Rev Est: {rev_str}\n"
                entry += f"└ {time_emoji} {time_str}\n\n"

                description_lines.append(entry)

            # Join all entries for this day
            description = "".join(description_lines)

            # Check if description exceeds Discord limit (4096 chars)
            if len(description) > 4000:
                # Split into multiple embeds if needed
                current_desc = ""
                part_num = 1

                for entry in description_lines:
                    if len(current_desc) + len(entry) > 4000:
                        # Create embed for current part
                        embed = discord.Embed(
                            title=f"📊 Earnings Calendar - {day_name} (Part {part_num})",
                            description=current_desc,
                            color=discord.Color.blue(),
                            timestamp=datetime.now(pytz.timezone('America/New_York'))
                        )
                        embed.set_footer(text=f"Data from Finnhub | Min Market Cap: ${min_market_cap_billions:.0f}B")
                        embeds.append(embed)

                        # Start new part
                        current_desc = entry
                        part_num += 1
                    else:
                        current_desc += entry

                # Add final part
                if current_desc:
                    embed = discord.Embed(
                        title=f"📊 Earnings Calendar - {day_name} (Part {part_num})",
                        description=current_desc,
                        color=discord.Color.blue(),
                        timestamp=datetime.now(pytz.timezone('America/New_York'))
                    )
                    embed.set_footer(text=f"Data from Finnhub | Min Market Cap: ${min_market_cap_billions:.0f}B")
                    embeds.append(embed)
            else:
                # Single embed for this day
                embed = discord.Embed(
                    title=f"📊 Earnings Calendar - {day_name}",
                    description=description,
                    color=discord.Color.blue(),
                    timestamp=datetime.now(pytz.timezone('America/New_York'))
                )
                embed.set_footer(text=f"Data from Finnhub | Min Market Cap: ${min_market_cap_billions:.0f}B")
                embeds.append(embed)

        return embeds


# Create bot instance
bot = StockBot()


@bot.command(name='stock', aliases=['price'])
async def get_stock_price(ctx, ticker: str = None):
    """
    Manually query stock price for any ticker.

    Usage: !stock AAPL or !price TSLA
    """
    if not ticker:
        await ctx.send("Please provide a ticker symbol. Usage: `!stock AAPL`")
        return

    ticker = ticker.upper()
    logger.info(f"Manual stock query for {ticker} by {ctx.author}")

    stock_data = bot.stock_api.get_stock_price(ticker)
    if not stock_data:
        await ctx.send(f"Unable to retrieve data for ticker: {ticker}")
        return

    embed = bot.create_stock_embed(stock_data)

    # Generate chart using Yahoo Finance historical data
    chart_file = None
    candle_data = bot.yahoo_api.get_candles(ticker, days=30)
    if candle_data:
        chart_buffer = ChartGenerator.create_price_chart(
            ticker,
            candle_data,
            stock_data.get('company_name')
        )
        if chart_buffer:
            chart_file = discord.File(chart_buffer, filename="chart.png")

    # Send with chart attachment if available
    if chart_file:
        await ctx.send(embed=embed, file=chart_file)
    else:
        await ctx.send(embed=embed)


@bot.command(name='crypto', aliases=['c'])
async def get_crypto_price(ctx, symbol: str = None):
    """
    Query cryptocurrency price.

    Usage: !crypto BTC or !c ETH
    """
    if not symbol:
        await ctx.send("Please provide a crypto symbol. Usage: `!crypto BTC`")
        return

    symbol = symbol.upper()
    logger.info(f"Crypto query for {symbol} by {ctx.author}")

    crypto_data = bot.crypto_api.get_crypto_price(symbol)
    if not crypto_data:
        await ctx.send(f"Unable to retrieve data for crypto: {symbol}")
        return

    embed = bot.create_crypto_embed(crypto_data)
    await ctx.send(embed=embed)


@bot.command(name='earnings', aliases=['earn'])
async def get_earnings_calendar(ctx, timeframe: str = 'week'):
    """
    Display earnings calendar for current/next week.

    Usage: !earnings or !earnings next
    """
    from collections import defaultdict
    from datetime import timedelta
    import asyncio

    logger.info(f"Earnings calendar query by {ctx.author} (timeframe: {timeframe})")

    # Determine date range based on timeframe
    now_et = datetime.now(MarketHours.NYSE_TIMEZONE)

    # Determine if we should show current week or next week
    if timeframe.lower() == 'next':
        # Always show next week
        days_until_monday = (7 - now_et.weekday()) % 7
        if days_until_monday == 0:
            days_until_monday = 7
    else:
        # Show current week, or next week if it's late Friday/weekend
        if now_et.weekday() == 4 and now_et.hour >= 16:  # Friday after 4 PM
            days_until_monday = 3
        elif now_et.weekday() >= 5:  # Weekend
            days_until_monday = (7 - now_et.weekday()) % 7
        else:  # Sunday-Thursday or Friday before 4 PM
            days_until_monday = -now_et.weekday()  # Go back to Monday

    # Calculate Monday to Friday of that week
    start_date = now_et.date() + timedelta(days=days_until_monday)
    end_date = start_date + timedelta(days=4)  # Friday

    from_date_str = start_date.strftime('%Y-%m-%d')
    to_date_str = end_date.strftime('%Y-%m-%d')

    logger.info(f"Fetching earnings from {from_date_str} to {to_date_str}")

    # Check if using Finnhub API
    if Config.STOCK_API_PROVIDER != 'finnhub':
        await ctx.send("Earnings calendar is only available with Finnhub API provider.")
        return

    # Fetch earnings data
    min_market_cap = Config.EARNINGS_MIN_MARKET_CAP_BILLIONS * 1_000_000_000
    earnings_list = bot.stock_api.get_earnings_calendar(from_date_str, to_date_str, min_market_cap)

    if earnings_list is None:
        await ctx.send("Unable to retrieve earnings calendar. Please try again later.")
        return

    if not earnings_list:
        await ctx.send(f"No major earnings (>${ Config.EARNINGS_MIN_MARKET_CAP_BILLIONS:.0f}B) scheduled for {from_date_str} to {to_date_str}.")
        return

    # Group earnings by day
    earnings_by_day = defaultdict(list)
    for earning in earnings_list:
        date = earning['date']
        earnings_by_day[date].append(earning)

    # Create embeds
    embeds = bot.create_earnings_embeds(earnings_by_day, Config.EARNINGS_MIN_MARKET_CAP_BILLIONS)

    if not embeds:
        await ctx.send("No earnings data available to display.")
        return

    # Send embeds with delay to avoid rate limiting
    for embed in embeds:
        await ctx.send(embed=embed)
        await asyncio.sleep(0.5)

    logger.info(f"Sent {len(embeds)} earnings calendar embeds")


@bot.command(name='ping')
async def ping(ctx):
    """Check if bot is responsive."""
    await ctx.send(f'Pong! Latency: {round(bot.latency * 1000)}ms')


@bot.command(name='market')
async def market_status(ctx):
    """Check if market is currently open."""
    is_open = MarketHours.is_market_open()
    status = "🟢 Market is currently OPEN" if is_open else "🔴 Market is currently CLOSED"

    embed = discord.Embed(
        title="NYSE Market Status",
        description=status,
        color=discord.Color.green() if is_open else discord.Color.red(),
        timestamp=datetime.now(pytz.timezone('America/New_York'))
    )

    embed.add_field(
        name="Trading Hours",
        value="Monday-Friday\n9:30 AM - 4:00 PM ET",
        inline=False
    )

    if not is_open:
        next_open = MarketHours.get_next_market_open()
        embed.add_field(
            name="Next Open",
            value=next_open.strftime("%A, %B %d at %I:%M %p ET"),
            inline=False
        )

    await ctx.send(embed=embed)


def main():
    """Main entry point for the bot."""
    if not Config.validate():
        logger.error("Configuration validation failed")
        return

    logger.info("Starting bot...")
    bot.run(Config.DISCORD_TOKEN)


if __name__ == '__main__':
    main()