lambda_function.py

import datetime
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from threading import Lock

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

if os.environ.get("ENVIRONMENT") != "production":
    from dotenv import load_dotenv

    load_dotenv()

GLOBAL_QUOTE_LITERAL = "Global Quote"
PRICE_LITERAL = "05. price"
INFO_LITERAL = "Information"
DEFAULT_TIMEOUT_SECONDS = 10
COINGECKO_SYMBOL_CACHE = None
DEFAULT_NOTION_UPDATE_MAX_WORKERS = 4
DEFAULT_NOTION_UPDATE_RPS_LIMIT = 2.5
DEFAULT_NOTION_UPDATE_BURST = 1


def create_session():
    retry = Retry(
        total=3,
        backoff_factor=0.5,
        status_forcelist=[429, 500, 502, 503, 504],
        allowed_methods=["GET", "POST", "PATCH"],
    )
    session = requests.Session()
    adapter = HTTPAdapter(max_retries=retry)
    session.mount("https://", adapter)
    session.mount("http://", adapter)
    return session


def request_json(session, method, url, headers=None, payload=None, params=None):
    try:
        response = session.request(
            method,
            url,
            headers=headers,
            json=payload,
            params=params,
            timeout=DEFAULT_TIMEOUT_SECONDS,
        )
    except requests.RequestException as exc:
        print(f"Request failed for {url}: {exc}")
        return None

    if response.status_code >= 400:
        print(f"Request failed for {url}: {response.status_code} {response.text}")
        return None

    if response.content:
        return response.json()
    return None


def request_status(session, method, url, headers=None, payload=None):
    try:
        response = session.request(
            method,
            url,
            headers=headers,
            json=payload,
            timeout=DEFAULT_TIMEOUT_SECONDS,
        )
    except requests.RequestException as exc:
        print(f"Request failed for {url}: {exc}")
        return False

    if response.status_code >= 400:
        print(f"Request failed for {url}: {response.status_code} {response.text}")
        return False
    return True


def get_required_env(name):
    value = os.environ.get(name)
    if not value:
        raise ValueError(f"Missing required environment variable: {name}")
    return value


def fetch_crypto_prices(unique_coins, session):
    if not unique_coins:
        return {}
    coins_list_url = "https://api.coingecko.com/api/v3/coins/list"
    print("Retrieving coins list from CoinGecko")
    coins_list = get_cached_coins_list(session, coins_list_url)
    if not coins_list:
        return {}
    print("Coins list retrieved successfully")

    symbol_to_id = create_symbol_to_id_mapping(coins_list, unique_coins)

    coin_prices = {}
    coin_ids = [
        symbol_to_id.get(coin.lower())
        for coin in unique_coins
        if symbol_to_id.get(coin.lower())
    ]
    if coin_ids:
        crypto_price_url = f"https://api.coingecko.com/api/v3/simple/price?ids={','.join(coin_ids)}&vs_currencies=usd"
        print("Retrieving prices from CoinGecko")
        price_data = request_json(session, "GET", crypto_price_url)
        if not price_data:
            print("Failed to retrieve prices from CoinGecko")
            return {}
        print("Prices retrieved successfully from CoinGecko")
        for coin in unique_coins:
            coin_id = symbol_to_id.get(coin.lower())
            if coin_id and coin_id.lower() in price_data:
                usd_price = price_data[coin_id.lower()].get("usd")
                if usd_price is not None:
                    coin_prices[coin] = usd_price
                else:
                    print(f"Warning: USD price not available for {coin}")

    return coin_prices


def get_cached_coins_list(session, coins_list_url):
    global COINGECKO_SYMBOL_CACHE
    if COINGECKO_SYMBOL_CACHE is None:
        COINGECKO_SYMBOL_CACHE = request_json(session, "GET", coins_list_url)
    return COINGECKO_SYMBOL_CACHE


