state_fetch.py

#!/usr/bin/env python3
import asyncio
import json
from typing import Dict, Any, Optional
from pathlib import Path

from multiversx_sdk import ProxyNetworkProvider, SmartContractController
from utils.contract_data_fetchers import FarmContractDataFetcher, PairContractDataFetcher, MetaStakingContractDataFetcher
from utils.utils_chain import hex_to_string, WrapperAddress as Address

# Configuration
XEXCHANGE_PROXY = "https://gateway.multiversx.com"

# State folder configuration
STATE_FOLDER = Path("state")

# LP token configurations
LP_TOKEN_CONFIG = {
    'UTKWEGLD-c960d1': {
        'address': 'erd1qqqqqqqqqqqqqpgq0lzzvt2faev4upyf586tg38s84d7zsaj2jpsglugga',
        'state_filename': 'state_pool_UTKWEGLD-c960d1.json'
    },
    'UTKUSH-587b46': {
        'address': 'erd1qqqqqqqqqqqqqpgqlt8sksgnhk98pm2chnjwhz5cat7s5wy72jpsgdrmac',
        'state_filename': 'state_pool_UTKUSH-587b46.json'
    }
}

# Farm token configurations (excluding SUTK)
FARM_TOKENS = {
    'UTKWEGLDF-5b9d50': {
        'address': 'erd1qqqqqqqqqqqqqpgqmaj6qc4c6f9ldkuwllqh7pyd8e9yu3m82jpscc6f9s',
        'state_filename': 'state_farm_UTKWEGLDF-5b9d50.json'
    },
    'UTKWEGLDFL-082aec': {
        'address': 'erd1qqqqqqqqqqqqqpgq3f8jfeg34ujzy0muhe9dvn5yngwgud022jpsscjxgl',
        'state_filename': 'state_farm_UTKWEGLDFL-082aec.json'
    },
    'UTKWEGLDFL-ba26d2': {
        'address': 'erd1qqqqqqqqqqqqqpgq4acurmluezvmhna8tztgcrnwh0l70a2wkp2sfh6jkp',
        'state_filename': 'state_farm_UTKWEGLDFL-ba26d2.json'
    }
}

# Metastaking farm configurations
METASTAKING_FARMS = {
    'METAUTK-291e0e': {
        'address': 'erd1qqqqqqqqqqqqqpgqv692rgsd6teaq78nufvgkgjnmca6anqn2jpsphphj6',
        'state_filename': 'state_farm_metastaking_METAUTK-291e0e.json'
    },
    'METAUTKLK-e6a445': {
        'address': 'erd1qqqqqqqqqqqqqpgq6zat6jefz9z8n54yag8xmz7jw20vlv3x2jpsxms0sd',
        'state_filename': 'state_farm_metastaking_METAUTKLK-e6a445.json'
    },
    'METAUTKLK-112f52': {
        'address': 'erd1qqqqqqqqqqqqqpgqmgd0eu4z9kzvrputt4l4kw4fqf2wcjsekp2sftan7s',
        'state_filename': 'state_farm_metastaking_METAUTKLK-112f52.json'
    }
}

