UAV-CO2/src/gasflux/janitor.py

"""
Janitor Module
Handles background cleanup of expired tasks and their output directories.
"""

import threading
import time
import shutil
from pathlib import Path
import sqlite3
from flask import current_app


def start_janitor(app):
    """Start the background janitor thread for cleaning up expired tasks."""
    def worker():
        with app.app_context():
            while True:
                try:
                    cleanup_expired_tasks()
                except Exception as e:
                    app.logger.error(f"Janitor cleanup error: {str(e)}", exc_info=True)
                finally:
                    # Sleep for configured interval before next cleanup cycle
                    cleanup_interval = int(current_app.config.get('TASK_CLEANUP_INTERVAL', 30))
                    time.sleep(cleanup_interval)

    # Create daemon thread so it doesn't prevent app shutdown
    thread = threading.Thread(target=worker, daemon=True, name="janitor")
    thread.start()
    app.logger.info("Janitor thread started for background cleanup")


def cleanup_expired_tasks():
    """Clean up tasks that have exceeded their deletion time."""
    db_path = _get_db_path()
    dry_run_config = current_app.config.get('JANITOR_DRY_RUN', 'false')
    if isinstance(dry_run_config, str):
        dry_run = dry_run_config.lower() == 'true'
    else:
        dry_run = bool(dry_run_config)

    # Get cleanup intervals from config (in seconds)
    failed_task_cleanup_age = current_app.config.get('FAILED_TASK_CLEANUP_AGE', 86400)  # Default 24 hours for failed tasks
    current_app.logger.info(f"Janitor: FAILED_TASK_CLEANUP_AGE = {failed_task_cleanup_age} seconds")

    try:
        conn = sqlite3.connect(str(db_path), check_same_thread=False)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA foreign_keys=ON")

        # Log current time for debugging (using same timezone as setting)
        current_time = conn.execute("SELECT datetime('now', '+8 hours')").fetchone()[0]
        current_app.logger.info(f"Janitor cleanup check at: {current_time}")

        # Bases used by both sections
        upload_base = Path(current_app.config.get('UPLOAD_FOLDER') or '')
        output_base = Path(current_app.config.get('OUTPUT_FOLDER') or '')

        # Debug: Check for tasks with delete_after_at set
        debug_rows = conn.execute("""
            SELECT task_id, delete_after_at, downloaded_at
            FROM tasks
            WHERE delete_after_at IS NOT NULL
              AND deleted_at IS NULL
        """).fetchall()
        if debug_rows:
            current_app.logger.info(f"Found {len(debug_rows)} tasks with delete_after_at set:")
            for row in debug_rows:
                current_app.logger.info(f"  Task {row['task_id']}: delete_at={row['delete_after_at']}, downloaded_at={row['downloaded_at']}")

        # 1) Scheduled deletions (delete_after_at)
        rows = conn.execute("""
            SELECT task_id, output_dir
            FROM tasks
            WHERE delete_after_at IS NOT NULL
              AND deleted_at IS NULL
              AND delete_after_at <= datetime('now', '+8 hours')
        """).fetchall()

        if rows:
         current_app.logger.info(f"Janitor found {len(rows)} expired tasks to clean up")

        for row in rows:
            task_id = row['task_id']
            output_dir = row['output_dir']

            try:
                delete_targets = []

                # 记录在库中的 output_dir
                if output_dir:
                    p = Path(output_dir)
                    delete_targets.append(p)

                # 兜底：按约定 outputs/<task_id>
                derived_output_dir = output_base / task_id if output_base else None
                if derived_output_dir and derived_output_dir not in delete_targets:
                    delete_targets.append(derived_output_dir)

                # 同时删除 uploads/<task_id>
                derived_upload_dir = upload_base / task_id if upload_base else None
                if derived_upload_dir:
                    delete_targets.append(derived_upload_dir)

                if dry_run:
                    for tgt in delete_targets:
                        if tgt:
                            current_app.logger.info(f"[DRY RUN] Would delete task {task_id} path: {tgt}")
                else:
                    # 实际删除
                    for tgt in delete_targets:
                        try:
                            if tgt and tgt.exists():
                                shutil.rmtree(tgt, ignore_errors=True)
                                current_app.logger.info(f"Deleted path for task {task_id}: {tgt}")
                            else:
                                current_app.logger.warning(f"Path not found for task {task_id}: {tgt}")
                        except Exception as e:
                            current_app.logger.error(f"Failed to delete path {tgt} for task {task_id}: {e}")

                    # Hard delete from database
                    conn.execute(
                        "DELETE FROM tasks WHERE task_id = ?",
                        (task_id,)
                    )
                    conn.commit()
                    current_app.logger.info(f"Hard deleted task {task_id} from database")

            except Exception as e:
                current_app.logger.error(f"Failed to delete task {task_id}: {str(e)}", exc_info=True)
                # Continue with other tasks even if one fails

        # Debug: Check all failed tasks
        debug_failed = conn.execute("""
            SELECT task_id, status, created_at, started_at, finished_at, deleted_at
            FROM tasks
            WHERE status = 'failed' AND deleted_at IS NULL
        """).fetchall()
        if debug_failed:
            current_app.logger.info(f"Janitor found {len(debug_failed)} total failed tasks:")
            for row in debug_failed:
                current_app.logger.info(f"  Task {row['task_id']}: created={row['created_at']}, started={row['started_at']}, finished={row['finished_at']}")

        # Clean up failed tasks that are older than configured age
        failed_rows = conn.execute("""
            SELECT task_id, output_dir
            FROM tasks
            WHERE status = 'failed'
              AND deleted_at IS NULL
              AND COALESCE(finished_at, started_at, created_at) <= datetime('now', '+8 hours', '-' || ? || ' seconds')
        """, (failed_task_cleanup_age,)).fetchall()

        if failed_rows:
            current_app.logger.info(f"Janitor found {len(failed_rows)} failed tasks older than {failed_task_cleanup_age} seconds")
            for row in failed_rows:
                current_app.logger.info(f"  Failed task: {row['task_id']}")

            for row in failed_rows:
                task_id = row['task_id']
                output_dir = row['output_dir']

                try:
                    delete_targets = []

                    # 记录在库中的 output_dir
                    if output_dir:
                        p = Path(output_dir)
                        delete_targets.append(p)

                    # 兜底：按约定 outputs/<task_id>
                    derived_output_dir = output_base / task_id if output_base else None
                    if derived_output_dir and derived_output_dir not in delete_targets:
                        delete_targets.append(derived_output_dir)

                    # 同时删除 uploads/<task_id>
                    derived_upload_dir = upload_base / task_id if upload_base else None
                    if derived_upload_dir:
                        delete_targets.append(derived_upload_dir)

                    if dry_run:
                        for tgt in delete_targets:
                            if tgt:
                                current_app.logger.info(f"[DRY RUN] Would delete failed task {task_id} path: {tgt}")
                    else:
                        # 实际删除
                        for tgt in delete_targets:
                            try:
                                if tgt and tgt.exists():
                                    shutil.rmtree(tgt, ignore_errors=True)
                                    current_app.logger.info(f"Deleted failed task {task_id} path: {tgt}")
                                else:
                                    current_app.logger.warning(f"Path not found for failed task {task_id}: {tgt}")
                            except Exception as e:
                                current_app.logger.error(f"Failed to delete path {tgt} for failed task {task_id}: {e}")

                        # Hard delete from database
                        conn.execute(
                            "DELETE FROM tasks WHERE task_id = ?",
                            (task_id,)
                        )
                        conn.commit()
                        current_app.logger.info(f"Auto-deleted failed task {task_id} from database (older than {failed_task_cleanup_age} seconds)")

                except Exception as e:
                    current_app.logger.error(f"Failed to delete failed task {task_id}: {str(e)}", exc_info=True)
                    # Continue with other tasks even if one fails

    except Exception as e:
        current_app.logger.error(f"Database error during cleanup: {str(e)}", exc_info=True)
    finally:
        if 'conn' in locals():
            conn.close()