def create_symbol_to_id_mapping(coins_list, unique_coins):
    symbol_to_id = {}
    for coin in coins_list:
        symbol = coin["symbol"]
        coin_id = coin["id"]
        if symbol.upper() in unique_coins and not (
            (symbol == "dai" and coin_id != "dai")
            or (symbol == "mana" and coin_id != "decentraland")
            or (symbol == "eth" and coin_id != "ethereum")
            or (symbol == "btc" and coin_id != "bitcoin")
            or (symbol == "usdt" and coin_id != "tether")
            or (symbol == "bnb" and coin_id != "binancecoin")
        ):
            symbol_to_id[symbol] = coin_id
    return symbol_to_id


def fetch_stock_prices(unique_stocks, alpha_vantage_api_key, session):
    stock_prices = {"USD": 1.00}  # Initialize with a value for USD
    for stock_symbol in unique_stocks:
        if stock_symbol != "USD":
            price = get_stock_price(stock_symbol, alpha_vantage_api_key, session)
            if price is not None:
                print(
                    "Successfully fetched price for stock "
                    + stock_symbol
                    + ". New price: "
                    + str(price)
                )
                stock_prices[stock_symbol] = price
            time.sleep(1)  # Rate limit: 1 request per second for AlphaVantage free tier
    return stock_prices


def get_stock_price(stock_symbol, alpha_vantage_api_key, session):
    if stock_symbol == "CSPX":
        stock_symbol = "CSPX.LON"  # Adjust the symbol for CSPX
    alpha_vantage_url = build_url(stock_symbol, alpha_vantage_api_key)
    print("Fetching stock price for " + stock_symbol)
    data = request_json(session, "GET", alpha_vantage_url)
    return parse_data(data)


def build_url(stock_symbol, alpha_vantage_api_key):
    return f"https://www.alphavantage.co/query?apikey={alpha_vantage_api_key}&function=GLOBAL_QUOTE&symbol={stock_symbol}"


def parse_data(data):
    if not data:
        return None
    if GLOBAL_QUOTE_LITERAL in data:
        global_quote = data[GLOBAL_QUOTE_LITERAL]
        price = global_quote.get(PRICE_LITERAL)
        try:
            return float(price)
        except (TypeError, ValueError):
            print(f"Unexpected price value: {price}")
            return None
    elif INFO_LITERAL in data:
        print(
            f"Rate limit exceeded or other information received: {data[INFO_LITERAL]}"
        )
    return None


def get_stock_amount(result):
    amount = result.get("properties", {}).get("Amount", {}).get("number")
    return amount if isinstance(amount, (int, float)) else 0


def get_select_name(result, property_name):
    properties = result.get("properties", {})
    select = properties.get(property_name, {}).get("select")
    if not isinstance(select, dict):
        return None
    name = select.get("name")
    return name if isinstance(name, str) and name else None


def parse_int_env(name, default, minimum):
    value = os.environ.get(name)
    if value is None:
        return default
    try:
        parsed = int(value)
    except ValueError:
        print(f"Invalid value for {name}: {value}. Using default {default}.")
        return default
    if parsed < minimum:
        print(
            f"Value for {name} below minimum {minimum}: {parsed}. Using default {default}."
        )
        return default
    return parsed


def parse_float_env(name, default, minimum):
    value = os.environ.get(name)
    if value is None:
        return default
    try:
        parsed = float(value)
    except ValueError:
        print(f"Invalid value for {name}: {value}. Using default {default}.")
        return default
    if parsed < minimum:
        print(
            f"Value for {name} below minimum {minimum}: {parsed}. Using default {default}."
        )
        return default
    return parsed


class RateLimiter:
    def __init__(self, rps_limit, burst):
        self.interval = 1.0 / rps_limit if rps_limit > 0 else 0
        self.burst = max(1, burst)
        self._lock = Lock()
        self._next_allowed_at = time.monotonic()

    def wait_for_slot(self):
        if self.interval <= 0:
            return
        with self._lock:
            now = time.monotonic()
            floor = now - ((self.burst - 1) * self.interval)
            if self._next_allowed_at < floor:
                self._next_allowed_at = floor
            delay = self._next_allowed_at - now
            if delay > 0:
                time.sleep(delay)
                now = time.monotonic()
            self._next_allowed_at = max(self._next_allowed_at, now) + self.interval


