common.py

#!/usr/bin/env python3
"""Common functions."""

import asyncio
import csv
import json
import os
import shutil
import tempfile
import typing
from collections.abc import Iterable
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import date, datetime
from enum import StrEnum
from functools import reduce, wraps
from pathlib import Path
from typing import Any, ClassVar, Final, Generator, Optional, TypeVar

import duckdb
import pandas as pd
import schwab
import walrus
from authlib.integrations.starlette_client import OAuth, StarletteOAuth2App
from browser_use_sdk import AsyncBrowserUse, TaskResult
from loguru import logger
from playwright.sync_api import sync_playwright
from pydantic import BaseModel
from schwab.client import AsyncClient, Client
from schwab.orders.options import OptionSymbol

CODE_DIR = f"{Path.home()}/code/accounts"
PUBLIC_HTML = f"{CODE_DIR}/web/"
PREFIX = PUBLIC_HTML
HOURLY_LOGFILE = f"{PREFIX}/finance_hourly.log"
LOCK_TTL_SECONDS = 10 * 60
DUCKDB = f"{PREFIX}/db.duckdb"
DUCKDB_LOCK_NAME = "duckdb"
SCHWAB_LOCK_NAME = "schwab"
SCRIPT_LOCK_NAME = "script"
LEDGER_BIN = "ledger"
LEDGER_DIR = f"{Path.home()}/code/ledger"
LEDGER_DAT = f"{LEDGER_DIR}/ledger.ledger"
LEDGER_PRICES_DB = f"{LEDGER_DIR}/prices.db"
LEDGER_PREFIX = f"{LEDGER_BIN} -f {LEDGER_DAT} --price-db {LEDGER_PRICES_DB} -X '$' -c --no-revalued"
GET_TICKER_TIMEOUT = 30
PLOTLY_THEME = "plotly_dark"
# Include currency equivalents like money markets.
CURRENCIES_REGEX = r"^(\\$|CHF|EUR|GBP|SGD|SWVXX|SGOV|VBIL)$"
CASH_EQUIVALENTS = r"^(SWVXX|SGOV|VBIL)$"
LEDGER_CURRENCIES_CMD = f"{LEDGER_PREFIX} --limit 'commodity=~/{CURRENCIES_REGEX}/'"
SUBPLOT_MARGIN = {"l": 0, "r": 50, "b": 0, "t": 50}
SCHWAB_PAL_INTEREST_SPREAD = 2.8


class Brokerage(StrEnum):
    SCHWAB = "Charles Schwab Brokerage"
    SCHWAB_PAL = "Charles Schwab PAL Brokerage"
    IBKR = "Interactive Brokers"


OPTIONS_BROKERAGES = (Brokerage.IBKR, Brokerage.SCHWAB)


@dataclass
class FutureSpec:
    multiplier: float
    margin_requirement_percent: dict[Brokerage, float]


