#!/opt/cloudlinux/venv/bin/python3 -sbb
# coding=utf-8
#
# Copyright © Cloud Linux GmbH & Cloud Linux Software, Inc 2010-2020 All Rights Reserved
#
# Licensed under CLOUD LINUX LICENSE AGREEMENT
# http://cloudlinux.com/docs/LICENCE.TXT
#

import os
import contextlib
import sqlite3
import logging

from datetime import datetime, timedelta
from sqlalchemy import (
    Column,
    Boolean,
    DateTime,
    Index,
    Integer,
    String,
    create_engine,
    event, func, text
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.orm import Session
from sqlalchemy.orm.session import close_all_sessions
from sqlalchemy.exc import DatabaseError

SSA_DB = '/var/lve/ssa.db'
OLD_SSA_DB = SSA_DB + '.old'
RETENTION_TIME_DAYS = 1

Base = declarative_base()

logger = logging.getLogger("cleanup_old_data")

class RequestResult(Base):
    """
    Describes processed request stored in database file.

    E.g.
    {
        "timestamp": "1650008727",
        "url": "http://mydomain.com/index.php",
        "duration": 162077,
        "hitting_limits": false,
        "throttled_time": 0,
        "io_throttled_time": 0,
        "wordpress": true
    }

    Note: created_at, updated_at is saved in local TZ format
    """
    __tablename__ = 'scrape_result'

    # Composite (domain, path) index: the daily-routine queries all filter
    # by domain and ORDER BY path. With this index SQLite streams rows in
    # index order and skips the external "TEMP B-TREE FOR ORDER BY" sort,
    # whose spill files would otherwise land in the SQLite temp dir — and
    # on hosts where that dir is tmpfs the spill is charged against the
    # MemoryMax=1G cgroup, recreating the OOM through a different path
    # (CLPRO-3077). It also subsumes the old standalone `domain` index as a
    # prefix, so we keep just this one to limit write amplification on the
    # insert-heavy table. NOTE: create_all() only fires for brand-new DBs;
    # existing multi-GB DBs get this index via ensure_indexes() instead.
    __table_args__ = (
        Index('ix_scrape_result_domain_path', 'domain', 'path'),
    )

    id = Column(Integer, primary_key=True)

    domain = Column(String, nullable=False)
    path = Column(String, nullable=False)

    timestamp = Column(Integer, nullable=False)

    duration = Column(Integer, nullable=False)
    is_slow_request = Column(Boolean, nullable=False)
    hitting_limits = Column(Boolean, nullable=False)
    throttled_time = Column(Integer, nullable=False)
    io_throttled_time = Column(Integer, nullable=False)
    wordpress = Column(Boolean, nullable=False)

    # Index on created_at speeds up cleanup_old_data queries significantly
    created_at = Column(DateTime(timezone=True), server_default=func.now(), index=True)
    updated_at = Column(DateTime(timezone=True), onupdate=func.now(), server_default=func.now())


def cleanup_old_data(engine, batch_size=10000):
    """
    Removes outdated records from database, saving disk space.
    Deletions are batched to avoid loading the entire database into
    the cgroup's page cache, which causes OOM on large databases.
    Runs VACUUM after deletion to reclaim disk space.
    """
    n_days_ago = datetime.today() - timedelta(days=RETENTION_TIME_DAYS)
    total_deleted = 0
    while True:
        with session_scope(engine) as session:
            result = session.execute(text(
                "DELETE FROM scrape_result WHERE rowid IN "
                "(SELECT rowid FROM scrape_result WHERE created_at < :cutoff LIMIT :batch)"
            ), {"cutoff": n_days_ago, "batch": batch_size})
            deleted = result.rowcount
            total_deleted += deleted
            if deleted == 0:
                break
    logger.info("Cleanup deleted %d old records", total_deleted)
    vacuum_database(engine)


def _redirect_temp_store_to_disk(cursor, operation):
    """
    Point SQLite's temp store at the database directory (on disk) for the
    given operation. SQLite spills large sorts / temp B-trees and VACUUM
    copies to the system temp directory, which can be tmpfs (RAM-backed).
    When it is, the spill is charged against the MemoryMax=1G cgroup and
    OOM-kills the agent on large databases — the same hazard for VACUUM
    and for index builds / ORDER BY sorts (CLPRO-2819, CLPRO-3077).
    """
    db_dir = os.path.dirname(os.path.abspath(SSA_DB))
    if os.access(db_dir, os.W_OK):
        cursor.execute(f"PRAGMA temp_store_directory = '{db_dir}'")
        logger.info("Set SQLite temp_store_directory to: %s", db_dir)
    else:
        logger.warning(
            "Directory %s not writable. %s will fall back to the default "
            "temp location.", db_dir, operation)


def vacuum_database(engine):
    """
    Run VACUUM command to reclaim disk space after deletions.
    VACUUM must be run outside of a transaction, so we use raw connection.
    """
    # Get a raw connection from the pool
    raw_conn = engine.raw_connection()
    original_isolation_level = raw_conn.isolation_level

    try:
        raw_conn.isolation_level = None
        cursor = raw_conn.cursor()

        _redirect_temp_store_to_disk(cursor, 'VACUUM')

        logger.info("Starting VACUUM operation...")
        cursor.execute("VACUUM")
        logger.info("VACUUM completed successfully.")

        cursor.close()
    except Exception as e:
        logger.error(f"VACUUM failed: {str(e)}")
    finally:
        raw_conn.isolation_level = original_isolation_level
        raw_conn.close()


def ensure_indexes(engine):
    """
    Create the composite (domain, path) index on existing databases.

    The schema is only materialised by ``create_all()`` in
    ``create_db_if_not_exist()``, which runs solely when the DB file does
    not yet exist — there are no migrations. So a model-level index reaches
    new installs only, while the multi-GB databases that actually OOM never
    get it. This explicit ``CREATE INDEX IF NOT EXISTS`` closes that gap; it
    is idempotent and a cheap no-op once the index exists (CLPRO-3077).

    Building the index on a large table sorts it, so we first redirect the
    temp store to disk to avoid the very tmpfs spill we are trying to fix.
    """
    raw_conn = engine.raw_connection()
    try:
        cursor = raw_conn.cursor()
        _redirect_temp_store_to_disk(cursor, 'CREATE INDEX')
        logger.info("Ensuring ix_scrape_result_domain_path exists...")
        cursor.execute(
            "CREATE INDEX IF NOT EXISTS ix_scrape_result_domain_path "
            "ON scrape_result(domain, path)")
        raw_conn.commit()
        logger.info("Index ix_scrape_result_domain_path is present.")
        cursor.close()
    except Exception as e:
        logger.error("Failed to ensure indexes: %s", str(e))
    finally:
        raw_conn.close()



def create_db_if_not_exist(engine):
    if not is_db_present(engine):
        Base.metadata.create_all(engine)


def is_db_present(engine):
    if not os.path.isfile(SSA_DB):
        return False
    database_inspection = Inspector.from_engine(engine)
    tables = [table for table in database_inspection.get_table_names()]
    return len(tables) > 0


def setup_wal_mode(dbapi_con, con_record):
    dbapi_con.execute('PRAGMA journal_mode = WAL')


def setup_temp_store_dir(dbapi_con, con_record):
    """
    Keep SQLite's temp spills (sort / GROUP BY temp B-trees) on disk rather
    than the default temp dir, which may be tmpfs (RAM-backed) and thus
    charged against the MemoryMax=1G cgroup. ``temp_store_directory`` is a
    process-wide setting, but applying it on every connect is cheap and
    guarantees it is in force before the daily routine runs (CLPRO-3077).

    Complements the (domain, path) index: the index removes the per-domain
    ORDER BY sort entirely, this catches whatever spill remains (e.g. the
    GROUP BY on the computed hour). Best-effort — never block a connection.
    """
    db_dir = os.path.dirname(os.path.abspath(SSA_DB))
    if not os.access(db_dir, os.W_OK):
        return
    try:
        dbapi_con.execute(f"PRAGMA temp_store_directory = '{db_dir}'")
    except sqlite3.Error as e:
        logger.warning("Could not set temp_store_directory to %s: %s",
                       db_dir, e)


def _setup_database(readonly):
    connection_string = f'file:{SSA_DB}'
    if readonly:
        connection_string = f'{connection_string}?mode=ro'
    creator = lambda: sqlite3.connect(connection_string, uri=True)
    engine = create_engine(
        'sqlite:////', creator=creator, echo=False,
    )
    event.listen(engine, 'connect', setup_wal_mode)
    event.listen(engine, 'connect', setup_temp_store_dir)
    create_db_if_not_exist(engine)
    return engine


def setup_database(readonly=False):
    return _setup_database(readonly)


def restore_database(engine):
    """
    Restore database by establish connections to old and new databases,
    merge data to new one if possible and delete old one.
    """
    if os.path.exists(SSA_DB):
        # Closing all sessions to ensure that no sessions is using database during replacing
        close_all_sessions()
        os.replace(SSA_DB, OLD_SSA_DB)

    new_engine = setup_database()
    # Dispose of the existing engine to close and refresh all connections, ensuring it connects to the new database
    engine.dispose()
    old_engine = create_engine(f'sqlite:////{OLD_SSA_DB}')

    try:
        with session_scope(old_engine) as session_old, session_scope(new_engine) as session_new:
            # Check if old database is able to read and it make sense to try save unharmed data
            session_old.query(RequestResult).first()
            merge_unharmed_data_from_database(session_old, session_new)
    except DatabaseError:
        pass

    for path in (OLD_SSA_DB, SSA_DB + "-wal", SSA_DB + "-shm"):
        if os.path.exists(path):
            os.remove(path)


def merge_unharmed_data_from_database(session_old, session_new):
    """
    Scrape all unharmed records from malformed database and merge them into new database.
    """
    offset = 0
    batch_size = 10
    while True:
        query = session_old.query(RequestResult).offset(offset).limit(batch_size)
        try:
            records_to_save = query.all()
            if not records_to_save:
                break
            for record in records_to_save:
                session_new.merge(record)
        except DatabaseError:
            for pos_in_batch in range(batch_size):
                try:
                    record_to_save = query.offset(offset + pos_in_batch).first()
                    if not record_to_save:
                        break
                    session_new.merge(record_to_save)
                except DatabaseError:
                    pass
        session_new.commit()
        offset += batch_size


def is_malformed_database(engine):
    """
    Try integrity check of database file to see if it is malformed.
    If database unable to execute it, will also count as malformed.
    """
    if os.path.exists(OLD_SSA_DB) and os.path.exists(SSA_DB):
        os.remove(OLD_SSA_DB)
    try:
        with session_scope(engine) as db:
            result = db.execute(text("PRAGMA integrity_check"))
            errors = result.fetchall()
            return errors[0][0] != 'ok'
    except DatabaseError:
        return True


@contextlib.contextmanager
def session_scope(engine) -> Session:
    """
    Provide a transactional scope around a series of operations.
    """
    session = Session(bind=engine)
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()