import requests
import asyncio
import aiohttp
import pandas as pd
import urllib
import gradio as gr
import os

addresses = ["0x398d6075a3Dec0cb8C987893Fc3Ab92F72376310",
             "0x42A2D148Df3021bb541c5834AdD699db9c8cc2ba",
             "0x3cB1ad37FE5C5ab2900dEf2f35B939acFf5924DF",
             "0x55d6f5dF162fd93408e08f88180c323810690Bd7"]

addy_map = {"0x398d6075a3Dec0cb8C987893Fc3Ab92F72376310":"LP Wallet",
            "0x42A2D148Df3021bb541c5834AdD699db9c8cc2ba":"Taker 1",
            "0x3cB1ad37FE5C5ab2900dEf2f35B939acFf5924DF":"Taker 2",
            "0x55d6f5dF162fd93408e08f88180c323810690Bd7":"Taker 3"
            }

addy_map = {k.lower(): v for k, v in addy_map.items()}

price_map = {}

lp_wallets = ["0x398d6075a3Dec0cb8C987893Fc3Ab92F72376310"]

def chunk_addresses(addresses, size):
    return [addresses[i:i + size] for i in range(0, len(addresses), size)]

API_URL = "https://api.g.alchemy.com/data/v1/yhe6L3PXmiENzS1sP9Fu4_T5E3l0QyeB/assets/tokens/by-address"

async def fetch_batch(session, batch,chain):
    json_payload = {
        "addresses": [{"address": addr, "networks": [chain]} for addr in batch],
        "withMetadata": True,
        "withPrices": True,
        "includeNativeTokens": True
    }
    async with session.post(API_URL, json=json_payload) as response:
        response = await response.json()
        return response['data']['tokens']

def dollar_values(amount,symbol):
    price = price_map[symbol]
    amount = float(amount)
    price = float(price)
    res = f"{amount} (${amount*price})"
    return res


async def get_wallet_base_balances(addresses):
    async with aiohttp.ClientSession() as session:
        batches = chunk_addresses(addresses, 3)
        tasks = [fetch_batch(session, batch,"base-mainnet") for batch in batches]
        responses = await asyncio.gather(*tasks)
        first_parts = []
        rest_parts = []

        for res in responses:
            first_parts.extend(res[:3])
            rest_parts.extend(res[3:])

        responses = first_parts + rest_parts
        df = pd.DataFrame(responses)
        for i in range(len(addresses)):
            df.at[i,'tokenMetadata'] = {}

        df['tokenBalance'] = df['tokenBalance'].apply(lambda x : int(x,16))
        df['decimals'] = df['tokenMetadata'].apply(lambda x : x.get('decimals',0))
        df['tokenBalance'] = df['tokenBalance']/(10**df['decimals'])
        for i in range(len(addresses)):
            df.at[i, 'tokenBalance'] = df.at[i, 'tokenBalance'] / (10**18)
        df['symbol'] = df['tokenMetadata'].apply(lambda x : x.get('symbol',''))
        for i in range(len(addresses)):
            df.at[i,'symbol'] = 'ETH'
        df['currentPrice'] = df['tokenPrices'].apply(lambda x : x[0].get('value',None) if len(x) > 0 else None)
        df.drop(columns=['tokenPrices','tokenMetadata','network','tokenAddress'], inplace=True)

        df = df.dropna()
        df['tokenBalance'] = pd.to_numeric(df['tokenBalance'], errors='coerce')
        df = df[df['tokenBalance'] != 0]
        df['tokenBalance'] = pd.to_numeric(df['tokenBalance'], errors='coerce')
        df['currentPrice'] = pd.to_numeric(df['currentPrice'], errors='coerce')
        for i in range(len(df)):
            symbol = df['symbol'].iloc[i]
            price = df['currentPrice'].iloc[i]
            price_map[symbol] = price

        df['value'] = df['tokenBalance'] * df['currentPrice']
        df.drop(columns=['decimals'],inplace=True)

        df = df[df['symbol'].isin(['HYB', 'ETH', 'USDT', 'USDC'])]
        df= df[['address','symbol','tokenBalance']]
        df = df.rename(columns={'tokenBalance':'amount','value':'usd_value'})

        df = df.pivot_table(index='address', columns='symbol', values='amount', aggfunc='sum').reset_index()
        df = df.fillna(0)
        df['name'] = df['address'].str.lower().map(addy_map)
        cols = ['name'] + list(df.columns[:-1])
        df = df[cols]
        total_row = df[df.columns[2:]].sum()
        final_row = pd.DataFrame([{"name":"Total","address":""}])
        total_row = pd.DataFrame(total_row).T
        final_row = pd.concat([final_row,total_row],axis=1)
        df = pd.concat([df,final_row],axis=0)
        df_formatted = df

        cols_to_format = df_formatted.columns[df.columns.get_loc("address") + 1:]

        for col in cols_to_format:
            df_formatted[col] = df_formatted[col].apply(lambda x: dollar_values(x, col))

        #df = df.groupby("symbol", as_index=False).sum()

        return df,df_formatted