class StateFetcher:
    """Centralized state fetching for XExchange contracts"""
    
    def __init__(self, proxy_url: str = XEXCHANGE_PROXY):
        self.proxy_url = proxy_url
    
    async def get_lp_pool_state(self, token_identifier: str) -> Optional[Dict[str, Any]]:
        """Get LP pool contract state using SDK"""
        if token_identifier not in LP_TOKEN_CONFIG:
            print(f"Unknown LP token: {token_identifier}")
            return None
        
        try:
            pool_address = LP_TOKEN_CONFIG[token_identifier]['address']
            print(f"Getting LP pool state for {token_identifier} at {pool_address}")
            
            # Create pair data fetcher
            data_fetcher = PairContractDataFetcher(Address(pool_address), self.proxy_url)
            
            # Get reserves and total supply
            reserves_and_supply = data_fetcher.get_data("getReservesAndTotalSupply")
            
            if len(reserves_and_supply) >= 3:
                pool_state = {
                    'token': token_identifier,
                    'address': pool_address,
                    'reserveFirstToken': str(reserves_and_supply[0]),
                    'reserveSecondToken': str(reserves_and_supply[1]),
                    'lpTokenSupply': str(reserves_and_supply[2]),
                    'firstTokenId': data_fetcher.get_data("getFirstTokenId"),
                    'secondTokenId': data_fetcher.get_data("getSecondTokenId"),
                    'lpTokenIdentifier': data_fetcher.get_data("getLpTokenIdentifier")
                }
                
                print(f"LP pool state retrieved successfully")
                return pool_state
            else:
                print(f"Invalid reserves data: {reserves_and_supply}")
                return None
                
        except Exception as e:
            print(f"Error getting LP pool state for {token_identifier}: {e}")
            return None
    
    async def get_farm_state(self, token_identifier: str) -> Optional[Dict[str, Any]]:
        """Get farm contract state using SDK"""
        if token_identifier not in FARM_TOKENS:
            print(f"Unknown farm token: {token_identifier}")
            return None
        
        try:
            farm_address = FARM_TOKENS[token_identifier]['address']
            print(f"Getting farm state for {token_identifier} at {farm_address}")
            
            # Create farm data fetcher
            data_fetcher = FarmContractDataFetcher(Address(farm_address), self.proxy_url)
            
            # Get farm data
            farm_state = {
                'token': token_identifier,
                'address': farm_address,
                'farmTokenSupply': data_fetcher.get_data("getFarmTokenSupply"),
                'farmingTokenReserve': data_fetcher.get_data("getFarmingTokenReserve"),
                'rewardReserve': data_fetcher.get_data("getRewardReserve"),
                'farmTokenId': data_fetcher.get_data("getFarmTokenId"),
                'farmingTokenId': data_fetcher.get_data("getFarmingTokenId"),
                'rewardTokenId': data_fetcher.get_data("getRewardTokenId"),
                'state': data_fetcher.get_data("getState"),
                'pairContractManagedAddress': data_fetcher.get_data("getPairContractManagedAddress")
            }
            
            print(f"Farm state retrieved successfully")
            return farm_state
            
        except Exception as e:
            print(f"Error getting farm state for {token_identifier}: {e}")
            return None
    
    async def get_metastaking_farm_state(self, token_identifier: str) -> Optional[Dict[str, Any]]:
        """Get metastaking farm contract state using SDK"""
        if token_identifier not in METASTAKING_FARMS:
            print(f"Unknown metastaking farm token: {token_identifier}")
            return None
        
        try:
            farm_address = METASTAKING_FARMS[token_identifier]['address']
            print(f"Getting metastaking farm state for {token_identifier} at {farm_address}")
            
            # Create metastaking data fetcher
            data_fetcher = MetaStakingContractDataFetcher(Address(farm_address), self.proxy_url)
            
            # Get metastaking-specific data
            farm_state = {
                'token': token_identifier,
                'address': farm_address,
                'dualYieldTokenId': data_fetcher.get_data("getDualYieldTokenId"),
                'farmTokenId': data_fetcher.get_data("getFarmTokenId"),
                'lpFarmAddress': data_fetcher.get_data("getLpFarmAddress"),
                'lpFarmTokenId': data_fetcher.get_data("getLpFarmTokenId"),
                'lpTokenId': data_fetcher.get_data("getLpTokenId"),
                'pairAddress': data_fetcher.get_data("getPairAddress"),
                'stakingFarmAddress': data_fetcher.get_data("getStakingFarmAddress"),
                'stakingTokenId': data_fetcher.get_data("getStakingTokenId"),
                'energyFactoryAddress': data_fetcher.get_data("getEnergyFactoryAddress")
            }
            
            print(f"Metastaking farm state retrieved successfully")
            return farm_state
            
        except Exception as e:
            print(f"Error getting metastaking farm state for {token_identifier}: {e}")
            return None
    
    async def save_state_to_json(self, state_data: Dict[str, Any], filename: str):
        """Save state data to JSON file in the state folder"""
        try:
            # Ensure state folder exists
            STATE_FOLDER.mkdir(exist_ok=True)
            
            # Create full path to state folder
            file_path = STATE_FOLDER / filename
            
            with open(file_path, 'w') as f:
                json.dump(state_data, f, indent=2)
            print(f"State saved to {file_path}")
        except Exception as e:
            print(f"Error saving state to {filename}: {e}")
    
    async def fetch_and_save_lp_state(self, token_identifier: str) -> Optional[Dict[str, Any]]:
        """Fetch LP state and save to JSON file"""
        state = await self.get_lp_pool_state(token_identifier)
        if state:
            filename = LP_TOKEN_CONFIG[token_identifier]['state_filename']
            await self.save_state_to_json(state, filename)
        return state
    
    async def fetch_and_save_farm_state(self, token_identifier: str) -> Optional[Dict[str, Any]]:
        """Fetch farm state and save to JSON file"""
        state = await self.get_farm_state(token_identifier)
        if state:
            filename = FARM_TOKENS[token_identifier]['state_filename']
            await self.save_state_to_json(state, filename)
        return state
    
    async def fetch_and_save_metastaking_farm_state(self, token_identifier: str) -> Optional[Dict[str, Any]]:
        """Fetch metastaking farm state and save to JSON file"""
        state = await self.get_metastaking_farm_state(token_identifier)
        if state:
            filename = METASTAKING_FARMS[token_identifier]['state_filename']
            await self.save_state_to_json(state, filename)
            
            # Also fetch the underlying farm states that this metastaking farm depends on
            if state.get('stakingFarmAddress'):
                print(f"Fetching underlying staking farm at {state['stakingFarmAddress']}")
                # TODO: Add logic to fetch staking farm state
                
            if state.get('lpFarmAddress'):
                print(f"Fetching underlying LP farm at {state['lpFarmAddress']}")
                # TODO: Add logic to fetch LP farm state
                
        return state
    
    async def fetch_all_states(self) -> Dict[str, Any]:
        print("=== Fetching All Contract States ===")
        
        states = {}
        
        # Fetch LP pool states
        for token_identifier in LP_TOKEN_CONFIG.keys():
            print(f"\nProcessing LP token: {token_identifier}")
            state = await self.fetch_and_save_lp_state(token_identifier)
            if state:
                states[f"lp_{token_identifier}"] = state
        
        # Fetch farm states
        for token_identifier in FARM_TOKENS.keys():
            print(f"\nProcessing farm token: {token_identifier}")
            state = await self.fetch_and_save_farm_state(token_identifier)
            if state:
                states[f"farm_{token_identifier}"] = state
        
        # Fetch metastaking farm states
        for token_identifier in METASTAKING_FARMS.keys():
            print(f"\nProcessing metastaking farm token: {token_identifier}")
            state = await self.fetch_and_save_metastaking_farm_state(token_identifier)
            if state:
                states[f"metastaking_farm_{token_identifier}"] = state
        
        print(f"\nTotal states fetched: {len(states)}")
        return states

# Global state fetcher instance
state_fetcher = StateFetcher()

# Convenience functions for backward compatibility
async def get_pool_data(token_identifier: str) -> Optional[Dict[str, Any]]:
    return await state_fetcher.get_lp_pool_state(token_identifier)

async def get_farm_data(token_identifier: str) -> Optional[Dict[str, Any]]:
    return await state_fetcher.get_farm_state(token_identifier)

async def state_fetch():
    
    result = await state_fetcher.fetch_all_states()
    print(f"Test completed. Fetched {len(result)} states.")
    return result

if __name__ == "__main__":
    asyncio.run(state_fetch())