"""Performance and connection tests for BigQueryClient."""

import time
import unittest
from unittest.mock import Mock, patch

import pandas as pd

from mlfastflow.bigqueryclient import BigQueryClient


class TestBigQueryClientPerformance(unittest.TestCase):
    """Tests for BigQueryClient connection management and performance."""

    @patch("mlfastflow.bigqueryclient.service_account")
    @patch("mlfastflow.bigqueryclient.bigquery")
    def _make_client(self, mock_bigquery, mock_service_account, **kwargs):
        """Helper: create a mocked BigQueryClient."""
        mock_creds = Mock()
        mock_creds.project_id = "test-project"
        mock_service_account.Credentials.from_service_account_file.return_value = (
            mock_creds
        )
        mock_bq = Mock()
        mock_bigquery.Client.return_value = mock_bq
        # Expose the mocks on the client for assertions
        client = BigQueryClient(
            project_id="test-project",
            dataset_id="test-dataset",
            key_file="/path/to/test-key.json",
            **kwargs,
        )
        client._mock_bigquery = mock_bigquery
        client._mock_bq = mock_bq
        return client

    @patch("mlfastflow.bigqueryclient.service_account")
    @patch("mlfastflow.bigqueryclient.bigquery")
    def test_connection_reuse_performance(self, mock_bigquery, mock_service_account):
        """Connections are reused across iterations and don't degrade over time."""
        mock_creds = Mock()
        mock_creds.project_id = "test-project"
        mock_service_account.Credentials.from_service_account_file.return_value = (
            mock_creds
        )

        mock_bq_client = Mock()
        mock_bigquery.Client.return_value = mock_bq_client

        mock_job = Mock()
        mock_df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
        mock_job.to_dataframe.return_value = mock_df
        mock_job.total_bytes_processed = 1024
        mock_bq_client.query.return_value = mock_job

        client = BigQueryClient(
            project_id="test-project",
            dataset_id="test-dataset",
            key_file="/path/to/test-key.json",
        )

        query_times = []
        num_iterations = 20

        for i in range(num_iterations):
            start = time.time()
            result = client.sql2df(f"SELECT * FROM test_table WHERE id = {i}")
            query_times.append(time.time() - start)
            self.assertIsInstance(result, pd.DataFrame)
            self.assertEqual(len(result), 3)

        self.assertEqual(mock_bq_client.query.call_count, num_iterations)
        self.assertEqual(client._query_count, num_iterations)

        # Later queries should not be >50% slower than early ones
        early_avg = sum(query_times[:5]) / 5
        late_avg = sum(query_times[-5:]) / 5
        if early_avg > 0:
            degradation = (late_avg - early_avg) / early_avg
            self.assertLess(
                degradation, 0.5,
                f"Performance degraded by {degradation * 100:.1f}%",
            )

    @patch("mlfastflow.bigqueryclient.service_account")
    @patch("mlfastflow.bigqueryclient.bigquery")
    @patch("mlfastflow.bigqueryclient.storage")
    def test_storage_client_caching(self, mock_storage, mock_bigquery, mock_service_account):
        """The storage client is created lazily and cached on first access."""
        mock_creds = Mock()
        mock_creds.project_id = "test-project"
        mock_service_account.Credentials.from_service_account_file.return_value = (
            mock_creds
        )
        mock_bigquery.Client.return_value = Mock()

        mock_storage_client = Mock()
        mock_storage.Client.return_value = mock_storage_client

        client = BigQueryClient(
            project_id="test-project",
            dataset_id="test-dataset",
            key_file="/path/to/test-key.json",
        )

        # Call via the public property and via the private helper — must be same object
        sc1 = client._get_storage_client()
        sc2 = client._get_storage_client()
        sc3 = client.storage_client

        self.assertIs(sc1, sc2)
        self.assertIs(sc2, sc3)
        # storage.Client() should have been called only once
        self.assertEqual(mock_storage.Client.call_count, 1)

    @patch("mlfastflow.bigqueryclient.service_account")
    @patch("mlfastflow.bigqueryclient.bigquery")
    def test_connection_refresh_mechanism(self, mock_bigquery, mock_service_account):
        """Connection is refreshed when query count hits the threshold."""
        mock_creds = Mock()
        mock_creds.project_id = "test-project"
        mock_service_account.Credentials.from_service_account_file.return_value = (
            mock_creds
        )

        mock_bq_client = Mock()
        mock_bigquery.Client.return_value = mock_bq_client

        client = BigQueryClient(
            project_id="test-project",
            dataset_id="test-dataset",
            key_file="/path/to/test-key.json",
        )

        initial_calls = mock_bigquery.Client.call_count
        self.assertGreaterEqual(initial_calls, 1)

        # Prime query count above the refresh threshold
        client._query_count = BigQueryClient._MAX_QUERIES_BEFORE_REFRESH

        mock_job = Mock()
        mock_job.to_dataframe.return_value = pd.DataFrame({"test": [1]})
        mock_job.total_bytes_processed = 0
        mock_bq_client.query.return_value = mock_job

        client.sql2df("SELECT 1 AS test")

        # A new bigquery.Client should have been created during refresh
        self.assertGreater(mock_bigquery.Client.call_count, initial_calls)
        # Counter resets to 0 in _refresh_connection, then +1 in _check_and_refresh
        self.assertEqual(client._query_count, 1)

    def test_memory_cleanup(self):
        """All resources are cleared after close()."""
        with (
            patch("mlfastflow.bigqueryclient.service_account") as mock_sa,
            patch("mlfastflow.bigqueryclient.bigquery") as mock_bq,
        ):
            mock_creds = Mock()
            mock_creds.project_id = "test-project"
            mock_sa.Credentials.from_service_account_file.return_value = mock_creds
            mock_bq.Client.return_value = Mock()

            client = BigQueryClient(
                project_id="test-project",
                dataset_id="test-dataset",
                key_file="/path/to/test-key.json",
            )

            self.assertIsNotNone(client.client)
            self.assertIsNotNone(client.credentials)
            self.assertIsNotNone(client._connection_created_at)

            client.close()

            self.assertIsNone(client.client)
            self.assertIsNone(client.credentials)
            self.assertIsNone(client._storage_client)
            self.assertIsNone(client._connection_created_at)
            self.assertEqual(client._query_count, 0)


if __name__ == "__main__":
    unittest.main()