def build_update_jobs(type, database_results, prices):
    jobs = []
    for result in database_results:
        page_id = result["id"]
        notion_page_url = f"https://api.notion.com/v1/pages/{page_id}"
        if type == "crypto":
            symbol = get_select_name(result, "Coin")
            if not symbol:
                print(
                    "Warning: Skipping crypto entry with missing Coin select. Page id: "
                    + page_id
                )
                continue
        else:
            symbol = result["properties"]["Stock"]["select"]["name"]
        current_price = result["properties"]["Price"]["number"]
        new_price = prices.get(symbol, 0)
        resolved_price = new_price if new_price is not None else current_price
        update_payload = {
            "properties": {
                "Price": {
                    "number": float(resolved_price)
                    if resolved_price is not None
                    else current_price
                }
            }
        }
        jobs.append(
            {
                "page_id": page_id,
                "symbol": symbol,
                "url": notion_page_url,
                "payload": update_payload,
            }
        )
    return jobs


def rate_limited_request_status(limiter, session, method, url, headers=None, payload=None):
    limiter.wait_for_slot()
    start = time.monotonic()
    success = request_status(
        session, method, url, headers=headers, payload=payload
    )
    return success, round(time.monotonic() - start, 3)


def run_notion_updates_concurrently(jobs, headers, session, max_workers, limiter):
    outcomes = {"ok": 0, "fail": 0}

    def worker(job):
        print("Updating price in Notion for " + job["symbol"])
        ok, elapsed = rate_limited_request_status(
            limiter,
            session,
            "PATCH",
            job["url"],
            headers=headers,
            payload=job["payload"],
        )
        status = "ok" if ok else "fail"
        print(
            f"Notion update outcome for {job['symbol']}: {status} "
            f"(page_id={job['page_id']}, elapsed={elapsed}s)"
        )
        return ok

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [executor.submit(worker, job) for job in jobs]
        for future in as_completed(futures):
            if future.result():
                outcomes["ok"] += 1
            else:
                outcomes["fail"] += 1
    return outcomes


def update_notion_prices(type, database_results, prices, headers, session):
    jobs = build_update_jobs(type, database_results, prices)
    if not jobs:
        print("No Notion updates to apply")
        return

    max_workers = parse_int_env(
        "NOTION_UPDATE_MAX_WORKERS", DEFAULT_NOTION_UPDATE_MAX_WORKERS, 1
    )
    rps_limit = parse_float_env(
        "NOTION_UPDATE_RPS_LIMIT", DEFAULT_NOTION_UPDATE_RPS_LIMIT, 0.1
    )
    burst = parse_int_env("NOTION_UPDATE_BURST", DEFAULT_NOTION_UPDATE_BURST, 1)
    limiter = RateLimiter(rps_limit, burst)

    print(
        f"Starting Notion updates: count={len(jobs)}, workers={max_workers}, "
        f"rps_limit={rps_limit}, burst={burst}"
    )
    outcomes = run_notion_updates_concurrently(
        jobs, headers, session, max_workers, limiter
    )
    print(
        f"Completed Notion updates: ok={outcomes['ok']}, fail={outcomes['fail']}"
    )


def query_notion_database(database_id, headers, session):
    notion_db_url = f"https://api.notion.com/v1/databases/{database_id}/query"
    results = []
    payload = {}
    while True:
        database = request_json(
            session, "POST", notion_db_url, headers=headers, payload=payload
        )
        if not database:
            break
        results.extend(database.get("results", []))
        if database.get("has_more"):
            payload = {"start_cursor": database.get("next_cursor")}
        else:
            break
    return results


def calculate_total_assets(databases, headers, session):
    print("Calculating total assets")
    total = 0
    for database_id in databases:
        results = query_notion_database(database_id, headers, session)
        for result in results:
            if result["parent"]["database_id"] == os.environ["FIAT_DB_ID"]:
                total += result["properties"]["Total"]["number"]
            else:
                total += result["properties"]["Total"]["formula"]["number"]
    return round(total, 2)