# Get margin requirements from Schwab or IBKR
FUTURE_SPEC: dict[str, FutureSpec] = {
    "10Y": FutureSpec(
        multiplier=1000,
        margin_requirement_percent={Brokerage.SCHWAB: 9, Brokerage.IBKR: 15},
    ),
    "M2K": FutureSpec(
        multiplier=5,
        margin_requirement_percent={Brokerage.SCHWAB: 9, Brokerage.IBKR: 9},
    ),
    "MBT": FutureSpec(
        multiplier=0.1,
        margin_requirement_percent={Brokerage.SCHWAB: 40, Brokerage.IBKR: 40},
    ),
    "MES": FutureSpec(
        multiplier=5,
        margin_requirement_percent={Brokerage.SCHWAB: 7, Brokerage.IBKR: 7},
    ),
    "MFS": FutureSpec(multiplier=50, margin_requirement_percent={Brokerage.IBKR: 6}),
    "MGC": FutureSpec(
        multiplier=10,
        margin_requirement_percent={Brokerage.SCHWAB: 16, Brokerage.IBKR: 11},
    ),
    "MTN": FutureSpec(multiplier=100, margin_requirement_percent={Brokerage.IBKR: 5}),
    # Silver 2500oz
    "QI": FutureSpec(
        multiplier=2500,
        margin_requirement_percent={Brokerage.SCHWAB: 30, Brokerage.IBKR: 21},
    ),
    # Silver 1000oz
    "SIL": FutureSpec(
        multiplier=1000,
        margin_requirement_percent={Brokerage.SCHWAB: 30},
    ),
    # Silver 5000oz
    "SI": FutureSpec(
        multiplier=5000,
        margin_requirement_percent={Brokerage.SCHWAB: 30, Brokerage.IBKR: 21},
    ),
    "TN": FutureSpec(multiplier=1000, margin_requirement_percent={Brokerage.IBKR: 4}),
    "ZN": FutureSpec(
        multiplier=1000,
        margin_requirement_percent={Brokerage.SCHWAB: 2, Brokerage.IBKR: 2},
    ),
}


def get_future_spec(ticker: str) -> FutureSpec:
    # Ticker ends with M26: /TNM26
    ticker = ticker[1:-3]
    return FUTURE_SPEC[ticker]


class GetTickerError(Exception):
    """Error getting ticker."""


class WalrusDb:
    def __init__(self):
        self.db = walrus.Database(host=os.environ.get("REDIS_HOST", "localhost"))
        self.cache = walrus.Cache(self.db, default_timeout=5 * 60)
        self.duckdb_lock = self.db.lock(DUCKDB_LOCK_NAME, ttl=LOCK_TTL_SECONDS * 1000)
        self.schwab_lock = self.db.lock(SCHWAB_LOCK_NAME, ttl=LOCK_TTL_SECONDS * 1000)


walrus_db: Final[WalrusDb] = WalrusDb()


class TickerOption(typing.NamedTuple):
    ticker: str
    expiration: date
    contract_type: str
    strike: float


class FutureQuote(typing.NamedTuple):
    mark: float
    multiplier: float


class OptionQuote(typing.NamedTuple):
    mark: float
    delta: float
    underlying_price: float
    symbol: str


