POTE/scripts/health_check.py

#!/usr/bin/env python3
"""
POTE Health Check Script

Checks the health of the POTE system and reports status.

Usage:
    python scripts/health_check.py
    python scripts/health_check.py --json
"""

import argparse
import json
import logging
import sys
from datetime import date, datetime, timedelta
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy import func

from pote.db import engine, get_session
from pote.db.models import MarketAlert, Official, Price, Security, Trade

logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger(__name__)


def check_database_connection() -> dict:
    """Check if database is accessible."""
    try:
        with engine.connect() as conn:
            conn.execute("SELECT 1")
        return {"status": "ok", "message": "Database connection successful"}
    except Exception as e:
        return {"status": "error", "message": f"Database connection failed: {str(e)}"}


def check_data_freshness() -> dict:
    """Check if data has been updated recently."""
    with get_session() as session:
        # Check most recent trade filing date
        latest_trade = (
            session.query(Trade).order_by(Trade.filing_date.desc()).first()
        )

        if not latest_trade:
            return {
                "status": "warning",
                "message": "No trades found in database",
                "latest_trade_date": None,
                "days_since_update": None,
            }

        days_since = (date.today() - latest_trade.filing_date).days

        if days_since > 7:
            status = "warning"
            message = f"Latest trade is {days_since} days old (may need update)"
        elif days_since > 14:
            status = "error"
            message = f"Latest trade is {days_since} days old (stale data)"
        else:
            status = "ok"
            message = f"Data is fresh ({days_since} days old)"

        return {
            "status": status,
            "message": message,
            "latest_trade_date": str(latest_trade.filing_date),
            "days_since_update": days_since,
        }


def check_data_counts() -> dict:
    """Check counts of key entities."""
    with get_session() as session:
        counts = {
            "officials": session.query(Official).count(),
            "securities": session.query(Security).count(),
            "trades": session.query(Trade).count(),
            "prices": session.query(Price).count(),
            "market_alerts": session.query(MarketAlert).count(),
        }

        if counts["trades"] == 0:
            status = "error"
            message = "No trades in database"
        elif counts["trades"] < 10:
            status = "warning"
            message = "Very few trades in database (< 10)"
        else:
            status = "ok"
            message = f"Database has {counts['trades']} trades"

        return {"status": status, "message": message, "counts": counts}


def check_recent_alerts() -> dict:
    """Check for recent market alerts."""
    with get_session() as session:
        yesterday = datetime.now() - timedelta(days=1)
        recent_alerts = (
            session.query(MarketAlert).filter(MarketAlert.timestamp >= yesterday).count()
        )

        return {
            "status": "ok",
            "message": f"{recent_alerts} alerts in last 24 hours",
            "recent_alerts_count": recent_alerts,
        }


def main():
    parser = argparse.ArgumentParser(description="POTE health check")
    parser.add_argument(
        "--json", action="store_true", help="Output results as JSON"
    )
    args = parser.parse_args()

    # Run all checks
    checks = {
        "database_connection": check_database_connection(),
        "data_freshness": check_data_freshness(),
        "data_counts": check_data_counts(),
        "recent_alerts": check_recent_alerts(),
    }

    # Determine overall status
    statuses = [check["status"] for check in checks.values()]
    if "error" in statuses:
        overall_status = "error"
    elif "warning" in statuses:
        overall_status = "warning"
    else:
        overall_status = "ok"

    result = {
        "timestamp": datetime.now().isoformat(),
        "overall_status": overall_status,
        "checks": checks,
    }

    if args.json:
        print(json.dumps(result, indent=2))
    else:
        # Human-readable output
        status_emoji = {"ok": "✓", "warning": "⚠", "error": "✗"}

        print("\n" + "=" * 60)
        print("POTE HEALTH CHECK")
        print("=" * 60)
        print(f"Timestamp: {result['timestamp']}")
        print(f"Overall Status: {status_emoji.get(overall_status, '?')} {overall_status.upper()}")
        print()

        for check_name, check_result in checks.items():
            status = check_result["status"]
            emoji = status_emoji.get(status, "?")
            print(f"{emoji} {check_name.replace('_', ' ').title()}: {check_result['message']}")

            # Print additional details if present
            if "counts" in check_result:
                for key, value in check_result["counts"].items():
                    print(f"    {key}: {value:,}")

        print("=" * 60 + "\n")

    # Exit with appropriate code
    if overall_status == "error":
        sys.exit(2)
    elif overall_status == "warning":
        sys.exit(1)
    else:
        sys.exit(0)


if __name__ == "__main__":
    main()