from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
from typing import Any
import requests
import pytz
import os
import yaml
from tools.final_answer import FinalAnswerTool
from moralis import evm_api
from Gradio_UI import GradioUI


VALID_NETWORKS = ["eth", "base", "bsc", "gnosis"]
MAX_TOKENS_CONTEXT = 2096


# Below is an example of a tool that does nothing. Amaze us with your creativity !
def process_wallet_activity(activity_data):
    """Process and summarize wallet activity data to reduce token size
    Args:
        activity_data: Raw activity data from Moralis API
    Returns:
        dict: Summarized activity data
    """
    summary = {
        'total_transactions': len(activity_data['result']),
        'tokens': {},
        'date_range': {
            'from': None,
            'to': None
        }
    }
    
    for tx in activity_data['result']:
        token_symbol = tx['token_symbol']
        if token_symbol not in summary['tokens']:
            summary['tokens'][token_symbol] = {
                'count': 0,
                'total_value': 0,
                'token_name': tx['token_name'],
                'token_decimals': tx['token_decimals'],
                'verified': tx['verified_contract'],
                'security_score': tx['security_score']
            }
        
        summary['tokens'][token_symbol]['count'] += 1
        summary['tokens'][token_symbol]['total_value'] += float(tx['value_decimal'] or 0)
        
        tx_date = tx['block_timestamp']
        if not summary['date_range']['from'] or tx_date < summary['date_range']['from']:
            summary['date_range']['from'] = tx_date
        if not summary['date_range']['to'] or tx_date > summary['date_range']['to']:
            summary['date_range']['to'] = tx_date
    
    return summary

@tool
def get_wallet_activity(
    address: str, network: str = "eth"
) -> dict[str, Any]:  
    """A tool that gets the tokens transactions of ERC20 of a wallet in a specific blockchain network and returns detailed information.
    The returned data includes:
    - total_transactions: Total number of transactions
    - tokens: Dictionary with token details including:
        - count: Number of transactions for this token
        - total_value: Total value transferred
        - token_name: Full name of the token
        - token_decimals: Number of decimals for the token
        - verified: Whether the token contract is verified
        - security_score: Security score of the token if available
    - date_range: Time range of the transactions
    
    Args:
        address: the address of the wallet
        network: the blockchain network to get the activity from (default: eth) valid options: "eth", "base", "bsc", "gnosis"
    Returns:
        dict: Detailed wallet activity information that can be analyzed to understand the wallet's behavior
    """
    params = {
        "address": address,
        "chain": network,
        "order": "DESC",
        "limit": 50 # free models have a low context window
    }

    api_key = os.getenv("MORALIS_API_KEY")
    result = evm_api.token.get_wallet_token_transfers(
        api_key=api_key,
        params=params,
    )
    
    # Process and return the detailed data for the LLM to analyze
    return process_wallet_activity(result)


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'

# Get environment variables - works both locally and in Hugging Face Spaces
os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")


model = HfApiModel(
    max_tokens=MAX_TOKENS_CONTEXT,
    temperature=0.5,
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",  # it is possible that this model may be overloaded
    custom_role_conversions=None,
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", "r") as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[final_answer, get_wallet_activity, get_current_time_in_timezone],  ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates,
)


GradioUI(agent).launch()