database.py

"""
AddressTrackerBot v3 - SQLite database layer.
Thread-safe: Lock around writes, concurrent reads OK.
Uses thread-local connections for reuse.
"""

import atexit
import sqlite3
import time
import threading
from threading import Lock
from typing import Any, Dict, List, Optional, Tuple

from config import DB_PATH
from utils import logger

_write_lock = Lock()
_local = threading.local()

SCHEMA_VERSION = 1

# Rate limiting (in-memory, bounded)
_rate_limit_store: Dict[int, List[float]] = {}
_rate_limit_lock = Lock()
_MAX_RATE_LIMIT_USERS = 10000


def get_connection() -> sqlite3.Connection:
    """Return a thread-local reusable connection."""
    conn = getattr(_local, 'conn', None)
    if conn is None:
        conn = sqlite3.connect(DB_PATH, check_same_thread=False)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA foreign_keys=ON")
        _local.conn = conn
    return conn


def close_connection() -> None:
    """Close the thread-local connection if open."""
    conn = getattr(_local, 'conn', None)
    if conn is not None:
        try:
            conn.close()
        except Exception:
            pass
        _local.conn = None


atexit.register(close_connection)


def init_db() -> None:
    """Create tables if they don't exist and seed chain data."""
    with _write_lock:
        conn = get_connection()
        try:
            conn.executescript("""
                CREATE TABLE IF NOT EXISTS schema_version (
                    version INTEGER NOT NULL
                );

                CREATE TABLE IF NOT EXISTS users (
                    chat_id     INTEGER PRIMARY KEY,
                    username    TEXT NOT NULL,
                    is_admin    INTEGER NOT NULL DEFAULT 0,
                    created_at  TEXT NOT NULL DEFAULT (datetime('now'))
                );

                CREATE TABLE IF NOT EXISTS chains (
                    chain_id        INTEGER PRIMARY KEY,
                    name            TEXT NOT NULL,
                    native_symbol   TEXT NOT NULL,
                    explorer_url    TEXT NOT NULL,
                    icon            TEXT NOT NULL DEFAULT ''
                );

                CREATE TABLE IF NOT EXISTS addresses (
                    id              INTEGER PRIMARY KEY AUTOINCREMENT,
                    chat_id         INTEGER NOT NULL,
                    chain_id        INTEGER NOT NULL DEFAULT 1,
                    address         TEXT NOT NULL,
                    name            TEXT NOT NULL,
                    last_seen_block INTEGER NOT NULL DEFAULT 0,
                    threshold_wei   TEXT NOT NULL DEFAULT '0',
                    created_at      TEXT NOT NULL DEFAULT (datetime('now')),
                    FOREIGN KEY (chat_id) REFERENCES users(chat_id),
                    FOREIGN KEY (chain_id) REFERENCES chains(chain_id),
                    UNIQUE(chat_id, name)
                );

                CREATE TABLE IF NOT EXISTS transactions (
                    id              INTEGER PRIMARY KEY AUTOINCREMENT,
                    address_id      INTEGER NOT NULL,
                    chain_id        INTEGER NOT NULL,
                    tx_hash         TEXT NOT NULL,
                    block_number    INTEGER NOT NULL,
                    from_addr       TEXT NOT NULL,
                    to_addr         TEXT NOT NULL DEFAULT '',
                    value_wei       TEXT NOT NULL DEFAULT '0',
                    tx_type         TEXT NOT NULL DEFAULT 'native',
                    token_address   TEXT,
                    token_symbol    TEXT,
                    token_decimals  INTEGER,
                    direction       TEXT NOT NULL DEFAULT 'unknown',
                    timestamp       INTEGER,
                    func_name       TEXT,
                    created_at      TEXT NOT NULL DEFAULT (datetime('now')),
                    FOREIGN KEY (address_id) REFERENCES addresses(id),
                    FOREIGN KEY (chain_id) REFERENCES chains(chain_id),
                    UNIQUE(tx_hash, address_id, tx_type, token_address)
                );

                CREATE TABLE IF NOT EXISTS settings (
                    chat_id     INTEGER NOT NULL,
                    key         TEXT NOT NULL,
                    value       TEXT NOT NULL,
                    PRIMARY KEY (chat_id, key),
                    FOREIGN KEY (chat_id) REFERENCES users(chat_id)
                );

                CREATE INDEX IF NOT EXISTS idx_addresses_chat ON addresses(chat_id);
                CREATE INDEX IF NOT EXISTS idx_addresses_chain ON addresses(chain_id);
                CREATE INDEX IF NOT EXISTS idx_transactions_address ON transactions(address_id);
                CREATE INDEX IF NOT EXISTS idx_transactions_block ON transactions(block_number);
            """)

            # Seed schema version
            row = conn.execute("SELECT COUNT(*) FROM schema_version").fetchone()
            if row[0] == 0:
                conn.execute("INSERT INTO schema_version (version) VALUES (?)", (SCHEMA_VERSION,))

            # Seed chains
            from config import CHAINS
            for chain_id, info in CHAINS.items():
                conn.execute(
                    """INSERT OR IGNORE INTO chains (chain_id, name, native_symbol, explorer_url, icon)
                       VALUES (?, ?, ?, ?, ?)""",
                    (chain_id, info['name'], info['native_symbol'], info['explorer_url'], info.get('icon', ''))
                )

            conn.commit()
            logger.info("Database initialized successfully")
        except Exception as e:
            conn.rollback()
            logger.error(f"Database init error: {e}")
            raise