class Schwab:
    SCHWAB_TOKEN_FILE: ClassVar[str] = f"{CODE_DIR}/.schwab_token.json"
    DUCKDB_OPTION_HISTORY: str = "option_prices"
    # Used for both stocks and futures
    DUCKDB_TICKER_HISTORY: str = "ticker_prices"

    def __init__(self):
        self.api_key: str = os.environ.get("SCHWAB_API_KEY", "")
        self.secret: str = os.environ.get("SCHWAB_SECRET", "")
        self._oauth: Optional[StarletteOAuth2App] = None

    def client(self, asyncio=False) -> Client | AsyncClient:
        if not all([self.api_key, self.secret]):
            raise GetTickerError("No schwab environment variables found")
        return schwab.auth.client_from_token_file(
            self.SCHWAB_TOKEN_FILE, self.api_key, self.secret, asyncio=asyncio
        )

    @property
    def oauth(self) -> StarletteOAuth2App:
        if not self._oauth:
            if not all([self.api_key, self.secret]):
                raise GetTickerError("No schwab environment variables found")
            self._oauth = OAuth().register(
                name="schwab",
                client_id=self.api_key,
                client_secret=self.secret,
                access_token_url="https://api.schwabapi.com/v1/oauth/token",
                authorize_url="https://api.schwabapi.com/v1/oauth/authorize",
                client_kwargs={
                    "scope": "read",
                },
            )
            if self._oauth is None:
                raise GetTickerError("Cannot create oauth")
        return self._oauth

    def write_token(self, token: dict):
        logger.info("Writing token")
        try:
            with open(self.SCHWAB_TOKEN_FILE) as f:
                data = json.load(f)
        except FileNotFoundError:
            data = {}
        data["creation_timestamp"] = int(datetime.now().timestamp())
        data["token"] = token
        with open(self.SCHWAB_TOKEN_FILE, "w") as f:
            f.write(json.dumps(data))

    def get_quotes(self, ts: Iterable[str]) -> dict[str, float]:
        r: dict[str, float] = {}
        j = self.client().get_quotes(ts).json()
        for t in ts:
            try:
                p = j[t]
            except KeyError:
                logger.error(f"Cannot find {t} in quote: {j}")
                raise GetTickerError(f"{t=} ticker not found")
            if "quote" in p:
                q = p["quote"]["lastPrice"]
            elif "regular" in p:
                q = p["regular"]["regularMarketLastPrice"]
            else:
                logger.error(p)
                raise GetTickerError(f"{t=} cannot find schwab price field")
            if q == 0:
                raise GetTickerError(f"{t=} received 0 as quote")
            logger.info(f"{t=} {q=}")
            r[t] = q
        if r:
            df = pd.DataFrame(r.items(), columns=["ticker", "price"])
            df["date"] = pd.Timestamp.now()
            df = df.set_index("date")
            to_sql(df, self.DUCKDB_TICKER_HISTORY)
        return r

    def get_delta_override(self, symbol: str) -> float:
        if (p := Path(f"{PUBLIC_HTML}delta_overrides")).exists():
            with p.open("r") as f:
                reader = csv.reader(f, delimiter=":")
                for row in reader:
                    if row[0] == symbol:
                        return float(row[1])
        return 0

    def get_option_quotes(
        self, ts: Iterable[TickerOption]
    ) -> dict[TickerOption, OptionQuote]:
        results: dict[TickerOption, OptionQuote] = {}
        fetch_tickers: dict[str, TickerOption] = {}
        for t in ts:
            symbol = OptionSymbol(
                t.ticker, t.expiration, t.contract_type[0], str(t.strike)
            ).build()
            fetch_tickers[symbol] = t
        if not fetch_tickers:
            return results
        js = self.client().get_quotes(fetch_tickers.keys()).json()
        for symbol, j in js.items():
            q = j["quote"]
            mark = q["mark"]
            delta = q["delta"]
            underlying_price = q["underlyingPrice"]
            if new_delta := self.get_delta_override(symbol):
                delta = new_delta
                logger.info(f"{symbol=} overriding {delta=}")
            logger.info(f"{symbol=} {mark=} {delta=} {underlying_price=}")
            if abs(delta) > 1:
                raise GetTickerError(f"Invalid delta value: {delta=} {symbol=}")
            contract_type = fetch_tickers[symbol].contract_type
            if contract_type == "PUT" and delta > 0:
                raise GetTickerError(f"Invalid PUT delta value: {delta=} {symbol=}")
            if contract_type == "CALL" and delta < 0:
                raise GetTickerError(f"Invalid CALL delta value: {delta=} {symbol=}")
            results[fetch_tickers[symbol]] = OptionQuote(
                mark=mark,
                delta=delta,
                underlying_price=underlying_price,
                symbol=symbol,
            )
        if results:
            df = pd.DataFrame(
                [
                    (q.symbol, q.mark, q.delta, q.underlying_price)
                    for q in results.values()
                ],
                columns=["symbol", "mark", "delta", "underlying_price"],
            )
            df["date"] = pd.Timestamp.now()
            df = df.set_index("date")
            to_sql(df, self.DUCKDB_OPTION_HISTORY)
        return results

    def get_future_quotes(self, ts: Iterable[str]) -> dict[str, FutureQuote]:
        r: dict[str, FutureQuote] = {}
        j = self.client().get_quotes(ts).json()
        for t, p in j.items():
            mark = p["quote"]["mark"]
            multiplier = p["reference"]["futureMultiplier"]
            q = FutureQuote(mark=mark, multiplier=multiplier)
            logger.info(f"{t=} {q=}")
            r[t] = q
        if r:
            df = pd.DataFrame(
                [(k, v.mark) for k, v in r.items()],
                columns=["ticker", "price"],
            )
            df["date"] = pd.Timestamp.now()
            df = df.set_index("date")
            to_sql(df, self.DUCKDB_TICKER_HISTORY)
        return r

    async def get_option_chains(
        self, tickers: list[str], from_date: date, to_date: date
    ) -> dict[str, pd.DataFrame]:
        async_client = self.client(asyncio=True)
        tg_results = {}
        async with asyncio.TaskGroup() as tg:
            for ticker in tickers:
                tg_results[ticker] = tg.create_task(
                    self.get_option_chain(async_client, ticker, from_date, to_date)
                )
        return {t: r.result() for t, r in tg_results.items()}

    async def get_option_chain(
        self, client: Client | AsyncClient, ticker: str, from_date: date, to_date: date
    ) -> pd.DataFrame:
        """Get option chain for a ticker between dates using async client.

        Returns DataFrame with columns:
        - expiration: date
        - strike: float
        - put_delta: float
        - call_delta: float
        - put_mark: float
        - call_mark: float
        - underlying_price: float
        """
        logger.info(f"Fetching {ticker=} from {from_date} to {to_date}")
        response = await client.get_option_chain(
            ticker,
            contract_type=schwab.client.Client.Options.ContractType.ALL,
            from_date=from_date,
            to_date=to_date,
        )
        data = response.json()

        rows = []
        underlying_price = data.get("underlyingPrice", 0)

        for exp_date_str, exp_data in data.get("callExpDateMap", {}).items():
            exp_date = datetime.strptime(exp_date_str.split(":")[0], "%Y-%m-%d").date()
            for strike_str, options in exp_data.items():
                strike = float(strike_str)
                for option in options:
                    oi = option.get("openInterest", 0)
                    call_delta = option.get("delta", 0)
                    call_mark = option.get("mark", 0)
                    rows.append(
                        {
                            "expiration": exp_date,
                            "strike": strike,
                            "call_delta": call_delta,
                            "call_mark": call_mark,
                            "underlying_price": underlying_price,
                            "call_oi": oi,
                        }
                    )

        for exp_date_str, exp_data in data.get("putExpDateMap", {}).items():
            exp_date = datetime.strptime(exp_date_str.split(":")[0], "%Y-%m-%d").date()
            for strike_str, options in exp_data.items():
                strike = float(strike_str)
                for option in options:
                    oi = option.get("openInterest", 0)
                    put_delta = option.get("delta", 0)
                    put_mark = option.get("mark", 0)
                    # Find matching row or create new
                    existing = [
                        r
                        for r in rows
                        if r["expiration"] == exp_date and r["strike"] == strike
                    ]
                    if existing:
                        existing[0]["put_delta"] = put_delta
                        existing[0]["put_mark"] = put_mark
                        existing[0]["put_oi"] = oi
                    else:
                        rows.append(
                            {
                                "expiration": exp_date,
                                "strike": strike,
                                "put_delta": put_delta,
                                "put_mark": put_mark,
                                "underlying_price": underlying_price,
                                "put_oi": oi,
                            }
                        )

        df = pd.DataFrame(rows)
        for col in (
            "put_delta",
            "call_delta",
            "put_mark",
            "call_mark",
            "put_oi",
            "call_oi",
        ):
            if col not in df.columns:
                df[col] = 0.0

        return df


