from flask import Blueprint, current_app, request
from flask_healthz import HealthError
from sqlalchemy.exc import SQLAlchemyError
from celery.exceptions import TimeoutError as CeleryTimeoutError
from common.extensions import db, metrics, minio_client
from common.utils.celery_utils import current_celery
from eveai_chat.socket_handlers.chat_handler import socketio_message_counter, socketio_message_latency

healthz_bp = Blueprint('healthz', __name__, url_prefix='/_healthz')


def liveness():
    try:
        # Basic check to see if the app is running
        return True
    except Exception:
        raise HealthError("Liveness check failed")


def readiness():
    checks = {
        "database": check_database(),
        "celery": check_celery(),
        # Add more checks as needed
    }

    if not all(checks.values()):
        raise HealthError("Readiness check failed")


def check_database():
    try:
        # Perform a simple database query
        db.session.execute("SELECT 1")
        return True
    except SQLAlchemyError:
        current_app.logger.error("Database check failed", exc_info=True)
        return False


def check_celery():
    try:
        # Send a simple task to Celery
        result = current_celery.send_task('ping', queue='llm_interactions')
        response = result.get(timeout=10)  # Wait for up to 10 seconds for a response
        return response == 'pong'
    except CeleryTimeoutError:
        current_app.logger.error("Celery check timed out", exc_info=True)
        return False
    except Exception as e:
        current_app.logger.error(f"Celery check failed: {str(e)}", exc_info=True)
        return False


@healthz_bp.route('/metrics')
@metrics.do_not_track()
def prometheus_metrics():
    return metrics.generate_latest()


def init_healtz(app):
    app.config.update(
        HEALTHZ={
            "live": "healthz_views.liveness",
            "ready": "healthz_views.readiness",
        }
    )
    # Register SocketIO metrics with Prometheus
    metrics.register(socketio_message_counter)
    metrics.register(socketio_message_latency)