"""Tests for joining slim price estimates back onto properties.parquet.

estimate.py emits (Postcode, coalesced address, estimate columns) and
join_estimates attaches them by that natural key. These tests pin the
properties that make the key safe: it maps estimates onto the right rows
regardless of order (a shuffled estimates frame is the worst case), it is
idempotent, and it refuses a partial/foreign estimates file rather than
silently nulling prices.
"""

from pathlib import Path

import polars as pl
import pytest

from pipeline.transform.join_price_estimates import join_estimates
from pipeline.transform.price_estimation.utils import (
    ESTIMATE_COLUMNS,
    JOIN_ADDRESS,
    JOIN_KEYS,
    join_address_expr,
)

N = 200


def _write_merged(path: Path) -> pl.DataFrame:
    """properties.parquet with the natural-key columns, a sentinel order column,
    and no estimates. Half the rows are sale-addressed, half EPC-only, so the
    coalesce in the key is exercised; every coalesced address is unique."""
    df = pl.DataFrame(
        {
            "Postcode": [f"AA{i % 7} {i % 9}AA" for i in range(N)],
            "Address per Property Register": [
                f"reg-{i}" if i % 2 == 0 else None for i in range(N)
            ],
            "Address per EPC": [f"epc-{i}" if i % 2 == 1 else None for i in range(N)],
            "order": list(range(N)),
            "junk": [f"x{i}" for i in range(N)],
        }
    )
    df.write_parquet(path)
    return df


def _write_estimates(path: Path, merged_path: Path, *, shuffle: bool = True) -> None:
    """Estimates keyed by the natural key, derived from the merged file the way
    estimate.py does. Estimate = order * 1000 so each row is checkable. Shuffled
    by default to prove order-independence."""
    est = (
        pl.read_parquet(merged_path)
        .with_columns(join_address_expr())
        .with_columns(
            (pl.col("order") * 1000).cast(pl.Float64).alias("Estimated current price"),
            (pl.col("order") * 10).cast(pl.Int32).alias("Est. price per sqm"),
        )
        .select(*JOIN_KEYS, *ESTIMATE_COLUMNS)
    )
    if shuffle:
        est = est.sample(fraction=1.0, shuffle=True, seed=7)
    est.write_parquet(path)


def test_join_attaches_estimates_to_the_right_rows(tmp_path: Path):
    props = tmp_path / "properties.parquet"
    estimates = tmp_path / "price_estimates.parquet"
    _write_merged(props)
    _write_estimates(estimates, props)

    written = join_estimates(props, estimates)
    out = pl.read_parquet(props)

    assert written == N
    assert out.height == N
    # Order preserved and the address-half of the key is not left behind.
    assert out["order"].to_list() == list(range(N))
    assert out["junk"].to_list() == [f"x{i}" for i in range(N)]
    assert JOIN_ADDRESS not in out.columns
    # Every row carries its own estimate, matched by key despite the shuffle.
    assert out["Estimated current price"].to_list() == [float(i * 1000) for i in range(N)]
    assert out["Est. price per sqm"].to_list() == [i * 10 for i in range(N)]
    assert out["Estimated current price"].null_count() == 0


def test_rerun_is_idempotent(tmp_path: Path):
    props = tmp_path / "properties.parquet"
    estimates = tmp_path / "price_estimates.parquet"
    _write_merged(props)
    _write_estimates(estimates, props)

    join_estimates(props, estimates)
    first = pl.read_parquet(props)
    join_estimates(props, estimates)  # second run on the augmented file
    second = pl.read_parquet(props)

    assert second.equals(first)
    assert second.columns.count("Estimated current price") == 1
    assert second.columns.count("Est. price per sqm") == 1


def test_missing_estimate_is_rejected(tmp_path: Path):
    """A property with no matching estimate (diverged dwelling universe) must
    fail loudly rather than silently leave its price null."""
    props = tmp_path / "properties.parquet"
    estimates = tmp_path / "price_estimates.parquet"
    _write_merged(props)
    _write_estimates(estimates, props)
    # Drop one estimate so a property key is no longer covered.
    pl.read_parquet(estimates).head(N - 1).write_parquet(estimates)

    with pytest.raises(ValueError, match="no matching estimate"):
        join_estimates(props, estimates)


def test_duplicate_key_is_rejected(tmp_path: Path):
    props = tmp_path / "properties.parquet"
    estimates = tmp_path / "price_estimates.parquet"
    _write_merged(props)
    _write_estimates(estimates, props)
    # Force row 1's key to collide with row 0's.
    est = pl.read_parquet(estimates).sort("Estimated current price")
    row0 = est.row(0, named=True)
    est = est.with_columns(
        pl.when(pl.int_range(pl.len()) == 1)
        .then(pl.lit(row0["Postcode"]))
        .otherwise(pl.col("Postcode"))
        .alias("Postcode"),
        pl.when(pl.int_range(pl.len()) == 1)
        .then(pl.lit(row0[JOIN_ADDRESS]))
        .otherwise(pl.col(JOIN_ADDRESS))
        .alias(JOIN_ADDRESS),
    )
    est.write_parquet(estimates)

    with pytest.raises(ValueError, match="not unique"):
        join_estimates(props, estimates)