"""
Utils for tidb
"""

from datetime import datetime, timedelta, timezone

import pymysql


def find_last_timestamp(connection: pymysql.Connection, table: str) -> float:
    """Find the last timestamp
    1. Try to find the last timestamp from `last_timestamp` table
    2. Try to find the last timestamp in the last 3 days
    3. If not found, return the timestamp of 30 minutes ago
    """
    with connection.cursor() as cursor:
        sql = f"""
        SELECT timestamp
        FROM `last_timestamp`
        WHERE table_name="{table}"
        """
        cursor.execute(sql)
        row = cursor.fetchone()
        if row is not None:
            return row["timestamp"]

        start_time = datetime.now(timezone.utc).replace(
            day=1, hour=0, minute=0, second=0
        )
        for _ in range(3):
            start_ts = start_time.timestamp()
            sql = f"""
                SELECT timestamp
                FROM `{table}`
                WHERE timestamp>{start_ts}
                ORDER BY timestamp DESC
                LIMIT 1
            """
            cursor.execute(sql)
            row = cursor.fetchone()
            if row is not None:
                return row["timestamp"]
            start_time = start_time - timedelta(days=1)
        return (datetime.now(timezone.utc) - timedelta(minutes=30)).timestamp()


def update_last_timestamp(
    connection: pymysql.Connection, table: str, timestamp: float
) -> None:
    """Update the last timestamp"""
    with connection.cursor() as cursor:
        sql = f"""
        REPLACE INTO `last_timestamp` (
            table_name, timestamp
        )
        VALUES ('{table}', {timestamp})
        """
        cursor.execute(sql)
        connection.commit()