from sqlalchemy import create_engine

from database.manager import Base


AI_OBSERVABILITY_TABLES = {
    "ai_calls": {
        "caller",
        "provider",
        "model",
        "status",
        "request_id",
        "meta",
    },
    "mcp_calls": {
        "caller",
        "server",
        "tool",
        "status",
        "request_id",
        "input_args",
    },
    "ai_call_budgets": {
        "period",
        "provider",
        "budget_usd",
        "alert_pct",
    },
    "rag_query_log": {
        "caller",
        "query_text",
        "query_embedding",
        "embedding_signature",
        "used_results",
        "saved_call",
    },
    "learning_episodes": {
        "episode_type",
        "source_table",
        "distilled_text",
        "embedding",
        "embedding_signature",
        "promotion_status",
    },
    "host_health_probes": {
        "host_label",
        "host_url",
        "healthy",
        "unhealthy_mark",
    },
    "ppt_audit_results": {
        "pptx_filename",
        "vision_enabled",
        "audit_status",
        "issues_found",
    },
}


def test_ai_observability_tables_registered_in_metadata():
    metadata_tables = set(Base.metadata.tables)

    assert AI_OBSERVABILITY_TABLES.keys() <= metadata_tables
    for table_name, expected_columns in AI_OBSERVABILITY_TABLES.items():
        actual_columns = set(Base.metadata.tables[table_name].columns.keys())
        assert expected_columns <= actual_columns


def test_ai_observability_tables_compile_for_sqlite_metadata_init():
    engine = create_engine("sqlite:///:memory:")

    for table_name in AI_OBSERVABILITY_TABLES:
        Base.metadata.tables[table_name].create(engine, checkfirst=True)