Ruby-Cogs/unbelievaboat/investments.py

import random
import asyncio
from typing import Optional
import discord
from redbot.core import commands, bank, Config
from redbot.core.utils.chat_formatting import humanize_number
from datetime import datetime, timedelta

class Investments(commands.Cog):
    """Investment commands for the UnbelievaBoat economy"""
    
    def __init__(self, bot):
        self.bot = bot
        self.config = Config.get_conf(self, identifier=95932766180344, force_registration=True)
        
        default_global = {
            "stocks": {
                "APPLE": {"price": 150.0, "volatility": 0.05},
                "GOOGLE": {"price": 2800.0, "volatility": 0.04},
                "TESLA": {"price": 900.0, "volatility": 0.08},
                "AMAZON": {"price": 3300.0, "volatility": 0.04},
                "MICROSOFT": {"price": 300.0, "volatility": 0.03}
            },
            "crypto": {
                "BTC": {"price": 45000.0, "volatility": 0.15},
                "ETH": {"price": 3000.0, "volatility": 0.12},
                "DOGE": {"price": 0.20, "volatility": 0.25}
            },
            "bonds": {
                "government": {"interest": 0.05, "term": 7},  # 5% interest, 7 day term
                "corporate": {"interest": 0.08, "term": 14},  # 8% interest, 14 day term
                "junk": {"interest": 0.15, "term": 30}  # 15% interest, 30 day term
            },
            "last_update": None
        }
        
        default_user = {
            "stocks": {},  # {"APPLE": {"amount": 10, "avg_price": 150.0}}
            "crypto": {},
            "bonds": [],  # [{"type": "government", "amount": 1000, "purchase_date": "2023-01-01", "maturity_date": "2023-01-08"}]
            "total_invested": 0
        }
        
        self.config.register_global(**default_global)
        self.config.register_user(**default_user)
        
        # Start the background task for updating prices
        self.price_update_task = self.bot.loop.create_task(self.update_prices())

    def cog_unload(self):
        if self.price_update_task:
            self.price_update_task.cancel()

    async def update_prices(self):
        """Background task to update stock and crypto prices"""
        while True:
            try:
                stocks = await self.config.stocks()
                crypto = await self.config.crypto()
                
                # Update stock prices
                for symbol in stocks:
                    volatility = stocks[symbol]["volatility"]
                    current_price = stocks[symbol]["price"]
                    change = random.uniform(-volatility, volatility)
                    new_price = current_price * (1 + change)
                    stocks[symbol]["price"] = round(new_price, 2)
                
                # Update crypto prices
                for symbol in crypto:
                    volatility = crypto[symbol]["volatility"]
                    current_price = crypto[symbol]["price"]
                    change = random.uniform(-volatility, volatility)
                    new_price = current_price * (1 + change)
                    crypto[symbol]["price"] = round(new_price, 2)
                
                await self.config.stocks.set(stocks)
                await self.config.crypto.set(crypto)
                await self.config.last_update.set(datetime.now().isoformat())
                
                await asyncio.sleep(300)  # Update every 5 minutes
            except Exception as e:
                print(f"Error in price update task: {e}")
                await asyncio.sleep(60)

    @commands.group()
    async def invest(self, ctx):
        """Investment commands"""
        pass

    @invest.command(name="stocks")
    async def view_stocks(self, ctx):
        """View current stock prices"""
        stocks = await self.config.stocks()
        
        embed = discord.Embed(title="Stock Market", color=discord.Color.blue())
        for symbol, data in stocks.items():
            embed.add_field(
                name=symbol,
                value=f"Price: ${data['price']:,.2f}\nVolatility: {data['volatility']*100}%",
                inline=True
            )
        
        await ctx.send(embed=embed)

    @invest.command(name="buy")
    async def buy_stock(self, ctx, symbol: str, amount: int):
        """Buy stocks"""
        symbol = symbol.upper()
        stocks = await self.config.stocks()
        
        if symbol not in stocks:
            return await ctx.send("Invalid stock symbol!")
        
        price = stocks[symbol]["price"]
        total_cost = price * amount
        
        if not await bank.can_spend(ctx.author, total_cost):
            return await ctx.send("You don't have enough money!")
        
        # Process the purchase
        await bank.withdraw_credits(ctx.author, total_cost)
        
        async with self.config.user(ctx.author).stocks() as user_stocks:
            if symbol not in user_stocks:
                user_stocks[symbol] = {"amount": amount, "avg_price": price}
            else:
                old_amount = user_stocks[symbol]["amount"]
                old_avg = user_stocks[symbol]["avg_price"]
                new_amount = old_amount + amount
                new_avg = ((old_amount * old_avg) + (amount * price)) / new_amount
                user_stocks[symbol] = {"amount": new_amount, "avg_price": new_avg}
        
        await ctx.send(f"Successfully bought {amount} shares of {symbol} at ${price:,.2f} each!")

    @invest.command(name="sell")
    async def sell_stock(self, ctx, symbol: str, amount: int):
        """Sell stocks"""
        symbol = symbol.upper()
        stocks = await self.config.stocks()
        
        if symbol not in stocks:
            return await ctx.send("Invalid stock symbol!")
        
        async with self.config.user(ctx.author).stocks() as user_stocks:
            if symbol not in user_stocks or user_stocks[symbol]["amount"] < amount:
                return await ctx.send("You don't own enough shares!")
            
            price = stocks[symbol]["price"]
            total_value = price * amount
            
            # Process the sale
            user_stocks[symbol]["amount"] -= amount
            if user_stocks[symbol]["amount"] == 0:
                del user_stocks[symbol]
            
            await bank.deposit_credits(ctx.author, total_value)
            
            profit = total_value - (amount * user_stocks[symbol]["avg_price"])
            
            await ctx.send(
                f"Successfully sold {amount} shares of {symbol} at ${price:,.2f} each!\n"
                f"Total profit: ${profit:,.2f}"
            )

    @invest.command(name="portfolio")
    async def view_portfolio(self, ctx):
        """View your investment portfolio"""
        stocks = await self.config.stocks()
        user_stocks = await self.config.user(ctx.author).stocks()
        
        if not user_stocks:
            return await ctx.send("You don't have any investments!")
        
        embed = discord.Embed(
            title=f"{ctx.author.name}'s Portfolio",
            color=discord.Color.green()
        )
        
        total_value = 0
        for symbol, data in user_stocks.items():
            current_price = stocks[symbol]["price"]
            amount = data["amount"]
            avg_price = data["avg_price"]
            
            value = current_price * amount
            profit = value - (avg_price * amount)
            profit_percent = (profit / (avg_price * amount)) * 100
            
            total_value += value
            
            embed.add_field(
                name=symbol,
                value=f"Shares: {amount}\n"
                      f"Avg Price: ${avg_price:,.2f}\n"
                      f"Current Price: ${current_price:,.2f}\n"
                      f"Profit: ${profit:,.2f} ({profit_percent:,.1f}%)",
                inline=False
            )
        
        embed.set_footer(text=f"Total Portfolio Value: ${total_value:,.2f}")
        await ctx.send(embed=embed)