def schwab_async_synchronized(method):
    @wraps(method)
    async def wrapper(*args, **kwargs):
        await asyncio.to_thread(walrus_db.schwab_lock.acquire)
        try:
            return await method(*args, **kwargs)
        finally:
            await asyncio.to_thread(walrus_db.schwab_lock.release)

    return wrapper


@contextmanager
def pandas_options():
    """Set pandas output options."""
    with pd.option_context(
        "display.max_rows", None, "display.max_columns", None, "display.width", 1000
    ):
        yield


def cache_tickers():
    from etfs import get_tickers as etfs_get_tickers
    from forex import TICKERS as FOREX_TICKERS
    from stock_options import get_options_and_spreads

    funcs = (
        lambda: get_tickers(etfs_get_tickers() | set(FOREX_TICKERS)),
        get_options_and_spreads,
    )
    with ThreadPoolExecutor() as executor:
        r = [executor.submit(f) for f in funcs]
    for f in r:
        f.result()


@walrus_db.schwab_lock
def get_tickers(ts: Iterable[str]) -> dict[str, float]:
    prefix = "get_ticker:"

    # Map each input ticker to its canonical fetch key and inverse flag
    ticker_map: dict[str, tuple[str, bool]] = {}
    for t in ts:
        if t.endswith("USD"):
            ticker_map[t] = (f"USD/{t[:3]}", True)
        elif t.startswith("SPX"):
            ticker_map[t] = ("$SPX", False)
        elif t.startswith("XSP"):
            ticker_map[t] = ("$XSP", False)
        else:
            ticker_map[t] = (t, False)

    fetch_keys = {k for k, _ in ticker_map.values()}

    # Check cache
    r: dict[str, float] = {}
    for key, val in walrus_db.cache.get_many(
        [f"{prefix}{k}" for k in fetch_keys]
    ).items():
        r[key.removeprefix(prefix)] = val

    # Fetch missing from Schwab
    remaining = fetch_keys - set(r)
    if remaining:
        qs = Schwab().get_quotes(remaining)
        walrus_db.cache.set_many({f"{prefix}{t}": v for t, v in qs.items()})
        r.update(qs)

    # Build output for original tickers
    return {
        t: (1 / r[key]) if invert else r[key] for t, (key, invert) in ticker_map.items()
    }


