scheduler/src/main.py

"""
The Scheduler - System-wide maintenance orchestration.
Handles backups, documentation mirroring, cleanup, and automated tasks.

Architecture: Hybrid APScheduler + DB-based priority system
- APScheduler runs a single job every minute
- Job queries DB for tasks scheduled in that minute
- Executes up to 5 tasks concurrently based on priority
"""
from fastapi import FastAPI, HTTPException, Depends, Header
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
from apscheduler.triggers.cron import CronTrigger
from contextlib import asynccontextmanager
import logging

from src.config import get_settings, Settings
from src.tasks.executor import TaskExecutor
from src.models import TaskCreate, TaskUpdate, TaskResponse

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Global instances
scheduler: AsyncIOScheduler | None = None
task_executor: TaskExecutor | None = None


def get_scheduler() -> AsyncIOScheduler:
    """Dependency to get scheduler instance."""
    if scheduler is None:
        raise HTTPException(500, "Scheduler not initialized")
    return scheduler


def get_task_executor() -> TaskExecutor:
    """Dependency to get task executor instance."""
    if task_executor is None:
        raise HTTPException(500, "Task executor not initialized")
    return task_executor

# Lifespan manager for startup/shutdown
@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage application lifecycle - startup and shutdown."""
    global scheduler, task_executor

    settings = get_settings()

    # Startup
    logger.info("=" * 60)
    logger.info("The Scheduler - System-wide Maintenance Orchestration")
    logger.info("=" * 60)
    logger.info(f"Architecture: Hybrid APScheduler + DB-based priority system")
    logger.info(f"Database: {settings.postgres_host}:{settings.postgres_port}/{settings.postgres_db}")
    logger.info(f"API: http://{settings.host}:{settings.port}")
    logger.info(f"Docs: http://{settings.host}:{settings.port}/docs")
    logger.info("=" * 60)

    # Initialize task executor
    task_executor = TaskExecutor(settings)
    logger.info("Task executor initialized (max 5 concurrent tasks)")

    # Initialize APScheduler with minimal configuration
    # No jobstore needed - we only have one in-memory job
    scheduler = AsyncIOScheduler(
        job_defaults={
            'coalesce': True,        # Combine missed runs
            'max_instances': 1,      # Only one instance running
            'misfire_grace_time': 30  # 30s grace period for minute-based execution
        }
    )

    # Add the single minute-based task processor
    scheduler.add_job(
        func=task_executor.process_minute,
        trigger=CronTrigger(minute='*'),  # Run every minute
        id='process_tasks',
        name='Process scheduled tasks',
        replace_existing=True
    )

    try:
        scheduler.start()
        logger.info("Scheduler started - processing tasks every minute")
        logger.info("Priority system: 1-5 (emergency/system), 10-30 (user), 40-70+ (maintenance)")
    except Exception as e:
        logger.error(f"Failed to start scheduler: {e}")
        raise

    yield

    # Shutdown
    logger.info("Stopping The Scheduler...")
    if scheduler:
        scheduler.shutdown(wait=True)
    logger.info("Scheduler stopped")


# FastAPI app
app = FastAPI(
    title="The Scheduler",
    version="1.0.0",
    description="System-wide maintenance orchestration - backups, doc mirroring, cleanup, task automation",
    lifespan=lifespan
)

# Dependencies
async def verify_api_key(
    authorization: str = Header(None),
    settings: Settings = Depends(get_settings)
):
    """Verify API key from Authorization header."""
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(401, "Missing API key")
    key = authorization.replace("Bearer ", "")
    if key != settings.scheduler_api_key:
        raise HTTPException(403, "Invalid API key")
    return key

# ============================================================================
# Public Endpoints (no auth required)
# ============================================================================

@app.get("/health")
async def health(
    settings: Settings = Depends(get_settings),
    sched: AsyncIOScheduler = Depends(get_scheduler)
):
    """Health check endpoint."""
    return {
        "status": "healthy",
        "scheduler_running": sched.running,
        "jobs_count": len(sched.get_jobs()),
        "database": settings.postgres_db
    }

# ============================================================================
# Protected Endpoints (require API key)
# ============================================================================

@app.get("/tasks")
async def list_tasks(
    enabled: bool = None,
    service: str = None,
    api_key: str = Depends(verify_api_key),
    executor: TaskExecutor = Depends(get_task_executor)
):
    """List all scheduled tasks from database."""
    import psycopg2.extras

    query = "SELECT * FROM scheduled_tasks WHERE 1=1"
    params = []

    if enabled is not None:
        query += " AND enabled = %s"
        params.append(enabled)

    if service:
        query += " AND service = %s"
        params.append(service)

    query += " ORDER BY priority ASC, task_name ASC"

    with executor.get_db_connection() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(query, params)
            tasks = [dict(task) for task in cur.fetchall()]

    return {
        "tasks": tasks,
        "count": len(tasks)
    }

@app.get("/tasks/{task_name}")
async def get_task_details(
    task_name: str,
    api_key: str = Depends(verify_api_key),
    executor: TaskExecutor = Depends(get_task_executor)
):
    """Get details for a specific task."""
    import psycopg2.extras

    with executor.get_db_connection() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute("SELECT * FROM scheduled_tasks WHERE task_name = %s", (task_name,))
            task = cur.fetchone()

    if not task:
        raise HTTPException(404, f"Task '{task_name}' not found")

    return dict(task)

@app.post("/tasks", response_model=TaskResponse, tags=["Task Management"])
async def create_task(
    task: TaskCreate,
    api_key: str = Depends(verify_api_key),
    executor: TaskExecutor = Depends(get_task_executor)
):
    """
    Create a new scheduled task.

    ## Schedule Pattern

    Use cron-style fields where `-1` means "every":
    - `minute: -1, hour: -1` → Runs every minute
    - `minute: 0, hour: -1` → Runs at minute 0 of every hour (hourly)
    - `minute: 0, hour: 2` → Runs at 2:00 AM every day
    - `minute: 0, hour: 2, day_of_week: 0` → Runs at 2:00 AM every Monday

    ## Priority Levels

    - **1-5**: Emergency/critical system tasks
    - **10-30**: User-initiated tasks
    - **40-70**: Background maintenance
    - **70+**: Low priority cleanup

    ## Executor Types

    ### `rest_api` Executor

    Calls HTTP endpoints. Config format:
    ```json
    {
      "method": "POST",
      "url": "http://service:port/endpoint",
      "headers": {
        "Authorization": "Bearer ${ENV_VAR}",
        "Content-Type": "application/json"
      },
      "body": {"key": "value"}
    }
    ```

    Environment variables can be referenced with `${VAR_NAME}` syntax.

    ## Example: Librarian Consolidation Task

    ```json
    {
      "task_name": "librarian_consolidation",
      "service": "library-desk",
      "executor": "rest_api",
      "priority": 25,
      "description": "Processes search queries and consolidates knowledge",
      "minute": 0,
      "hour": -1,
      "config": {
        "method": "POST",
        "url": "http://library-desk:8089/consolidation/run",
        "headers": {
          "Authorization": "Bearer ${LIBRARY_DESK_API_KEY}",
          "Content-Type": "application/json"
        },
        "body": {
          "dry_run": false,
          "process_limit": 10
        }
      }
    }
    ```

    This creates an hourly task that calls the consolidation endpoint.
    """
    import psycopg2.extras
    import json

    # Convert task model to dict
    task_data = task.model_dump()

    # Convert config dict to JSON string if present
    if 'config' in task_data and isinstance(task_data['config'], dict):
        task_data['config'] = json.dumps(task_data['config'])

    with executor.get_db_connection() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute("""
                INSERT INTO scheduled_tasks
                    (task_name, service, executor, priority, minute, hour,
                     day_of_month, month, day_of_week, enabled, description,
                     config, max_retries, timeout_seconds, created_by)
                VALUES
                    (%(task_name)s, %(service)s, %(executor)s, %(priority)s,
                     %(minute)s, %(hour)s, %(day_of_month)s, %(month)s,
                     %(day_of_week)s, %(enabled)s, %(description)s,
                     %(config)s::jsonb, %(max_retries)s, %(timeout_seconds)s,
                     %(created_by)s)
                RETURNING *
            """, task_data)
            new_task = dict(cur.fetchone())
            conn.commit()

    logger.info(f"Created task: {new_task['task_name']}")
    return new_task

@app.put("/tasks/{task_name}")
async def update_task(
    task_name: str,
    task_data: dict,
    api_key: str = Depends(verify_api_key),
    executor: TaskExecutor = Depends(get_task_executor)
):
    """Update an existing scheduled task."""
    import psycopg2.extras
    import json

    # Build update query dynamically
    allowed_fields = ['service', 'executor', 'priority', 'minute', 'hour',
                     'day_of_month', 'month', 'day_of_week', 'enabled',
                     'description', 'config', 'max_retries', 'timeout_seconds']

    updates = {k: v for k, v in task_data.items() if k in allowed_fields}
    if not updates:
        raise HTTPException(400, "No valid fields to update")

    # Convert config dict to JSON string if present
    if 'config' in updates and isinstance(updates['config'], dict):
        updates['config'] = json.dumps(updates['config'])

    set_clause = ', '.join([f"{k} = %({k})s" for k in updates.keys()])

    with executor.get_db_connection() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(f"""
                UPDATE scheduled_tasks
                SET {set_clause}, updated_at = NOW()
                WHERE task_name = %(task_name)s
                RETURNING *
            """, {**updates, 'task_name': task_name})

            updated_task = cur.fetchone()
            if not updated_task:
                raise HTTPException(404, f"Task '{task_name}' not found")

            conn.commit()

    logger.info(f"Updated task: {task_name}")
    return dict(updated_task)

@app.delete("/tasks/{task_name}")
async def delete_task(
    task_name: str,
    api_key: str = Depends(verify_api_key),
    executor: TaskExecutor = Depends(get_task_executor)
):
    """Delete a scheduled task."""
    with executor.get_db_connection() as conn:
        with conn.cursor() as cur:
            cur.execute("""
                DELETE FROM scheduled_tasks
                WHERE task_name = %s
                RETURNING task_name
            """, (task_name,))

            deleted = cur.fetchone()
            if not deleted:
                raise HTTPException(404, f"Task '{task_name}' not found")

            conn.commit()

    logger.info(f"Deleted task: {task_name}")
    return {"message": f"Task '{task_name}' deleted successfully"}

@app.post("/tasks/{task_name}/trigger")
async def trigger_task(
    task_name: str,
    api_key: str = Depends(verify_api_key),
    executor: TaskExecutor = Depends(get_task_executor)
):
    """Manually trigger a task to run immediately."""
    import psycopg2.extras

    # Get task details
    with executor.get_db_connection() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute("SELECT * FROM scheduled_tasks WHERE task_name = %s", (task_name,))
            task = cur.fetchone()

    if not task:
        raise HTTPException(404, f"Task '{task_name}' not found")

    if not task['enabled']:
        raise HTTPException(400, f"Task '{task_name}' is disabled")

    # Execute task immediately in background
    import asyncio
    asyncio.create_task(executor.execute_task(dict(task)))

    logger.info(f"Manually triggered task: {task_name}")
    return {
        "message": f"Task '{task_name}' triggered successfully",
        "task_name": task_name,
        "priority": task['priority'],
        "executor": task['executor']
    }

# Legacy endpoints (deprecated)
@app.post("/tasks/backup")
async def trigger_backup(api_key: str = Depends(verify_api_key)):
    """Trigger backup tasks manually (deprecated - use POST /tasks/{name}/trigger)"""
    # TODO: Implement backup executor
    logger.info("Manual backup triggered")
    return {
        "message": "Backup task triggered",
        "status": "not_implemented",
        "note": "Backup executor needs to be implemented"
    }

@app.post("/tasks/docs/update")
async def trigger_docs_update(
    project: str = None,
    api_key: str = Depends(verify_api_key)
):
    """Trigger documentation mirror update"""
    # TODO: Implement doc mirror executor
    logger.info(f"Doc mirror update triggered for project: {project or 'all'}")
    return {
        "message": f"Documentation update triggered for {project or 'all projects'}",
        "status": "not_implemented",
        "note": "Doc mirror executor needs to be implemented"
    }

@app.post("/tasks/docs/check-versions")
async def check_doc_versions(api_key: str = Depends(verify_api_key)):
    """Check for new documentation versions"""
    # TODO: Implement version check executor
    logger.info("Version check triggered")
    return {
        "message": "Version check triggered",
        "status": "not_implemented",
        "note": "Version check executor needs to be implemented"
    }

@app.post("/tasks/cleanup")
async def trigger_cleanup(api_key: str = Depends(verify_api_key)):
    """Run cleanup tasks"""
    # TODO: Implement cleanup executor
    logger.info("Cleanup task triggered")
    return {
        "message": "Cleanup task triggered",
        "status": "not_implemented",
        "note": "Cleanup executor needs to be implemented"
    }

@app.get("/executions")
async def task_history(
    limit: int = 20,
    task_name: str = None,
    service: str = None,
    status: str = None,
    api_key: str = Depends(verify_api_key),
    executor: TaskExecutor = Depends(get_task_executor)
):
    """View task execution history."""
    import psycopg2.extras

    query = "SELECT * FROM task_executions WHERE 1=1"
    params = []

    if task_name:
        query += " AND task_name = %s"
        params.append(task_name)

    if service:
        query += " AND service = %s"
        params.append(service)

    if status:
        query += " AND status = %s"
        params.append(status)

    query += " ORDER BY triggered_at DESC LIMIT %s"
    params.append(limit)

    with executor.get_db_connection() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(query, params)
            executions = [dict(ex) for ex in cur.fetchall()]

    return {
        "executions": executions,
        "count": len(executions),
        "limit": limit
    }

@app.get("/stats")
async def stats(
    api_key: str = Depends(verify_api_key),
    settings: Settings = Depends(get_settings),
    sched: AsyncIOScheduler = Depends(get_scheduler),
    task_exec: TaskExecutor = Depends(get_task_executor)
):
    """Get system statistics."""
    import psycopg2.extras

    # Query task stats from database
    with task_exec.get_db_connection() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            # Count enabled tasks
            cur.execute("SELECT COUNT(*) as count FROM scheduled_tasks WHERE enabled = true")
            enabled_tasks = cur.fetchone()['count']

            # Count running tasks
            cur.execute("SELECT COUNT(*) as count FROM task_executions WHERE status = 'running'")
            running_tasks = cur.fetchone()['count']

            # Recent execution stats (last 24 hours)
            cur.execute("""
                SELECT status, COUNT(*) as count
                FROM task_executions
                WHERE triggered_at > NOW() - INTERVAL '24 hours'
                GROUP BY status
            """)
            execution_stats = {row['status']: row['count'] for row in cur.fetchall()}

    return {
        "scheduler_running": sched.running,
        "minute_processor_active": True,  # If we got here, it's running
        "database": settings.postgres_db,
        "tasks_enabled": enabled_tasks,
        "tasks_currently_running": running_tasks,
        "concurrent_limit": 5,
        "execution_stats_24h": execution_stats,
        "priority_system": "1-5 (emergency/system), 10-30 (user), 40-70+ (maintenance)"
    }