import ccxt
import pandas as pd

def get_ohlcv(symbol, timeframe):
    exchange = ccxt.kraken()
    # Symbol mapping for Kraken
    symbol_map = {
        'BTC/USDT': 'BTC/USD',
        'ETH/USDT': 'ETH/USD',
        'XRP/USDT': 'XRP/USD',
        'LTC/USDT': 'LTC/USD',
        'ADA/USDT': 'ADA/USD',
        'SOL/USDT': 'SOL/USD',
        'DOGE/USDT': 'DOGE/USD',
    }

    try:
        # First, try the original symbol
        ohlcv = exchange.fetch_ohlcv(symbol, timeframe)
    except ccxt.BadSymbol:
        # If the symbol is not found, try the mapped symbol
        mapped_symbol = symbol_map.get(symbol)
        if mapped_symbol:
            try:
                ohlcv = exchange.fetch_ohlcv(mapped_symbol, timeframe)
            except ccxt.ExchangeError as e:
                raise Exception(f"Could not fetch data for {symbol} or {mapped_symbol} from Kraken: {e})")
        else:
            raise Exception(f"Symbol {symbol} not found on Kraken and no mapping available.")

    df = pd.DataFrame(ohlcv, columns=['timestamp', 'open', 'high', 'low', 'close', 'volume'])
    df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
    return df

def get_trend(df):
    df['ema20'] = df['close'].ewm(span=20, adjust=False).mean()
    df['ema50'] = df['close'].ewm(span=50, adjust=False).mean()
    df['rsi'] = 100 - (100 / (1 + df['close'].diff().apply(lambda x: x if x > 0 else 0).ewm(alpha=1/14, adjust=False).mean() / df['close'].diff().apply(lambda x: abs(x) if x < 0 else 0).ewm(alpha=1/14, adjust=False).mean()))

    last_row = df.iloc[-1]
    if last_row['ema20'] > last_row['ema50'] and last_row['close'] > last_row['ema20'] and last_row['rsi'] > 50:
        return 'UP'
    elif last_row['ema20'] < last_row['ema50'] and last_row['close'] < last_row['ema20'] and last_row['rsi'] < 50:
        return 'DOWN'
    else:
        return None

def get_setup(df):
    df['rsi'] = 100 - (100 / (1 + df['close'].diff().apply(lambda x: x if x > 0 else 0).ewm(alpha=1/14, adjust=False).mean() / df['close'].diff().apply(lambda x: abs(x) if x < 0 else 0).ewm(alpha=1/14, adjust=False).mean()))

    last_row = df.iloc[-1]
    if last_row['rsi'] > 40 and last_row['rsi'] < 60:
        return True
    else:
        return False

def get_entry(df):
    df['ema20'] = df['close'].ewm(span=20, adjust=False).mean()
    df['ema50'] = df['close'].ewm(span=50, adjust=False).mean()
    df['rsi'] = 100 - (100 / (1 + df['close'].diff().apply(lambda x: x if x > 0 else 0).ewm(alpha=1/14, adjust=False).mean() / df['close'].diff().apply(lambda x: abs(x) if x < 0 else 0).ewm(alpha=1/14, adjust=False).mean()))
    df['volume_sma20'] = df['volume'].rolling(window=20).mean()

    last_row = df.iloc[-1]
    if last_row['ema20'] > last_row['ema50'] and last_row['rsi'] > 50 and last_row['close'] > last_row['ema20'] and last_row['volume'] > last_row['volume_sma20']:
        return 'BUY'
    elif last_row['ema20'] < last_row['ema50'] and last_row['rsi'] < 50 and last_row['close'] < last_row['ema20'] and last_row['volume'] > last_row['volume_sma20']:
        return 'SELL'
    else:
        return 'WAIT'

def generate_signal(symbol):
    # 1h trend
    df_1h = get_ohlcv(symbol, '1h')
    trend = get_trend(df_1h)

    if trend:
        # 15m setup
        df_15m = get_ohlcv(symbol, '15m')
        setup = get_setup(df_15m)

        if setup:
            # 5m entry
            df_5m = get_ohlcv(symbol, '5m')
            entry = get_entry(df_5m)

            if (trend == 'UP' and entry == 'BUY') or (trend == 'DOWN' and entry == 'SELL'):
                return entry

    return 'WAIT'