def reconcile_tasks_on_startup():
    """Reconcile task states on application startup."""
    db_path = _get_db_path()

    # Get cleanup ages from config (in seconds)
    successful_task_cleanup_age = current_app.config.get('SUCCESSFUL_TASK_CLEANUP_AGE', 3600)
    failed_task_cleanup_age = current_app.config.get('FAILED_TASK_CLEANUP_AGE', 86400)

    try:
        conn = sqlite3.connect(str(db_path), check_same_thread=False)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA foreign_keys=ON")

        # Fix tasks that were downloaded but don't have delete_after_at set
        # This handles cases where the app crashed after marking downloaded but before setting delete time
        conn.execute("""
            UPDATE tasks
            SET delete_after_at = COALESCE(delete_after_at, datetime(downloaded_at, '+' || ? || ' seconds'))
            WHERE downloaded_at IS NOT NULL
              AND deleted_at IS NULL
              AND delete_after_at IS NULL
        """, (successful_task_cleanup_age,))
        updated_count = conn.total_changes

        if updated_count > 0:
            current_app.logger.info(f"Startup reconciliation: Fixed delete_after_at for {updated_count} downloaded tasks")

        # Clean up failed tasks that are older than configured age
        failed_rows = conn.execute("""
            SELECT task_id, output_dir
            FROM tasks
            WHERE status = 'failed'
              AND deleted_at IS NULL
              AND COALESCE(finished_at, started_at, created_at) <= datetime('now', '+8 hours', '-' || ? || ' seconds')
        """, (failed_task_cleanup_age,)).fetchall()

        if failed_rows:
            current_app.logger.info(f"Startup reconciliation: Found {len(failed_rows)} failed tasks older than {failed_task_cleanup_age} seconds")

            upload_base = Path(current_app.config.get('UPLOAD_FOLDER') or '')
            output_base = Path(current_app.config.get('OUTPUT_FOLDER') or '')
            dry_run_config = current_app.config.get('JANITOR_DRY_RUN', 'false')

            if isinstance(dry_run_config, str):
                dry_run = dry_run_config.lower() == 'true'
            else:
                dry_run = bool(dry_run_config)

            for row in failed_rows:
                task_id = row['task_id']
                output_dir = row['output_dir']

                try:
                    delete_targets = []

                    # 记录在库中的 output_dir
                    if output_dir:
                        p = Path(output_dir)
                        delete_targets.append(p)

                    # 兜底：按约定 outputs/<task_id>
                    derived_output_dir = output_base / task_id if output_base else None
                    if derived_output_dir and derived_output_dir not in delete_targets:
                        delete_targets.append(derived_output_dir)

                    # 同时删除 uploads/<task_id>
                    derived_upload_dir = upload_base / task_id if upload_base else None
                    if derived_upload_dir:
                        delete_targets.append(derived_upload_dir)

                    if dry_run:
                        for tgt in delete_targets:
                            if tgt:
                                current_app.logger.info(f"[DRY RUN] Would delete failed task {task_id} path: {tgt}")
                    else:
                        # 实际删除
                        for tgt in delete_targets:
                            try:
                                if tgt and tgt.exists():
                                    import shutil
                                    shutil.rmtree(tgt, ignore_errors=True)
                                    current_app.logger.info(f"Deleted failed task {task_id} path: {tgt}")
                                else:
                                    current_app.logger.warning(f"Path not found for failed task {task_id}: {tgt}")
                            except Exception as e:
                                current_app.logger.error(f"Failed to delete path {tgt} for failed task {task_id}: {e}")

                        # Hard delete from database
                        conn.execute(
                            "DELETE FROM tasks WHERE task_id = ?",
                            (task_id,)
                        )
                        current_app.logger.info(f"Startup reconciliation: Auto-deleted failed task {task_id} from database (older than {failed_task_cleanup_age} seconds)")

                except Exception as e:
                    current_app.logger.error(f"Failed to delete failed task {task_id}: {str(e)}", exc_info=True)

        # 检查 output_dir 不存在的任务（仅警告，不自动删除）
        rows = conn.execute("""
            SELECT task_id, output_dir
            FROM tasks
            WHERE output_dir IS NOT NULL
              AND deleted_at IS NULL
        """).fetchall()

        orphaned_count = 0
        for row in rows:
            task_id_from_db = row[0]
            output_dir_from_db = row[1]

            if not Path(output_dir_from_db).exists():
                current_app.logger.warning(
                    f"Startup reconciliation: Task {task_id_from_db} output directory "
                    f"not found on disk: {output_dir_from_db}. Database record preserved."
                )
                orphaned_count += 1

        if orphaned_count > 0:
            current_app.logger.info(
                f"Startup reconciliation: Found {orphaned_count} tasks with missing directories "
                f"(records preserved, use DELETE /task/<id> to manually clean up)"
            )

        # Backfill output_dir for rows with NULL, using OUTPUT_FOLDER/<task_id>
        output_base = Path(current_app.config.get('OUTPUT_FOLDER') or '')
        if output_base:
            rows2 = conn.execute("""
                SELECT task_id FROM tasks
                WHERE (output_dir IS NULL OR output_dir = '')
                  AND deleted_at IS NULL
            """).fetchall()
            for r in rows2:
                task_id_from_db = r[0]  # task_id是第一个字段
                guess = output_base / task_id_from_db
                conn.execute("UPDATE tasks SET output_dir = ? WHERE task_id = ?", (str(guess), task_id_from_db))
                current_app.logger.info(f"Backfilled output_dir for task {task_id_from_db}: {guess}")

        conn.commit()

    except Exception as e:
        current_app.logger.error(f"Startup reconciliation error: {str(e)}", exc_info=True)
    finally:
        if 'conn' in locals():
            conn.close()


def _get_db_path():
    """Get database file path from app config."""
    from .db import get_db_path as get_config_db_path
    return get_config_db_path(current_app)