def update_total_assets_callout(block_id, total_assets, headers, session):
    notion_block_url = f"https://api.notion.com/v1/blocks/{block_id}"
    block = request_json(session, "GET", notion_block_url, headers=headers)
    if not block:
        return

    callout = block.get("callout")
    if not callout:
        print("Callout block not found in response")
        return

    rich_text = callout.get("rich_text", [])
    if len(rich_text) < 1:
        rich_text.append({"type": "text", "text": {"content": ""}})
    if len(rich_text) < 2:
        rich_text.append({"type": "text", "text": {"content": ""}})

    if "text" not in rich_text[1]:
        rich_text[1]["text"] = {"content": ""}
    rich_text[1]["text"]["content"] = f": ${total_assets:.2f}"
    callout["rich_text"] = rich_text
    print("Updating total assets")
    request_status(
        session,
        "PATCH",
        notion_block_url,
        headers=headers,
        payload={"callout": {"rich_text": callout["rich_text"]}},
    )


def lambda_handler(event, context):
    notion_api_key = get_required_env("NOTION_API_KEY")
    alpha_vantage_api_key = get_required_env("ALPHA_VANTAGE_API_KEY")
    crypto_database_id = get_required_env("CRYPTO_DB_ID")
    stock_database_id = get_required_env("STOCK_DB_ID")
    fiat_database_id = get_required_env("FIAT_DB_ID")
    session = create_session()

    headers = {
        "Authorization": "Bearer " + notion_api_key,
        "accept": "application/json",
        "Notion-Version": "2022-06-28",
        "content-type": "application/json",
    }

    # CRYPTOCURRENCY PRICES
    print("Getting CRYPTO database information")
    crypto_results = query_notion_database(crypto_database_id, headers, session)
    unique_coins_set = set()
    for result in crypto_results:
        coin_name = get_select_name(result, "Coin")
        if coin_name:
            unique_coins_set.add(coin_name)
        else:
            page_id = result.get("id", "unknown")
            print(
                "Warning: Skipping crypto entry with missing Coin select. Page id: "
                + page_id
            )
    unique_coins = list(unique_coins_set)
    crypto_prices = fetch_crypto_prices(unique_coins, session)
    if crypto_prices:
        update_notion_prices("crypto", crypto_results, crypto_prices, headers, session)

    # STOCK PRICES
    # Get current UTC hour
    current_utc_hour = datetime.datetime.utcnow().hour

    # Define the hour at which to perform the stock price update (e.g., 11 UTC)
    # This is because the free tier of the API has a limit of 25 requests per day
    stock_update_hour = 11

    # Only execute stock price update if the current hour matches the specified hour
    if current_utc_hour == stock_update_hour:
        stock_results = query_notion_database(stock_database_id, headers, session)
        filtered_stock_results = [
            result
            for result in stock_results
            if get_stock_amount(result) > 0
            and result["properties"]["Stock"]["select"]["name"] != "USD"
        ]
        filtered_out_count = len(stock_results) - len(filtered_stock_results)
        if filtered_out_count:
            print(
                f"Skipping {filtered_out_count} stock entries with Amount <= 0 or USD currency"
            )
        unique_stocks = list(
            set(
                [
                    result["properties"]["Stock"]["select"]["name"]
                    for result in filtered_stock_results
                ]
            )
        )
        stock_prices = fetch_stock_prices(unique_stocks, alpha_vantage_api_key, session)
        if stock_prices:
            update_notion_prices(
                "stock", filtered_stock_results, stock_prices, headers, session
            )
    else:
        print("Skipping stock price updates")

    # CALCULATE TOTAL ASSETS
    block_id = get_required_env("TOTAL_CALLOUT_BLOCK_ID")
    total_assets = calculate_total_assets(
        [crypto_database_id, stock_database_id, fiat_database_id], headers, session
    )
    update_total_assets_callout(block_id, total_assets, headers, session)


# Main execution
if __name__ == "__main__":
    lambda_handler(None, None)