"""Join the slim price estimates back onto properties.parquet.

Price estimation runs on ``price_inputs.parquet`` (built by ``property_base``
straight from epc_pp + arcgis, independently of merge's area features) and emits
``price_estimates.parquet``: the natural key (Postcode + coalesced address) plus
``Estimated current price`` / ``Est. price per sqm``. This step joins those two
columns onto properties.parquet to produce the file the server consumes.

Why the natural key
-------------------
Estimates and properties are built by separate runs, so a positional row index
would not line up. Instead both derive the key ``(Postcode, coalesce(register
address, EPC address))``, which is unique and non-null on the deduped dwelling
universe (see ``property_base._dedupe_collapsed_properties``) and identical on
both sides because both start from that same universe. So estimates map onto
properties 1:1 regardless of row order.

Re-running is safe: any pre-existing estimate columns are dropped first, and the
join is keyed (not positional), so a second run reproduces the same result. The
join refuses if any property has no estimate (the dwelling universes diverged,
e.g. a stale price_inputs vs a newer epc_pp) rather than silently leaving prices
null. Output is written to a temp file and atomically renamed.
"""

import argparse
from pathlib import Path

import polars as pl

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


def join_estimates(properties: Path, estimates_path: Path) -> int:
    """Augment ``properties`` in place with the estimate columns; return rows.

    Joins the slim estimates onto properties by the natural key and atomically
    replaces properties.parquet. Idempotent: any estimate columns already on the
    file are dropped first.
    """
    estimates = pl.scan_parquet(estimates_path)
    est_cols = estimates.collect_schema().names()
    missing = [c for c in (*JOIN_KEYS, *ESTIMATE_COLUMNS) if c not in est_cols]
    if missing:
        raise ValueError(f"{estimates_path}: missing columns {missing}")

    stats = estimates.select(
        n=pl.len(), unique=pl.struct(JOIN_KEYS).n_unique()
    ).collect(engine="streaming")
    n_estimates, n_unique = stats["n"][0], stats["unique"][0]
    if n_unique != n_estimates:
        raise ValueError(
            f"{estimates_path}: natural key {JOIN_KEYS} is not unique "
            f"({n_estimates - n_unique:,} duplicate rows)"
        )

    n_properties = pl.scan_parquet(properties).select(pl.len()).collect().item()

    # Drop any estimate columns already present (idempotent re-run) and attach the
    # coalesced-address half of the natural key.
    properties_keyed = (
        pl.scan_parquet(properties)
        .drop(ESTIMATE_COLUMNS, strict=False)
        .with_columns(join_address_expr())
    )

    # Every property must have an estimate: estimates and properties come from the
    # same dwelling universe, so a gap means a stale/foreign price_inputs (e.g.
    # built from a different epc_pp). Fail loudly instead of nulling prices.
    #
    # This assumes properties.parquet contains ONLY epc_pp-derived dwellings, which
    # is true for the production merge output. Running merge with --actual-listings
    # appends listing seed rows whose (Postcode, address) keys are absent from
    # price_inputs (built straight from epc_pp), which would trip the guard below.
    # Enabling listing integration on the primary output therefore requires
    # price_inputs to include those seed rows too.
    unmatched = (
        properties_keyed.select(JOIN_KEYS)
        .join(estimates.select(JOIN_KEYS), on=JOIN_KEYS, how="anti")
        .select(pl.len())
        .collect(engine="streaming")
        .item()
    )
    if unmatched:
        raise ValueError(
            f"{properties}: {unmatched:,} of {n_properties:,} properties have no "
            "matching estimate; the price_inputs and properties dwelling universes "
            "differ (regenerate price_inputs.parquet from the current epc_pp)."
        )

    # maintain_order="left" keeps properties in merge's row order; the unique key
    # cannot fan the join out, so the row count is preserved.
    result = properties_keyed.join(
        estimates, on=JOIN_KEYS, how="left", maintain_order="left"
    ).drop(JOIN_ADDRESS)

    tmp = properties.with_name(properties.name + ".tmp")
    result.sink_parquet(tmp)

    written = pl.scan_parquet(tmp).select(pl.len()).collect().item()
    if written != n_properties:
        tmp.unlink(missing_ok=True)
        raise ValueError(
            f"{properties}: join changed the row count "
            f"({n_properties:,} -> {written:,})"
        )

    tmp.replace(properties)
    return written


def main():
    parser = argparse.ArgumentParser(
        description="Join price_estimates.parquet onto properties.parquet"
    )
    parser.add_argument(
        "--properties",
        type=Path,
        required=True,
        help="properties.parquet (read, then overwritten with the estimate "
        "columns joined in)",
    )
    parser.add_argument(
        "--estimates",
        type=Path,
        required=True,
        help="Slim price_estimates.parquet from price_estimation.estimate",
    )
    args = parser.parse_args()

    written = join_estimates(args.properties, args.estimates)
    size_mb = args.properties.stat().st_size / (1024 * 1024)
    n_priced = (
        pl.scan_parquet(args.properties)
        .filter(pl.col("Estimated current price").is_not_null())
        .select(pl.len())
        .collect()
        .item()
    )
    print(f"Wrote {args.properties} ({size_mb:.1f} MB)")
    print(f"  {written:,} rows, {n_priced:,} with an estimated current price")


if __name__ == "__main__":
    main()