bot.py

import os
import io
from telegram import Update
from telegram.ext import ApplicationBuilder, CommandHandler, ContextTypes

user_data = {}


def get_user_portfolio(user_id):
    return user_data.setdefault(user_id, {"balance": 1000, "portfolio": {}})


async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    get_user_portfolio(user_id)
    await update.message.reply_text(
        "Welcome to MiniTradeBot \nType /help to see commands."
    )


async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text("""
    Available commands:

    /balance - check your USD balance
    /portfolio - View your holdings
    /buy <symbol> <amount> - Buy asset (e.g /buy BTC 100)
    /sell <symbol> <amount> - Sell asset
    /chart <symbol> [period] [interval] [style] [options] - Generate LIVE pictorial charts
      - styles: line, candle (default: candle)
      - options: volume, ma20, ma50, ma200
      - examples:
        /chart BTC-USD 1d 5m candle volume
        /chart AAPL 5d 1h line ma20 ma50
        /chart SPY 1mo 1d candle ma200
        /chart ETH-USD 4h 15m
    """
    )


async def balance(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    data = get_user_portfolio(user_id)
    await update.message.reply_text(f"Balance: ${data['balance']:.2f}")


async def portfolio(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    data = get_user_portfolio(user_id)
    if not data["portfolio"]:
        await update.message.reply_text('Your portfolio is empty')
    else:
        msg = "Your portfolio:\n"
        for symbol, amount in data["portfolio"].items():
            msg += f"{symbol}: ${amount:.2f}\n"
        await update.message.reply_text(msg)


async def buy(update: Update, context: ContextTypes.DEFAULT_TYPE):
    try:
        user_id = update.effective_user.id
        data = get_user_portfolio(user_id)
        symbol = context.args[0].upper()
        amount = float(context.args[1])

        if amount > data["balance"]:
            await update.message.reply_text('Insufficient balance.')
            return

        data["balance"] -= amount
        data["portfolio"][symbol] = data["portfolio"].get(symbol, 0) + amount
        await update.message.reply_text(f"Bought ${amount:.2f} of {symbol}.")
    except (IndexError, ValueError):
        await update.message.reply_text("Usage: /buy <symbol> <amount>")


async def sell(update: Update, context: ContextTypes.DEFAULT_TYPE):
    try:
        user_id = update.effective_user.id
        data = get_user_portfolio(user_id)
        symbol = context.args[0].upper()
        amount = float(context.args[1])

        if symbol not in data["portfolio"] or data["portfolio"][symbol] < amount:
            await update.message.reply_text("Not enough holdings.")
            return

        data["portfolio"][symbol] -= amount
        if data["portfolio"][symbol] == 0:
            del data["portfolio"][symbol]
        data["balance"] += amount
        await update.message.reply_text(f"Sold ${amount:.2f} of {symbol}.")

    except (IndexError, ValueError):
        await update.message.reply_text("Usage: /sell <symbol> <amount>")


async def chart(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Generate live pictorial charts using TradingView-compatible data sources"""
    try:
        # Lazy import for chart dependencies
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
            import yfinance as yf
            import io
            import pandas as pd
            from datetime import datetime, timedelta
        except ImportError:
            await update.message.reply_text(
                "Chart dependencies missing. Please install: pip install yfinance matplotlib pandas\n"
                "Run: pip install yfinance matplotlib pandas"
            )
            return

        # Parse arguments
        symbol = context.args[0].upper() if len(context.args) >= 1 else "BTC-USD"
        period = context.args[1] if len(context.args) >= 2 else "1d"
        interval = context.args[2] if len(context.args) >= 3 else "5m"
        style = context.args[3] if len(context.args) >= 4 else "candle"
        
        # Parse additional options
        options = context.args[4:] if len(context.args) >= 5 else []
        options = [opt.lower() for opt in options]  # Normalize to lowercase
        show_volume = "volume" in options
        show_ma = any(opt.startswith("ma") for opt in options)
        
        # Debug info (can be removed later)
        print(f"Debug - Options: {options}, show_volume: {show_volume}, show_ma: {show_ma}")
        
        # Convert symbols to Yahoo Finance format (TradingView compatible)
        symbol_mapping = {
            "BTCUSD": "BTC-USD",
            "ETHUSD": "ETH-USD",
            "BTC-USD": "BTC-USD",
            "ETH-USD": "ETH-USD",
            "AAPL": "AAPL",
            "GOOGL": "GOOGL",
            "MSFT": "MSFT",
            "TSLA": "TSLA",
            "SPY": "SPY",
            "QQQ": "QQQ",
            "NVDA": "NVDA",
            "AMZN": "AMZN"
        }
        
        yf_symbol = symbol_mapping.get(symbol, symbol)
        
        # Fetch live data from Yahoo Finance (same data source as TradingView)
        try:
            data = yf.download(
                tickers=yf_symbol,
                period=period,
                interval=interval,
                progress=False,
                auto_adjust=True
            )
        except Exception as e:
            await update.message.reply_text(f"Error fetching data for {symbol}: {e}")
            return
        
        if data is None or data.empty:
            await update.message.reply_text(
                f"No data found for {symbol}. Try:\n"
                f"• `/chart BTC-USD 1d 5m`\n"
                f"• `/chart AAPL 5d 1h`\n"
                f"• `/chart SPY 1mo 1d`"
            )
            return

        # Create chart image
        img_bytes = io.BytesIO()
        
        if style.lower() in ["candle", "candles"]:
            try:
                import mplfinance as mpf
                
                # Prepare data for mplfinance
                df = data.copy()
                                # Handle multi-level columns from Yahoo Finance
                # Columns might be tuples like ('Close', 'BTC-USD') or simple strings
                if isinstance(df.columns, pd.MultiIndex):
                    # For multi-level columns, flatten them to single level
                    df.columns = df.columns.get_level_values(0)
                    print(f"Debug - Flattened multi-level columns: {list(df.columns)}")
                else:
                    # For simple columns, ensure they're strings
                    df.columns = [str(col) for col in df.columns]
                    print(f"Debug - Simple columns: {list(df.columns)}")

                # mplfinance expects specific column names: Open, High, Low, Close, Volume
                # Map the actual column names to what mplfinance expects
                column_mapping = {}
                for col in df.columns:
                    col_str = str(col).lower().strip()
                    if 'open' in col_str:
                        column_mapping[col] = 'Open'
                    elif 'high' in col_str:
                        column_mapping[col] = 'High'
                    elif 'low' in col_str:
                        column_mapping[col] = 'Low'
                    elif 'close' in col_str:
                        column_mapping[col] = 'Close'
                    elif 'volume' in col_str:
                        column_mapping[col] = 'Volume'
                    else:
                        column_mapping[col] = col_str.title()

                df = df.rename(columns=column_mapping)
                
                # Debug: print column names
                print(f"Debug - Original columns: {list(data.columns)}")
                print(f"Debug - Mapped columns: {list(df.columns)}")
                
                # Check if we have the required columns for candlestick charts
                required_columns = ['Open', 'High', 'Low', 'Close']
                missing_columns = [col for col in required_columns if col not in df.columns]
                
                if missing_columns:
                    await update.message.reply_text(
                        f"Missing required columns for candlestick chart: {missing_columns}\n"
                        f"Available columns: {list(df.columns)}\n"
                        f"Try using line chart instead: /chart {symbol} {period} {interval} line"
                    )
                    return
                
                # Clean the data - remove any rows with NaN values and ensure numeric data
                df_clean = df[required_columns].dropna()
                
                if df_clean.empty:
                    await update.message.reply_text(
                        f"No valid data after cleaning for {symbol}.\n"
                        f"Try a different period or interval."
                    )
                    return
                
                # Ensure all data is numeric
                for col in required_columns:
                    df_clean[col] = pd.to_numeric(df_clean[col], errors='coerce')
                
                # Remove any remaining NaN values after conversion
                df_clean = df_clean.dropna()
                
                if df_clean.empty:
                    await update.message.reply_text(
                        f"No valid numeric data for {symbol}.\n"
                        f"Try a different period or interval."
                    )
                    return
                
                print(f"Debug - Clean data shape: {df_clean.shape}")
                print(f"Debug - Sample data:\n{df_clean.head()}")
                print(f"Debug - Index type: {type(df_clean.index)}")
                print(f"Debug - Index sample: {df_clean.index[:5]}")
                
                # Ensure the index is datetime for mplfinance
                if not isinstance(df_clean.index, pd.DatetimeIndex):
                    try:
                        df_clean.index = pd.to_datetime(df_clean.index)
                        print("Debug - Converted index to datetime")
                    except Exception as e:
                        print(f"Debug - Error converting index: {e}")
                        await update.message.reply_text(
                            f"Error with data format for {symbol}.\n"
                            f"Try a different period or interval."
                        )
                        return
                
                # Add moving averages if requested
                add_plots = []
                if show_ma:
                    try:
                        if "ma20" in options or "ma" in options:
                            ma20 = df_clean['Close'].rolling(window=20).mean()
                            if not ma20.isna().all() and len(ma20.dropna()) > 0:  # Only add if we have valid data
                                add_plots.append(mpf.make_addplot(ma20, color='blue', width=1))
                                print(f"Debug - Added MA20, valid points: {len(ma20.dropna())}")
                        if "ma50" in options or "ma" in options:
                            ma50 = df_clean['Close'].rolling(window=50).mean()
                            if not ma50.isna().all() and len(ma50.dropna()) > 0:  # Only add if we have valid data
                                add_plots.append(mpf.make_addplot(ma50, color='red', width=1))
                                print(f"Debug - Added MA50, valid points: {len(ma50.dropna())}")
                        if "ma200" in options:
                            ma200 = df_clean['Close'].rolling(window=200).mean()
                            if not ma200.isna().all() and len(ma200.dropna()) > 0:  # Only add if we have valid data
                                add_plots.append(mpf.make_addplot(ma200, color='purple', width=1))
                                print(f"Debug - Added MA200, valid points: {len(ma200.dropna())}")
                    except Exception as e:
                        print(f"Debug - Error calculating moving averages: {e}")
                        # Continue without moving averages if they fail
                
                # Chart creation is now handled in the progressive plotting section below
                print(f"Debug - DataFrame info:\n{df_clean.info()}")
                print(f"Debug - Final columns: {list(df_clean.columns)}")
                print(f"Debug - Data types:\n{df_clean.dtypes}")
                print(f"Debug - Sample data:\n{df_clean.head()}")
                
                # Ensure all required columns exist and are numeric
                for col in ['Open', 'High', 'Low', 'Close']:
                    if col not in df_clean.columns:
                        await update.message.reply_text(
                            f"Missing required column: {col}\n"
                            f"Available columns: {list(df_clean.columns)}\n"
                            f"Try using line chart instead: /chart {symbol} {period} {interval} line"
                        )
                        return
                    if not pd.api.types.is_numeric_dtype(df_clean[col]):
                        await update.message.reply_text(
                            f"Column {col} is not numeric: {df_clean[col].dtype}\n"
                            f"Try using line chart instead: /chart {symbol} {period} {interval} line"
                        )
                        return
                
                # Final validation - check for any remaining non-numeric values
                for col in ['Open', 'High', 'Low', 'Close']:
                    if df_clean[col].dtype == 'object':
                        # Try to convert object columns to numeric
                        try:
                            df_clean[col] = pd.to_numeric(df_clean[col], errors='coerce')
                            df_clean = df_clean.dropna()  # Remove any new NaNs
                            if df_clean.empty:
                                await update.message.reply_text(
                                    f"All data became invalid after converting {col} to numeric.\n"
                                    f"Try using line chart instead: /chart {symbol} {period} {interval} line"
                                )
                                return
                        except Exception as e:
                            print(f"Debug - Error converting {col} to numeric: {e}")
                            await update.message.reply_text(
                                f"Column {col} contains non-numeric data that cannot be converted.\n"
                                f"Try using line chart instead: /chart {symbol} {period} {interval} line"
                            )
                            return
                
                # Final data cleanup - ensure everything is exactly right for mplfinance
                df_clean = df_clean.copy()  # Create a fresh copy
                
                # Convert index to datetime if it's not already
                if not isinstance(df_clean.index, pd.DatetimeIndex):
                    df_clean.index = pd.to_datetime(df_clean.index)
                
                # Ensure all numeric columns are float64
                for col in ['Open', 'High', 'Low', 'Close']:
                    df_clean[col] = df_clean[col].astype('float64')
                
                # Remove any remaining NaN values
                df_clean = df_clean.dropna()
                
                if df_clean.empty:
                    await update.message.reply_text(
                        f"No valid data remaining after final cleanup for {symbol}.\n"
                        f"Try a different period or interval."
                    )
                    return
                
                # Debug: print exact data being passed to mplfinance
                print(f"Debug - Final DataFrame shape: {df_clean.shape}")
                print(f"Debug - Final DataFrame index type: {type(df_clean.index)}")
                print(f"Debug - Final DataFrame index sample: {df_clean.index[:3]}")
                print(f"Debug - Final Close column sample: {df_clean['Close'].head()}")
                print(f"Debug - Final Close column dtype: {df_clean['Close'].dtype}")
                print(f"Debug - Final data types:\n{df_clean.dtypes}")
                
                # Create the plot with minimal, safe parameters
                try:
                    print(f"Debug - Creating plot with DataFrame shape: {df_clean.shape}")
                    
                    # Start with absolute minimal parameters
                    basic_kwargs = {
                        "type": "candle",
                        "returnfig": True,
                        "figsize": (10, 6)
                    }
                    
                    print(f"Debug - Basic plot kwargs: {basic_kwargs}")
                    fig, axlist = mpf.plot(df_clean, **basic_kwargs)
                    print("Debug - Basic plot successful")
                    
                    # Now add style if it works
                    try:
                        style_kwargs = {**basic_kwargs, "style": "yahoo"}
                        print(f"Debug - Style plot kwargs: {style_kwargs}")
                        fig, axlist = mpf.plot(df_clean, **style_kwargs)
                        print("Debug - Style plot successful")
                    except Exception as style_error:
                        print(f"Debug - Style plot failed, using basic: {style_error}")
                        fig, axlist = mpf.plot(df_clean, **basic_kwargs)
                    
                    # Now try to add volume if requested
                    if show_volume:
                        try:
                            volume_kwargs = {**style_kwargs, "volume": True}
                            print(f"Debug - Volume plot kwargs: {volume_kwargs}")
                            fig, axlist = mpf.plot(df_clean, **volume_kwargs)
                            print("Debug - Volume plot successful")
                        except Exception as volume_error:
                            print(f"Debug - Volume plot failed, using style: {volume_error}")
                            fig, axlist = mpf.plot(df_clean, **style_kwargs)
                    
                    # Finally try to add moving averages if requested
                    if add_plots and len(add_plots) > 0:
                        try:
                            ma_kwargs = {**style_kwargs}
                            if show_volume:
                                ma_kwargs["volume"] = True
                            ma_kwargs["addplot"] = add_plots
                            print(f"Debug - MA plot kwargs: {ma_kwargs}")
                            fig, axlist = mpf.plot(df_clean, **ma_kwargs)
                            print("Debug - MA plot successful")
                        except Exception as ma_error:
                            print(f"Debug - MA plot failed, using previous: {ma_error}")
                            # Use the previous successful plot
                            pass
                    
                except Exception as e:
                    print(f"Debug - Plot creation failed: {e}")
                    print(f"Debug - Error type: {type(e)}")
                    print(f"Debug - DataFrame info:\n{df_clean.info()}")
                    
                    # Try to create a simple line chart as fallback
                    try:
                        print("Debug - Attempting fallback line chart")
                        fig, ax = plt.subplots(figsize=(10, 6))
                        ax.plot(df_clean.index, df_clean['Close'], color='blue', linewidth=2)
                        ax.set_title(f"{symbol} - {period} {interval} (Line Chart Fallback)")
                        ax.set_ylabel("Price")
                        ax.grid(True, alpha=0.3)
                        
                        # Add current price
                        current_price = float(df_clean['Close'].iloc[-1])
                        ax.axhline(y=current_price, color='green', linestyle='--', alpha=0.7)
                        ax.text(df_clean.index[-1], current_price, f"${current_price:.2f}", 
                               verticalalignment='bottom', fontweight='bold', color='green')
                        
                        axlist = [ax]
                        print("Debug - Fallback line chart successful")
                        
                    except Exception as fallback_error:
                        print(f"Debug - Fallback chart also failed: {fallback_error}")
                        await update.message.reply_text(
                            f"Error creating chart: {e}\n"
                            f"Fallback also failed: {fallback_error}\n"
                            f"Please try a different symbol or time period."
                        )
                        return
                
                # Add current price annotation
                current_price = df_clean['Close'].iloc[-1]
                
                # Ensure current_price is a valid numeric value
                try:
                    current_price_float = float(current_price)
                    current_price_formatted = f"${current_price_float:.2f}"
                except (ValueError, TypeError):
                    print(f"Debug - Invalid current price: {current_price} (type: {type(current_price)})")
                    current_price_float = 0.0
                    current_price_formatted = "$0.00"
                
                # Handle both single axis and list of axes
                if isinstance(axlist, (list, tuple)):
                    ax = axlist[0]
                else:
                    ax = axlist
                
                try:
                    ax.axhline(y=current_price_float, color='green', linestyle='--', alpha=0.7)
                    ax.text(df_clean.index[-1], current_price_float, current_price_formatted, 
                           verticalalignment='bottom', fontweight='bold', color='green')
                except Exception as e:
                    print(f"Debug - Error adding price annotation: {e}")
                    # Continue without annotation if it fails
                
                fig.savefig(img_bytes, format="png", dpi=150, bbox_inches="tight")
                plt.close(fig)
                
            except ImportError:
                await update.message.reply_text(
                    "For candlestick charts, install mplfinance:\n"
                    "pip install mplfinance"
                )
                return
                
        else:
            # Line chart
            fig, ax = plt.subplots(figsize=(10, 6))
            
            # Plot price
            ax.plot(data.index, data['Close'], label='Close Price', color='#1f77b4', linewidth=2)
            
            # Add moving averages
            if show_ma:
                if "ma20" in options or "ma" in options:
                    ma20 = data['Close'].rolling(window=20).mean()
                    ax.plot(data.index, ma20, label='MA20', color='orange', linewidth=1.5)
                if "ma50" in options or "ma" in options:
                    ma50 = data['Close'].rolling(window=50).mean()
                    ax.plot(data.index, ma50, label='MA50', color='red', linewidth=1.5)
                if "ma200" in options:
                    ma200 = data['Close'].rolling(window=200).mean()
                    ax.plot(data.index, ma200, label='MA200', color='purple', linewidth=1.5)
            
            # Add volume if requested
            if show_volume and 'Volume' in data.columns:
                ax2 = ax.twinx()
                ax2.bar(data.index, data['Volume'], alpha=0.3, color='gray', label='Volume')
                ax2.set_ylabel('Volume', color='gray')
                ax2.tick_params(axis='y', labelcolor='gray')
            
            # Add current price line
            current_price = data['Close'].iloc[-1]
            ax.axhline(y=current_price, color='green', linestyle='--', alpha=0.7)
            ax.text(data.index[-1], current_price, f' ${current_price:.2f}', 
                   verticalalignment='bottom', fontweight='bold', color='green')
            
            # Chart styling
            ax.set_title(f"{symbol} - {period} {interval} (Live TradingView Data)", fontsize=14, fontweight='bold')
            ax.set_xlabel("Time", fontsize=12)
            ax.set_ylabel("Price ($)", fontsize=12)
            ax.grid(True, alpha=0.3)
            ax.legend(loc='upper left')
            
            # Format x-axis
            plt.xticks(rotation=45)
            fig.tight_layout()
            
            fig.savefig(img_bytes, format="png", dpi=150, bbox_inches="tight")
            plt.close(fig)
        
        # Send the chart image
        img_bytes.seek(0)
        
        # Create informative caption
        current_price = data['Close'].iloc[-1]
        price_change = data['Close'].iloc[-1] - data['Close'].iloc[-2] if len(data) > 1 else 0
        change_pct = (price_change / data['Close'].iloc[-2] * 100) if len(data) > 1 else 0
        
        caption = f"📊 **{symbol} Live Chart**\n"
        caption += f"💰 **Current Price:** ${current_price:.2f}\n"
        if len(data) > 1:
            change_symbol = "📈" if price_change >= 0 else "📉"
            caption += f"{change_symbol} **Change:** ${price_change:.2f} ({change_pct:.2f}%)\n"
        caption += f"⏰ **Period:** {period} | **Interval:** {interval}\n"
        caption += f"🎯 **Data Source:** TradingView Compatible (Yahoo Finance)"
        
        await update.message.reply_photo(
            photo=img_bytes, 
            caption=caption,
            parse_mode='Markdown'
        )
        
    except Exception as exc:
        await update.message.reply_text(f"Error generating chart: {exc}")


def main():
    token = "7241739709:AAFH6sF3DgsxakqcG1TjrpZu-PFRTn1upAA"

    app = ApplicationBuilder().token(token).build()

    app.add_handler(CommandHandler("start", start))
    app.add_handler(CommandHandler("help", help_command))
    app.add_handler(CommandHandler("balance", balance))
    app.add_handler(CommandHandler("portfolio", portfolio))
    app.add_handler(CommandHandler("buy", buy))
    app.add_handler(CommandHandler("sell", sell))
    app.add_handler(CommandHandler("chart", chart))

    print("Bot is running...")
    app.run_polling()


if __name__ == "__main__":
    main()