"""Tests for parquet null preservation.

Verifies that nullable columns in parquet files produce NaN/None in DataFrames,
not zero. Regression test for the null-to-zero bug where the Go API wrote
non-nullable parquet schemas, causing null values to be silently coerced to 0.
"""

from __future__ import annotations

from pathlib import Path

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import pytest


def _write_observation_parquet(path: Path) -> None:
    """Write a parquet file mimicking the fixed Go API output.

    Some columns are optional (nullable) — wind_gust_kt, sky_base_2_ft, etc.
    Required columns (station_code, temp_f in row 0) are non-null.
    """
    schema = pa.schema([
        pa.field("station_code", pa.binary(), nullable=False),
        pa.field("observed_at", pa.binary(), nullable=False),
        pa.field("temp_f", pa.float64(), nullable=True),
        pa.field("wind_speed_kt", pa.int64(), nullable=True),
        pa.field("wind_gust_kt", pa.int64(), nullable=True),
        pa.field("sky_base_1_ft", pa.int64(), nullable=True),
        pa.field("sky_base_2_ft", pa.int64(), nullable=True),
        pa.field("visibility_miles", pa.float64(), nullable=True),
        pa.field("feels_like_f", pa.float64(), nullable=True),
    ])

    table = pa.table(
        {
            "station_code": pa.array([b"DCA", b"DCA", b"DCA"], type=pa.binary()),
            "observed_at": pa.array(
                [b"2026-03-30T12:00:00Z", b"2026-03-30T13:00:00Z", b"2026-03-30T14:00:00Z"],
                type=pa.binary(),
            ),
            "temp_f": pa.array([55.0, None, 60.0], type=pa.float64()),
            "wind_speed_kt": pa.array([10, 15, None], type=pa.int64()),
            "wind_gust_kt": pa.array([None, 25, None], type=pa.int64()),
            "sky_base_1_ft": pa.array([12000, None, 8000], type=pa.int64()),
            "sky_base_2_ft": pa.array([None, None, None], type=pa.int64()),
            "visibility_miles": pa.array([10.0, None, 10.0], type=pa.float64()),
            "feels_like_f": pa.array([None, None, 58.0], type=pa.float64()),
        },
        schema=schema,
    )
    pq.write_table(table, path)


def _load_parquet_as_dataframe(path: Path) -> pd.DataFrame:
    """Load parquet using the same code path as therminal SDK's _http.py."""
    _str_types = {pa.binary(), pa.large_binary(), pa.utf8(), pa.large_utf8()}
    return pq.read_table(path).to_pandas(
        types_mapper=lambda arrow_type: pd.StringDtype() if arrow_type in _str_types else None,
    )


class TestParquetNullPreservation:
    """Regression tests: nullable parquet columns must produce NaN, not zero."""

    @pytest.fixture()
    def df(self, tmp_path: Path) -> pd.DataFrame:
        path = tmp_path / "observations.parquet"
        _write_observation_parquet(path)
        return _load_parquet_as_dataframe(path)

    def test_null_float_is_nan_not_zero(self, df: pd.DataFrame) -> None:
        """temp_f null in row 1 must be NaN, not 0.0."""
        assert pd.isna(df["temp_f"].iloc[1]), (
            f"Expected NaN for null temp_f, got {df['temp_f'].iloc[1]}"
        )

    def test_null_int_is_nan_not_zero(self, df: pd.DataFrame) -> None:
        """wind_gust_kt null in row 0 must be NaN, not 0."""
        assert pd.isna(df["wind_gust_kt"].iloc[0]), (
            f"Expected NaN for null wind_gust_kt, got {df['wind_gust_kt'].iloc[0]}"
        )

    def test_present_values_preserved(self, df: pd.DataFrame) -> None:
        """Non-null values must be correct."""
        assert df["temp_f"].iloc[0] == 55.0
        assert df["temp_f"].iloc[2] == 60.0
        assert df["wind_gust_kt"].iloc[1] == 25.0  # int → float due to NaN coexistence

    def test_all_null_column_is_all_nan(self, df: pd.DataFrame) -> None:
        """sky_base_2_ft is null in every row — must be all NaN, not all 0."""
        assert df["sky_base_2_ft"].isna().all(), (
            f"Expected all NaN for all-null column, got: {df['sky_base_2_ft'].tolist()}"
        )

    def test_no_spurious_zeros_in_nullable_columns(self, df: pd.DataFrame) -> None:
        """Nullable columns should have NaN where null, never zero (unless genuinely zero)."""
        # wind_gust_kt: [None, 25, None] — no real zeros
        gust = df["wind_gust_kt"]
        assert (gust == 0).sum() == 0, (
            f"wind_gust_kt has {(gust == 0).sum()} zeros — expected only NaN or 25"
        )
        # feels_like_f: [None, None, 58.0] — no real zeros
        feels = df["feels_like_f"]
        assert (feels == 0).sum() == 0, (
            f"feels_like_f has {(feels == 0).sum()} zeros — expected only NaN or 58.0"
        )

    def test_null_vs_zero_distinguishable(self, df: pd.DataFrame) -> None:
        """Confirm we can distinguish null from a real zero value.

        This is the core invariant broken by the old bug: null became 0,
        making it impossible to tell 'no reading' from '0 knots wind gust'.
        """
        # wind_speed_kt row 2 is null, rows 0/1 are 10/15
        assert pd.isna(df["wind_speed_kt"].iloc[2])
        assert df["wind_speed_kt"].iloc[0] == 10.0
        assert df["wind_speed_kt"].iloc[1] == 15.0


class TestOldBrokenParquet:
    """Demonstrates the bug: non-nullable schema coerces null → zero."""

    @pytest.fixture()
    def broken_df(self, tmp_path: Path) -> pd.DataFrame:
        """Write a parquet with the OLD (broken) non-nullable schema."""
        schema = pa.schema([
            pa.field("station_code", pa.binary(), nullable=False),
            pa.field("temp_f", pa.float64(), nullable=False),
            pa.field("wind_gust_kt", pa.int64(), nullable=False),
        ])
        # Non-nullable columns can't hold None — pyarrow refuses.
        # Simulate what the old Go writer did: write 0 instead of null.
        table = pa.table(
            {
                "station_code": pa.array([b"DCA", b"DCA"], type=pa.binary()),
                "temp_f": pa.array([55.0, 0.0], type=pa.float64()),  # 0.0 instead of null
                "wind_gust_kt": pa.array([0, 25], type=pa.int64()),  # 0 instead of null
            },
            schema=schema,
        )
        path = tmp_path / "broken.parquet"
        pq.write_table(table, path)
        return _load_parquet_as_dataframe(path)

    def test_broken_schema_has_zero_not_nan(self, broken_df: pd.DataFrame) -> None:
        """Old bug: null values became 0 — indistinguishable from real zeros."""
        # This test documents the bug. In the broken schema:
        assert broken_df["wind_gust_kt"].iloc[0] == 0  # Should be NaN!
        assert broken_df["temp_f"].iloc[1] == 0.0  # Should be NaN!
        # No NaN values exist — that's the bug.
        assert not broken_df["wind_gust_kt"].isna().any()
        assert not broken_df["temp_f"].isna().any()