async def get_wallet_bsc_balances(addresses):
    async with aiohttp.ClientSession() as session:
        batches = chunk_addresses(addresses, 3)
        tasks = [fetch_batch(session, batch,"bnb-mainnet") for batch in batches]
        responses = await asyncio.gather(*tasks)
        first_parts = []
        rest_parts = []

        for res in responses:
            first_parts.extend(res[:3])
            rest_parts.extend(res[3:])

        responses = first_parts + rest_parts


        df = pd.DataFrame(responses)
        for i in range(len(addresses)):
            df.at[i,'tokenMetadata'] = {}
        
        df['tokenBalance'] = df['tokenBalance'].apply(lambda x : int(x,16))
        df['decimals'] = df['tokenMetadata'].apply(lambda x : x.get('decimals',0))
        df['tokenBalance'] = df['tokenBalance']/(10**df['decimals'])
        for i in range(len(addresses)):
            df.at[i, 'tokenBalance'] = df.at[i, 'tokenBalance'] / (10**18)
        df['symbol'] = df['tokenMetadata'].apply(lambda x : x.get('symbol',''))
        for i in range(len(addresses)):
            df.at[i,'symbol'] = 'BNB'
        
        df['currentPrice'] = df['tokenPrices'].apply(lambda x : x[0].get('value',None) if len(x) > 0 else None)
        for i in range(len(df)):
            symbol = df['symbol'].iloc[i]
            price = df['currentPrice'].iloc[i]
            price_map[symbol] = price
        df.drop(columns=['tokenPrices','tokenMetadata','network','tokenAddress'], inplace=True)
   
        df = df.dropna()
        df['tokenBalance'] = pd.to_numeric(df['tokenBalance'], errors='coerce')
        df = df[df['tokenBalance'] != 0]
        df['tokenBalance'] = pd.to_numeric(df['tokenBalance'], errors='coerce')
        df['currentPrice'] = pd.to_numeric(df['currentPrice'], errors='coerce')
        df['value'] = df['tokenBalance'] * df['currentPrice']
        df.drop(columns=['decimals'],inplace=True)

        df = df[df['symbol'].isin(['HYB', 'BNB', 'USDT', 'USDC'])]

        df= df[['address','symbol','tokenBalance','value']]
        df = df.rename(columns={'tokenBalance':'amount','value':'usd_value'})
        df = df.pivot_table(index='address', columns='symbol', values='amount', aggfunc='sum').reset_index()
        df = df.fillna(0)
        df['name'] = df['address'].str.lower().map(addy_map)
        cols = ['name'] + list(df.columns[:-1])
        df = df[cols]
        total_row = df[df.columns[2:]].sum()
        final_row = pd.DataFrame([{"name":"Total","address":""}])
        total_row = pd.DataFrame(total_row).T
        final_row = pd.concat([final_row,total_row],axis=1)
        df = pd.concat([df,final_row],axis=0)
        df_formatted = df

        cols_to_format = df_formatted.columns[df.columns.get_loc("address") + 1:]

        for col in cols_to_format:
            df_formatted[col] = df_formatted[col].apply(lambda x: dollar_values(x, col))
        
        #df = df.groupby("symbol", as_index=False).sum()

        return df,df_formatted