def make_option_key(t: TickerOption) -> str:
    return f"{t.ticker} {t.expiration} {t.strike} {t.contract_type}"


@schwab_async_synchronized
async def get_option_chains(
    tickers: list[str], from_date: date, to_date: date
) -> dict[str, pd.DataFrame]:
    """Get option chains for multiple tickers concurrently using async interface.

    Args:
        tickers: List of ticker symbols
        from_date: Start date for option expiration
        to_date: End date for option expiration

    Returns:
        Dictionary mapping ticker symbols to their option chain DataFrames.
    """
    prefix = "get_option_chain:"
    results: dict[str, pd.DataFrame] = {}
    needed_tickers: list[str] = []

    # Check cache for each ticker
    cache_keys = [f"{prefix}{t}:{from_date}:{to_date}" for t in tickers]
    cached = walrus_db.cache.get_many(cache_keys)

    for ticker in tickers:
        cache_key = f"{prefix}{ticker}:{from_date}:{to_date}"
        if cache_key in cached:
            df = cached[cache_key]
            results[ticker] = df
        else:
            needed_tickers.append(ticker)

    if not needed_tickers:
        return results

    # Fetch missing tickers concurrently
    results.update(await Schwab().get_option_chains(needed_tickers, from_date, to_date))
    cache_qs = {f"{prefix}{t}:{from_date}:{to_date}": results[t] for t in results}
    walrus_db.cache.set_many(cache_qs)
    return results


@walrus_db.schwab_lock
def get_option_quotes(
    ts: Iterable[TickerOption],
) -> dict[TickerOption, Optional[OptionQuote]]:
    r: dict[TickerOption, Optional[OptionQuote]] = {}
    prefix = "get_option_quote:"
    cache_keys: list[str] = [f"{prefix}{make_option_key(t)}" for t in ts]
    cached: dict[str, OptionQuote] = walrus_db.cache.get_many(cache_keys)
    needed: set[TickerOption] = set()
    for t in ts:
        if o := cached.get(f"{prefix}{make_option_key(t)}"):
            r[t] = o
        else:
            needed.add(t)
    if needed:
        qs = Schwab().get_option_quotes(needed)
        cache_qs = {f"{prefix}{make_option_key(t)}": qs[t] for t in qs}
        walrus_db.cache.set_many(cache_qs)
        r.update(qs)
    return r