# --- User CRUD ---

def add_user(chat_id: int, username: str, is_admin: bool = False) -> bool:
    with _write_lock:
        conn = get_connection()
        try:
            conn.execute(
                "INSERT OR IGNORE INTO users (chat_id, username, is_admin) VALUES (?, ?, ?)",
                (chat_id, username, int(is_admin))
            )
            conn.commit()
            return True
        except Exception as e:
            logger.error(f"Error adding user {chat_id}: {e}")
            return False


def remove_user(chat_id: int) -> bool:
    with _write_lock:
        conn = get_connection()
        try:
            conn.execute("DELETE FROM settings WHERE chat_id = ?", (chat_id,))
            addr_ids = [r[0] for r in conn.execute(
                "SELECT id FROM addresses WHERE chat_id = ?", (chat_id,)
            ).fetchall()]
            if addr_ids:
                placeholders = ','.join('?' * len(addr_ids))
                conn.execute(
                    f"DELETE FROM transactions WHERE address_id IN ({placeholders})",
                    tuple(addr_ids)
                )
            conn.execute("DELETE FROM addresses WHERE chat_id = ?", (chat_id,))
            conn.execute("DELETE FROM users WHERE chat_id = ?", (chat_id,))
            conn.commit()
            return True
        except Exception as e:
            conn.rollback()
            logger.error(f"Error removing user {chat_id}: {e}")
            return False


def get_user(chat_id: int) -> Optional[Dict[str, Any]]:
    conn = get_connection()
    row = conn.execute("SELECT * FROM users WHERE chat_id = ?", (chat_id,)).fetchone()
    return dict(row) if row else None


def is_user_authorized(chat_id: int) -> bool:
    return get_user(chat_id) is not None


def is_user_admin(chat_id: int) -> bool:
    user = get_user(chat_id)
    return user is not None and user['is_admin'] == 1


def get_all_users() -> List[Dict[str, Any]]:
    conn = get_connection()
    rows = conn.execute("SELECT * FROM users").fetchall()
    return [dict(r) for r in rows]


# --- Address CRUD ---

def add_address(chat_id: int, chain_id: int, address: str, name: str, last_seen_block: int = 0) -> bool:
    with _write_lock:
        conn = get_connection()
        try:
            conn.execute(
                """INSERT INTO addresses (chat_id, chain_id, address, name, last_seen_block)
                   VALUES (?, ?, ?, ?, ?)""",
                (chat_id, chain_id, address, name, last_seen_block)
            )
            conn.commit()
            return True
        except sqlite3.IntegrityError:
            logger.warning(f"Address name '{name}' already exists for user {chat_id}")
            return False
        except Exception as e:
            logger.error(f"Error adding address: {e}")
            return False


def remove_address(chat_id: int, name: str) -> bool:
    with _write_lock:
        conn = get_connection()
        try:
            row = conn.execute(
                "SELECT id FROM addresses WHERE chat_id = ? AND name = ?",
                (chat_id, name)
            ).fetchone()
            if not row:
                return False
            addr_id = row[0]
            conn.execute("DELETE FROM transactions WHERE address_id = ?", (addr_id,))
            conn.execute("DELETE FROM addresses WHERE id = ?", (addr_id,))
            conn.commit()
            return True
        except Exception as e:
            conn.rollback()
            logger.error(f"Error removing address '{name}': {e}")
            return False


def get_address(chat_id: int, name: str) -> Optional[Dict[str, Any]]:
    conn = get_connection()
    row = conn.execute(
        "SELECT a.*, c.name as chain_name, c.native_symbol, c.explorer_url, c.icon "
        "FROM addresses a JOIN chains c ON a.chain_id = c.chain_id "
        "WHERE a.chat_id = ? AND a.name = ?",
        (chat_id, name)
    ).fetchone()
    return dict(row) if row else None


def get_user_addresses(chat_id: int) -> List[Dict[str, Any]]:
    conn = get_connection()
    rows = conn.execute(
        "SELECT a.*, c.name as chain_name, c.native_symbol, c.explorer_url, c.icon "
        "FROM addresses a JOIN chains c ON a.chain_id = c.chain_id "
        "WHERE a.chat_id = ? ORDER BY c.chain_id, a.name",
        (chat_id,)
    ).fetchall()
    return [dict(r) for r in rows]


def get_all_addresses() -> List[Dict[str, Any]]:
    """Get all monitored addresses across all users (for monitoring loop)."""
    conn = get_connection()
    rows = conn.execute(
        "SELECT a.*, c.name as chain_name, c.native_symbol, c.explorer_url, c.icon "
        "FROM addresses a JOIN chains c ON a.chain_id = c.chain_id "
        "ORDER BY a.chain_id, a.chat_id"
    ).fetchall()
    return [dict(r) for r in rows]