def agg_balances(lp):
    lps = []
    for i in range(len(lp)):
        two = lp['node.tokens'].iloc[i]
        more_lp = pd.json_normalize(two)
        more_lp['token.balance'] = more_lp['token.balance'].apply(pd.to_numeric,errors='coerce')
        more_lp['token.balanceUSD'] = more_lp['token.balanceUSD'].apply(pd.to_numeric,errors='coerce')
        x = more_lp.groupby('token.symbol')[['token.balance','token.balanceUSD']].sum().reset_index()
        #x.columns = x.iloc[0]  # Set first row as column headers
        #x = x.drop(x.index[0]).reset_index(drop=True)  # Drop the row that became header
        #x = x.apply(pd.to_numeric, errors='coerce')  # Convert all to numeric
        lps.append(x)
    bals_df = pd.concat(lps, ignore_index=True)
    bals_df['token.balance'] = bals_df['token.balance'].apply(pd.to_numeric,errors='coerce')
    bals_df['token.balanceUSD'] = bals_df['token.balanceUSD'].apply(pd.to_numeric,errors='coerce')
    bals_df.fillna(0, inplace=True)
    bals_df = bals_df.groupby('token.symbol').sum().reset_index()
    for i in range(len(bals_df)):
        token_balance = bals_df['token.balance'].iloc[i]
        usd_balance = bals_df['token.balanceUSD'].iloc[i]
        symbol = bals_df['token.symbol'].iloc[i]
        price = usd_balance/token_balance
        price_map[symbol] = price
    return bals_df

def get_dex_balances(df):
    balances = []
    for i in range(len(pd.json_normalize(df['node.positionBalances.edges']))):
        balances.append((df['node.app.displayName'].iloc[i],agg_balances(pd.json_normalize(df['node.positionBalances.edges'].iloc[i]))))

    final_balances = []
    for i in range(len(balances)):
        df = balances[i]
        dex_name = df[0]  # or extract from your grouped index
        df_pivot = df[1].pivot_table(index=None, columns='token.symbol', values='token.balance')
        # Add DEX column and reorder
        df_pivot.insert(0,'DEX', dex_name)
        # Remove the first level of row index (e.g., 'token.symbol')
        df_pivot = df_pivot.reset_index(drop=True)
        df_pivot.columns.name = None 
        final_balances.append(df_pivot)

    final_balances = pd.concat(final_balances, ignore_index=True)
    final_balances.fillna(0, inplace=True)
    return final_balances


async def get_lp_balances():
    url = "https://public.zapper.xyz/graphql"
    headers = {
        "Content-Type": "application/json",
        "x-zapper-api-key": "8fe2c210-66e0-4ef9-9505-96a901c9b042"
    }

    query = """
    query AppBalances($addresses: [Address!]!, $first: Int = 10) {
    portfolioV2(addresses: $addresses) {
        appBalances {
        totalBalanceUSD
        byApp(first: $first) {
            totalCount
            edges {
            node {
                balanceUSD
                app {
                displayName
                imgUrl
                description
                category { name }
                }
                network {
                name
                chainId
                }
                positionBalances(first: 10) {
                edges {
                    node {
                    ... on AppTokenPositionBalance {
                        type
                        symbol
                        balance
                        balanceUSD
                        price
                        groupLabel
                        displayProps {
                        label
                        images
                        }
                    }
                    ... on ContractPositionBalance {
                        type
                        balanceUSD
                        groupLabel
                        tokens {
                        metaType
                        token {
                            ... on BaseTokenPositionBalance {
                            symbol
                            balance
                            balanceUSD
                            }
                        }
                        }
                        displayProps {
                        label
                        images
                        }
                    }
                    }
                }
                }
            }
            }
        }
        }
    }
    }
    """

    variables = {
        "addresses": lp_wallets,
        "first": 5
    }

    payload = {
        "query": query,
        "variables": variables
    }

    response = requests.post(url, json=payload, headers=headers)
    response = response.json()
    df = pd.json_normalize(response['data']['portfolioV2']['appBalances']['byApp']['edges'])
    lp_df = get_dex_balances(df)

    lp_df_formatted = lp_df

    cols_to_format = lp_df_formatted.columns[1:]

    for col in cols_to_format:
        lp_df_formatted[col] = lp_df_formatted[col].apply(lambda x: dollar_values(x, col))
    

    return lp_df,lp_df_formatted

async def get_all_dfs():
    bsc_df,bsc_formatted  = await get_wallet_bsc_balances(addresses)
    base_df,base_formatted = await get_wallet_base_balances(addresses)
    lp_df,lp_df_formatted = await get_lp_balances()

    return bsc_formatted, base_formatted, lp_df_formatted


with gr.Blocks() as demo:
    gr.Markdown("## HYB DEX Balances")

    bsc_df = gr.Dataframe(label="Wallet Balances (BSC)")

    base_df = gr.Dataframe(label="Wallet Balances (Base)")

    lp_df = gr.Dataframe(label="LP Wallet Positions")
    
    # Load from MongoDB on app load (sync function)
    demo.load(
    fn= get_all_dfs, 
    inputs = [],
    outputs=[bsc_df, base_df, lp_df]
    )

demo.launch(debug=True, share=True)