def get_option_quotes_from_db(
    ts: Iterable[TickerOption],
) -> dict[TickerOption, Optional[OptionQuote]]:
    r: dict[TickerOption, Optional[OptionQuote]] = {}
    symbol_map: dict[str, TickerOption] = {}
    for t in ts:
        symbol = OptionSymbol(
            t.ticker, t.expiration, t.contract_type[0], str(t.strike)
        ).build()
        symbol_map[symbol] = t
    if not symbol_map:
        return r
    symbols_str = ",".join(f"'{s}'" for s in symbol_map)
    query = (
        f"SELECT symbol, mark, delta, underlying_price "
        f"FROM {Schwab.DUCKDB_OPTION_HISTORY} "
        f"WHERE symbol IN ({symbols_str}) "
        f"QUALIFY date = MAX(date) OVER (PARTITION BY symbol)"
    )
    with duckdb_lock(read_only=True) as con:
        df = con.sql(query).df()
    for _, row in df.iterrows():
        t = symbol_map[row["symbol"]]
        r[t] = OptionQuote(
            mark=row["mark"],
            delta=row["delta"],
            underlying_price=row["underlying_price"],
            symbol=row["symbol"],
        )
    for symbol, t in symbol_map.items():
        if t not in r:
            r[t] = None
    return r


@walrus_db.schwab_lock
def get_future_quotes(ts: Iterable[str]) -> dict[str, FutureQuote]:
    prefix = "get_future_quote:"

    # Map each input ticker to its canonical fetch key
    ticker_map: dict[str, str] = {}
    for t in ts:
        if t.startswith("/MTN"):
            ticker_map[t] = t.replace("/MTN", "/TN")
        else:
            ticker_map[t] = t

    fetch_keys = set(ticker_map.values())

    # Check cache
    r: dict[str, FutureQuote] = {}
    for key, val in walrus_db.cache.get_many(
        [f"{prefix}{k}" for k in fetch_keys]
    ).items():
        r[key.removeprefix(prefix)] = val

    # Fetch missing from Schwab
    remaining = fetch_keys - set(r)
    if remaining:
        qs = Schwab().get_future_quotes(remaining)
        walrus_db.cache.set_many({f"{prefix}{t}": v for t, v in qs.items()})
        r.update(qs)

    # Build output for original tickers
    return {
        t: FutureQuote(mark=r[key].mark, multiplier=get_future_spec(t).multiplier)
        if t.startswith("/MTN")
        else r[key]
        for t, key in ticker_map.items()
    }


@contextmanager
def duckdb_lock(
    read_only: bool = False,
) -> Generator[duckdb.DuckDBPyConnection, None, None]:
    with walrus_db.duckdb_lock:
        with duckdb.connect(DUCKDB, read_only=read_only) as con:
            yield con


def compact_db():
    # Remove duplicates from large timeseries
    table_specs = {
        "option_prices": ("(mark, delta, underlying_price)", "symbol"),
        "ticker_prices": ("price", "ticker"),
    }
    with duckdb_lock() as con:
        con.execute("BEGIN TRANSACTION")
        for table, (cols, partition) in table_specs.items():
            con.execute(
                f"CREATE TABLE {table}_new AS "
                f"SELECT * FROM {table} "
                f"QUALIFY {cols} IS DISTINCT FROM "
                f"LAG({cols}) OVER (PARTITION BY {partition} ORDER BY date) "
                f"ORDER BY (date, {partition})"
            )
            con.execute(f"DROP TABLE {table}")
            con.execute(f"ALTER TABLE {table}_new RENAME TO {table}")
        con.execute("COMMIT")
    with walrus_db.duckdb_lock:
        with temporary_file_move(DUCKDB) as new_db:
            with duckdb.connect() as con:
                con.execute(f"ATTACH '{DUCKDB}' AS old")
                os.unlink(new_db.name)
                con.execute(f"ATTACH '{new_db.name}' AS new")
                con.execute("COPY FROM DATABASE old TO new")


