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

base_addresses = ["0x45e7035499bc860dA0b5666A93B97AFc51F3cD3F",'0x8Ff84fa392340d4683C57A602AC400C3D2fF5cdD','0xA6bbEFBB13BBbcA0022a270c734674146DCcf668','0x43D04107E14C19DCDd2C0D0afC33cDa458b349A3','0xA19f7F0643FF7f655D3E39D0689caa3D53944AAE','0x9772B62cC13E86A670a907Ae00D51bfd5674EA78','0xC314D03f9674B31B70e63e19d464ECa6b01C0176','0xF24B37E86867BeDD11b9923172c3FBf0140140f2']

lp_wallets = ["0x45e7035499bc860dA0b5666A93B97AFc51F3cD3F",'0x8Ff84fa392340d4683C57A602AC400C3D2fF5cdD','0xA6bbEFBB13BBbcA0022a270c734674146DCcf668','0x43D04107E14C19DCDd2C0D0afC33cDa458b349A3','0xA19f7F0643FF7f655D3E39D0689caa3D53944AAE']

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

        vice_data = {
          "decimals": 18,
          "logo": 'null',
          "name": "",
          "symbol": "VICE"
        }
        df['tokenMetadata'] = np.where(df['tokenAddress'] == '0xfd409bc96d126bc8a56479d4c7672015d539f96c',vice_data,df['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',0) if len(x) > 0 else 0)
        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')
        eth_price = df.loc[df['symbol']=='ETH','currentPrice'].values[0]
        base_price = df.loc[df['symbol']=='VICE','currentPrice'].values[0]
        df['value'] = df['tokenBalance'] * df['currentPrice']

        df = df[df['symbol'].isin(['VICE', 'ETH', 'USDT', 'USDC'])]
        print(df)
        df_lp_only = df[df['address'] == '0x45e7035499bc860da0b5666a93b97afc51f3cd3f'].copy()
        df_lp_only = df_lp_only[['symbol','tokenBalance','value']]
        df_lp_only = df_lp_only.rename(columns={'tokenBalance':'amount','value':'usd_value'})
        taker_df = df[df['address'] != '0x45e7035499bc860da0b5666a93b97afc51f3cd3f'].copy()
        taker_df= taker_df[['symbol','tokenBalance','value']]
        taker_df = taker_df.rename(columns={'tokenBalance':'amount','value':'usd_value'})
        
        taker_df = taker_df.groupby("symbol", as_index=False).sum()

        return taker_df, taker_df['usd_value'].sum(), eth_price, base_price, df_lp_only

async def fetch_lp():
    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'])
    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 = []
        if df.empty:
            return
        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
    final_balances = get_dex_balances(df)
    return final_balances.T.drop('DEX').reset_index().rename(columns={'index':'symbol',0:'amount'}) if not(final_balances.empty) else pd.DataFrame([{}])
    

async def get_all_dfs():
    starting_df = pd.DataFrame({
    "symbol": ["ETH", "VICE"],
    "amount": [23.77, 11066328.78],
    "usd_value": [111264.43, 448971.41]
    })
    df, total, eth_price, base_price, df_lp_only = await get_wallet_base_balances(base_addresses)
    symbols = {"VICE":base_price,"WETH":eth_price}
    lp_df = await fetch_lp()
    if not lp_df.empty:
        lp_df['price'] = lp_df['symbol'].map(symbols)
        lp_df['price'] = lp_df['price'] * lp_df['amount']
        lp_df.rename(columns={"price":"usd_value"},inplace=True)

    total_current_balances = pd.concat([df,lp_df,df_lp_only]).groupby('symbol').agg({'amount':'sum','usd_value':'sum'}).reset_index()
    for col in ['amount','usd_value']:
        df[col] = pd.to_numeric(df[col], errors='coerce')
        df[col] = df[col].apply(lambda x: f"{x:,.2f}")
        df_lp_only[col] = pd.to_numeric(df_lp_only[col], errors='coerce')
        df_lp_only[col] = df_lp_only[col].apply(lambda x: f"{x:,.2f}")
        lp_df[col] = pd.to_numeric(lp_df[col], errors='coerce')
        lp_df[col] = lp_df[col].apply(lambda x: f"{x:,.2f}")
        total_current_balances[col] = pd.to_numeric(total_current_balances[col], errors='coerce')
        total_current_balances[col] = total_current_balances[col].apply(lambda x: f"{x:,.2f}")

    
    return starting_df,df,lp_df,df_lp_only,total_current_balances

with gr.Blocks() as demo:
    gr.Markdown("## VICE Balances & PnL")

    df0_out = gr.Dataframe(label="Starting Balances")
    df1_out = gr.Dataframe(label="Taker Balances")
    df3_out = gr.Dataframe(label="LP Wallet Balances")
    df2_out = gr.Dataframe(label="LP Position Balances")
    df4_out = gr.Dataframe(label="Total Current Balances")
    
    demo.load(
    fn= get_all_dfs, 
    inputs = [],
    outputs=[df0_out,df1_out,df2_out,df3_out,df4_out]
    )


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