#!/usr/bin/env python
#
# tyr/tools/query_tyr_audit.py - MCP Tool for querying audit logs
#

from google.cloud import bigquery


def query_tyr_audit(p: dict) -> dict:
    """Executes a read-only SQL query against the TYR audit log dataset."""
    query = p.get("query")
    if not query:
        raise ValueError("Missing required parameter: 'query'")

    # Initialize the BigQuery client
    client = bigquery.Client()

    # Construct the full table name (assuming standard log sink naming)
    # This will need to be adjusted with the actual table name once logs are flowing.
    table_id = "propane-will-491900-m5.tyr_audit_logs.cloudaudit_googleapis_com_activity"

    # For security, ensure the query is a SELECT statement
    if not query.strip().upper().startswith("SELECT"):
        raise ValueError(
            "Security violation: Only SELECT queries are allowed.")

    # Construct the full query
    full_query = query.replace("FROM activity", f"FROM `{table_id}`")

    try:
        query_job = client.query(full_query)
        results = query_job.result()  # Waits for the job to complete

        # Convert rows to a list of dictionaries
        rows = [dict(row) for row in results]

        return {"status": "success", "row_count": len(rows), "rows": rows}
    except Exception as e:
        print(f"An error occurred: {e}")
        raise