from ..base_gcp_agent import BaseGCPAgent
from ..gcp_connector import GCPConnector
from google.cloud import bigquery

class BigQueryAgent(BaseGCPAgent):
    """
    An agent specialized in handling GCP BigQuery tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to BigQuery.
        """
        if command == 'run_query':
            return self._run_query(**kwargs)
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by BigQueryAgent.")

    def _run_query(self, project_id: str, query_string: str):
        """
        Runs a SQL query in BigQuery and retrieves the results.
        """
        print(f"BigQueryAgent: Running query in project '{project_id}'...")
        try:
            client = bigquery.Client(credentials=GCPConnector.get_credentials(), project=project_id)

            query_job = client.query(query_string)
            results = query_job.result()

            rows = []
            for row in results:
                rows.append(dict(row))

            return {"status": "success", "message": "BigQuery query executed successfully.", "results": rows}
        except Exception as e:
            print(f"Error running BigQuery query: {e}")
            return {"status": "error", "message": str(e)}