import requests
import asyncio
import aiohttp
import pandas as pd
import urllib
import gradio as gr
import os
from pymongo import MongoClient

base_addresses = ["0xCaFe02A66B7fa9E63D5Ea13128cC8A82477AA418",
                  "0x33C65AbDD4DaAc415738e05caFc38574bEfC1278",
                  "0xC13b4E80208c24BaE4fCfC3b6DBc35b818b910D9",
                  "0xcFcd6d25aAD11E40e4B0A158a92fD086A0CCDd69"]

bsc_addresses = ["0x57AbFf97abe0c873893B71809b67826170934F27",
                 "0x60BE905be793C2260d13945dfd4CEEAf45449332",
                 "0x9f07773a51330eAED5DD580A30dDa7b5760C8396",
                 "0x8dC9cF38E22229168F9400ee2427aA6b3F3bDCF8"]

lp_wallets = ["0x57AbFf97abe0c873893B71809b67826170934F27",
             "0x9f07773a51330eAED5DD580A30dDa7b5760C8396",
             "0xC13b4E80208c24BaE4fCfC3b6DBc35b818b910D9",
             "0xcFcd6d25aAD11E40e4B0A158a92fD086A0CCDd69"]

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']

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'] = {}

        new_metadata = {
            'decimals': 18,
            'logo': None,
            'name': "",
            'symbol': "LGCT"
        }

        df.loc[df['tokenAddress'] == '0xd38b305cac06990c0887032a02c03d6839f770a8', 'tokenMetadata'] = [new_metadata]

        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')
        base_price = df.loc[df['symbol']=='LGCT','currentPrice'].values[0]
        df['value'] = df['tokenBalance'] * df['currentPrice']
        df.drop(columns=['address','decimals'],inplace=True)

        df = df[df['symbol'].isin(['LGCT', 'ETH', 'USDT', 'USDC'])]
        df= df[['symbol','tokenBalance','value']]
        df = df.rename(columns={'tokenBalance':'amount','value':'usd_value'})
        
        df = df.groupby("symbol", as_index=False).sum()

        return df, df['usd_value'].sum(),base_price

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)
        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=['address','decimals'],inplace=True)

        df = df[df['symbol'].isin(['LGCT', 'BNB', 'USDT', 'USDC'])]
        df= df[['symbol','tokenBalance','value']]
        df = df.rename(columns={'tokenBalance':'amount','value':'usd_value'})
        
        df = df.groupby("symbol", as_index=False).sum()

        return df, df['usd_value'].sum()

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()
    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)

    return lp_df

async def get_all_dfs():
    bsc_df,bsc_total = await get_wallet_bsc_balances(bsc_addresses)
    base_df,base_total,base_price = await get_wallet_base_balances(base_addresses)
    lp_df = await get_lp_balances()
    placeholder_df = lp_df.drop(columns=['DEX']).sum()
    placeholder_df = pd.DataFrame(placeholder_df)
    placeholder_df = placeholder_df.rename(columns={0: 'amount'}).reset_index().rename(columns={'index': 'symbol'})
    prices = {
        'USDT': 1.0,
        'USDC': 1.0,
        'LGCT': base_price
    }
    placeholder_df['usd_value'] = placeholder_df['symbol'].map(prices) * placeholder_df['amount']
    final_df = pd.concat([bsc_df,base_df,placeholder_df])
    final_df = final_df.groupby("symbol", as_index=False).sum()

    return bsc_df, bsc_total, base_df, base_total, lp_df, placeholder_df['usd_value'].sum(), final_df, final_df['usd_value'].sum()


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

    df0_out = gr.Dataframe(label="BSC Wallet Balances")
    txt0_out = gr.Textbox(label="Total Value")

    df1_out = gr.Dataframe(label="Base Wallet Balances")
    txt1_out = gr.Textbox(label="Total Value")

    df2_out = gr.Dataframe(label="LP Balances")
    txt2_out = gr.Textbox(label="Total Value")

    df3_out = gr.Dataframe(label="Total Balances")
    txt3_out = gr.Textbox(label="Total Value")

    # Load from MongoDB on app load (sync function)
    demo.load(
    fn= get_all_dfs, 
    inputs = [],
    outputs=[df0_out, txt0_out, df1_out, txt1_out, df2_out, txt2_out,df3_out,txt3_out]
    )

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