def insert_sql(table: str, data: dict[str, Any], timestamp: Optional[datetime] = None):
    """Insert data into sql table."""
    cols = ["date"]
    values = [timestamp]
    if timestamp is None:
        values = [datetime.now()]
    prepared = ["?"]
    for col, value in data.items():
        cols.append(f'"{col}"')
        values.append(value)
        prepared.append("?")
    with duckdb_lock() as con:
        con.execute(
            f"INSERT INTO {table} ({', '.join(cols)}) VALUES ({', '.join(prepared)})",
            values,
        )


def read_sql_table(table, index_col="date") -> pd.DataFrame:
    with duckdb_lock(read_only=True) as con:
        rel = con.table(table)
        return rel.df().set_index(index_col)


def read_sql_query(query: str, index: str = "date") -> pd.DataFrame:
    """Load table from sql query."""
    with duckdb_lock(read_only=True) as con:
        return con.sql(query).df().set_index(index)


def read_sql_last(table: str) -> pd.DataFrame:
    return read_sql_query(f"SELECT * FROM {table} ORDER BY date DESC LIMIT 1")


def to_sql(dataframe, table, if_exists="append"):
    """Write dataframe to sql table."""
    dataframe = dataframe.reset_index()
    with duckdb_lock() as con:
        if if_exists == "replace":
            con.execute("BEGIN TRANSACTION")
            con.execute(f"DROP TABLE IF EXISTS {table}")
            con.execute(f"CREATE TABLE {table} AS SELECT * FROM dataframe")
            con.execute("COMMIT")
        else:
            exists = (
                con.execute(
                    f"SELECT count(*) FROM information_schema.tables WHERE table_name = '{table}'"
                ).fetchone()
                or (0,)
            )[0]
            if exists:
                sql = f"INSERT INTO {table} BY NAME SELECT * FROM dataframe"
                con.execute(sql)
            else:
                con.execute(f"CREATE TABLE {table} AS SELECT * FROM dataframe")


@contextmanager
def temporary_file_move(dest_file):
    """Provides a temporary file that is moved in place after context."""
    with tempfile.NamedTemporaryFile(mode="w", delete=False) as write_file:
        yield write_file
    shutil.move(write_file.name, dest_file)


BaseModelType = TypeVar("BaseModelType", bound=BaseModel)


async def run_browser_use(
    task: str, model: typing.Type[BaseModelType]
) -> TaskResult[BaseModelType]:
    rate_limit = walrus_db.db.rate_limit("BrowserUse", limit=1, per=24 * 60 * 60)
    if rate_limit.limit(task):
        raise walrus.RateLimitException("Browser Use rate limited to once per day")
    client = AsyncBrowserUse()
    return await client.run(task=task, output_schema=model)


@contextmanager
def run_with_browser_page(url):
    """Run code with a Chromium browser page."""
    with sync_playwright() as p:
        browser = p.chromium.launch()
        page = browser.new_page()
        try:
            page.goto(url)
            yield page
        finally:
            browser.close()


def reduce_merge_asof(dataframes):
    """Reduce and merge date tables."""
    return reduce(
        lambda L, r: pd.merge_asof(L, r, left_index=True, right_index=True),
        dataframes,
    )


def load_sql_and_rename_col(table, rename_cols=None):
    """Load resampled table from sql and rename columns."""
    dataframe = read_sql_table(table)
    if rename_cols:
        dataframe = dataframe.rename(columns=rename_cols)
    return dataframe