def update_last_seen_block(address_id: int, block_number: int) -> None:
    with _write_lock:
        conn = get_connection()
        try:
            conn.execute(
                "UPDATE addresses SET last_seen_block = ? WHERE id = ?",
                (block_number, address_id)
            )
            conn.commit()
        except Exception as e:
            logger.error(f"Error updating last_seen_block for address {address_id}: {e}")


def set_threshold(chat_id: int, name: str, threshold_wei: str) -> bool:
    with _write_lock:
        conn = get_connection()
        try:
            result = conn.execute(
                "UPDATE addresses SET threshold_wei = ? WHERE chat_id = ? AND name = ?",
                (threshold_wei, chat_id, name)
            )
            conn.commit()
            return result.rowcount > 0
        except Exception as e:
            logger.error(f"Error setting threshold: {e}")
            return False


# --- Transaction CRUD ---

def store_transaction(
    address_id: int, chain_id: int, tx_hash: str, block_number: int,
    from_addr: str, to_addr: str, value_wei: str,
    tx_type: str = 'native', token_address: str = None,
    token_symbol: str = None, token_decimals: int = None,
    direction: str = 'unknown', timestamp: int = None, func_name: str = None
) -> bool:
    with _write_lock:
        conn = get_connection()
        try:
            conn.execute(
                """INSERT OR IGNORE INTO transactions
                   (address_id, chain_id, tx_hash, block_number, from_addr, to_addr,
                    value_wei, tx_type, token_address, token_symbol, token_decimals,
                    direction, timestamp, func_name)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (address_id, chain_id, tx_hash, block_number, from_addr, to_addr,
                 value_wei, tx_type, token_address, token_symbol, token_decimals,
                 direction, timestamp, func_name)
            )
            conn.commit()
            return True
        except Exception as e:
            logger.error(f"Error storing transaction {tx_hash}: {e}")
            return False


def get_known_tokens(address_id: int) -> List[Dict[str, Any]]:
    """Get distinct ERC-20 tokens seen in transactions for an address."""
    conn = get_connection()
    rows = conn.execute(
        "SELECT DISTINCT token_address, token_symbol, token_decimals "
        "FROM transactions WHERE address_id = ? AND tx_type = 'erc20' AND token_address IS NOT NULL",
        (address_id,)
    ).fetchall()
    return [dict(r) for r in rows]


def get_transactions(address_id: int, limit: int = 10, offset: int = 0) -> Tuple[List[Dict[str, Any]], int]:
    """Return (transactions, total_count) for pagination."""
    conn = get_connection()
    total = conn.execute(
        "SELECT COUNT(*) FROM transactions WHERE address_id = ?",
        (address_id,)
    ).fetchone()[0]

    rows = conn.execute(
        "SELECT * FROM transactions WHERE address_id = ? ORDER BY block_number DESC LIMIT ? OFFSET ?",
        (address_id, limit, offset)
    ).fetchall()

    return [dict(r) for r in rows], total


# --- Settings CRUD ---

def get_setting(chat_id: int, key: str, default: str = None) -> Optional[str]:
    conn = get_connection()
    row = conn.execute(
        "SELECT value FROM settings WHERE chat_id = ? AND key = ?",
        (chat_id, key)
    ).fetchone()
    return row[0] if row else default


def set_setting(chat_id: int, key: str, value: str) -> None:
    with _write_lock:
        conn = get_connection()
        try:
            conn.execute(
                "INSERT OR REPLACE INTO settings (chat_id, key, value) VALUES (?, ?, ?)",
                (chat_id, key, value)
            )
            conn.commit()
        except Exception as e:
            logger.error(f"Error setting {key} for {chat_id}: {e}")


# --- Rate limiting ---

def check_rate_limit(chat_id: int, max_commands: int = 30, window: int = 60) -> bool:
    """Return True if the user is within rate limits."""
    now = time.time()
    with _rate_limit_lock:
        # Evict stale users if store is too large
        if len(_rate_limit_store) > _MAX_RATE_LIMIT_USERS:
            stale = [
                uid for uid, timestamps in _rate_limit_store.items()
                if not timestamps or now - timestamps[-1] > window
            ]
            for uid in stale:
                del _rate_limit_store[uid]

        if chat_id not in _rate_limit_store:
            _rate_limit_store[chat_id] = []

        # Prune old entries
        _rate_limit_store[chat_id] = [
            t for t in _rate_limit_store[chat_id] if now - t < window
        ]

        if len(_rate_limit_store[chat_id]) >= max_commands:
            return False

        _rate_limit_store[chat_id].append(now)
        return True


# --- Chain helpers ---

def get_chains() -> List[Dict[str, Any]]:
    conn = get_connection()
    rows = conn.execute("SELECT * FROM chains ORDER BY chain_id").fetchall()
    return [dict(r) for r in rows]


def get_chain(chain_id: int) -> Optional[Dict[str, Any]]:
    conn = get_connection()
    row = conn.execute("SELECT * FROM chains WHERE chain_id = ?", (chain_id,)).fetchone()
    return dict(row) if row else None