"""Unit tests for GCP parquet partition-scoping helpers."""

from pathlib import Path
from unittest.mock import Mock, patch

import pytest

from src.services.gcp_parquet import (
    _cache_dir_for_partition,
    _ensure_partition_cache,
    _extract_single_date_partition,
    _is_cache_ready,
    _strip_partition_date_filter,
    read_gcp_delta_parquet,
)
from src.utils.errors import DeltaError


class TestExtractSingleDatePartition:
    def test_accepts_compact_date(self):
        day = _extract_single_date_partition("SELECT * FROM __t__ WHERE date='20260318'")
        assert day == "20260318"

    def test_rejects_dash_date(self):
        with pytest.raises(DeltaError):
            _extract_single_date_partition("SELECT * FROM __t__ WHERE date='2026-03-18'")

    def test_requires_date_filter(self):
        with pytest.raises(DeltaError):
            _extract_single_date_partition("SELECT * FROM __t__ LIMIT 10")

    def test_rejects_multiple_dates(self):
        with pytest.raises(DeltaError):
            _extract_single_date_partition(
                "SELECT * FROM __t__ WHERE date='20260318' OR date='20260319'"
            )


class TestStripPartitionDateFilter:
    def test_removes_date_term_and_keeps_other_predicates(self):
        q = (
            "SELECT query_time FROM __t__ "
            "WHERE date='20260318' AND query_time > 1 "
            "ORDER BY query_time DESC LIMIT 10"
        )
        out = _strip_partition_date_filter(q)
        assert "date='20260318'" not in out
        assert "query_time > 1" in out
        assert "ORDER BY query_time DESC" in out
        assert "LIMIT 10" in out

    def test_removes_where_when_only_date_term(self):
        q = "SELECT * FROM __t__ WHERE date='20260318' LIMIT 10"
        out = _strip_partition_date_filter(q)
        assert "WHERE" not in out.upper()
        assert out == "SELECT * FROM __t__ LIMIT 10"


class TestPartitionCache:
    def test_first_download_writes_cache(self, monkeypatch):
        monkeypatch.setenv("O11Y_GCP_PARQUET_CACHE_DIR", "/tmp/o11y-mcp-test-cache-a")
        cache_dir = _cache_dir_for_partition("b1", "deltalake/tidb/slowlogs/", "20260318")
        if cache_dir.exists():
            for f in cache_dir.glob("*"):
                f.unlink()
            cache_dir.rmdir()

        with patch(
            "src.services.gcp_parquet._list_parquet_objects",
            return_value=["deltalake/tidb/slowlogs/date=20260318/p1.parquet"],
        ):
            with patch("src.services.gcp_parquet._download_gcs_object") as mock_download:
                def _fake_download(_client, _bucket, _name, _token, dest_file: Path):
                    dest_file.write_bytes(b"PAR1")

                mock_download.side_effect = _fake_download
                out = _ensure_partition_cache(
                    client=Mock(),
                    bucket="b1",
                    prefix="deltalake/tidb/slowlogs/",
                    day="20260318",
                    access_token="ya29.token",
                )
        assert out == cache_dir
        assert _is_cache_ready(cache_dir)
        assert (cache_dir / "_cache_manifest.json").exists()

    def test_cache_hit_skips_download(self, monkeypatch):
        monkeypatch.setenv("O11Y_GCP_PARQUET_CACHE_DIR", "/tmp/o11y-mcp-test-cache-b")
        cache_dir = _cache_dir_for_partition("b1", "deltalake/tidb/slowlogs/", "20260318")
        cache_dir.mkdir(parents=True, exist_ok=True)
        (cache_dir / "part_0.parquet").write_bytes(b"PAR1")
        (cache_dir / "_cache_manifest.json").write_text(
            '{"bucket":"b1","prefix":"deltalake/tidb/slowlogs/","date":"20260318","parquet_count":1}',
            encoding="utf-8",
        )

        with patch("src.services.gcp_parquet._list_parquet_objects") as mock_list:
            with patch("src.services.gcp_parquet._download_gcs_object") as mock_download:
                out = _ensure_partition_cache(
                    client=Mock(),
                    bucket="b1",
                    prefix="deltalake/tidb/slowlogs/",
                    day="20260318",
                    access_token="ya29.token",
                )
        assert out == cache_dir
        mock_list.assert_not_called()
        mock_download.assert_not_called()


class TestReadGcpParquet:
    def test_builds_distinct_view_for_local_query(self):
        mock_client = Mock()
        mock_cm = Mock()
        mock_cm.__enter__ = Mock(return_value=mock_client)
        mock_cm.__exit__ = Mock(return_value=None)

        mock_df = Mock()
        mock_df.empty = False
        mock_df.to_dict.return_value = [{"x": 1}]
        mock_exec_result = Mock()
        mock_exec_result.fetchdf.return_value = mock_df
        mock_conn = Mock()
        mock_conn.execute.return_value = mock_exec_result

        with patch("src.services.gcp_parquet.httpx.Client", return_value=mock_cm):
            with patch(
                "src.services.gcp_parquet._ensure_partition_cache",
                return_value=Path("/tmp/o11y-gcp-cache"),
            ):
                with patch("src.services.gcp_parquet._is_cache_ready", return_value=True):
                    with patch("src.services.gcp_parquet.duckdb.connect", return_value=mock_conn):
                        read_gcp_delta_parquet(
                            "gs://bucket/deltalake/t1/c1/slowlogs",
                            "ya29.token",
                            "SELECT * FROM __t__ WHERE date='20260318' LIMIT 1",
                            1,
                        )

        create_view_calls = [
            c[0][0]
            for c in mock_conn.execute.call_args_list
            if "CREATE VIEW __t__" in str(c[0][0])
        ]
        assert create_view_calls
        assert "SELECT DISTINCT * FROM read_parquet" in create_view_calls[0]