# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""`RunSpec.data_dir` (the prospective bundle mechanism): the three parquet loads and the geo cache are redirected to
the given directory; the default layout (data/processed | data/synthetic) is untouched. Synthetic data, mock provider."""
from __future__ import annotations

import json
from pathlib import Path

import polars as pl
import pytest

from agentic_valuation.config import PROCESSED, SYNTHETIC
from agentic_valuation.experiments.run_arm import RunSpec, data_base, feature_backend, load_analytic, run_arm, select_subsample


def _bundle(tmp_path: Path, syn_analytic: pl.DataFrame, syn_splits: pl.DataFrame, syn_subsamples: pl.DataFrame) -> tuple[Path, list[str]]:
    """A bundle whose 'prospective' rows are the 6 latest synthetic sales relabelled (the history stays as it is)."""
    d = tmp_path / "bundle"
    d.mkdir()
    pros = syn_analytic.sort(["date", "tx_id"]).tail(6)["tx_id"].to_list()
    splits = syn_splits.with_columns(pl.when(pl.col("tx_id").is_in(pros)).then(pl.lit("prospective")).otherwise(pl.col("split")).alias("split"))
    subs = pl.concat([
        syn_subsamples.filter(~pl.col("tx_id").is_in(pros)),
        pl.DataFrame({"tx_id": sorted(pros), "split": ["prospective"] * 6, "stratum": ["s"] * 6, "draw_seed": [1] * 6}, schema=syn_subsamples.schema),
    ])
    syn_analytic.write_parquet(d / "analytic.parquet")
    splits.write_parquet(d / "splits.parquet")
    subs.write_parquet(d / "llm_subsamples.parquet")
    pl.DataFrame({"tx_id": pros, "flood_mapped": [True] * 6, "in_flood_zone": [False] * 6, "flood_zone_kind": [None] * 6, "transit_stops_within_1km": [3] * 6,
                  "dauid": ["x"] * 6, "da_population": [500] * 6, "da_population_density": [1000.0] * 6, "da_median_age": [40.0] * 6, "da_median_income": [50000.0] * 6,
                  "da_pct_renters": [30.0] * 6, "cbd_name": ["Montréal"] * 6, "distance_to_cbd_km": [5.0] * 6},
                 schema_overrides={"flood_zone_kind": pl.String, "transit_stops_within_1km": pl.Int32}).write_parquet(d / "geo_features.parquet")
    return d, sorted(pros)


def test_defaults_untouched():
    assert RunSpec().data_dir is None
    assert data_base(False) == PROCESSED and data_base(True) == SYNTHETIC
    assert data_base(False, None) == PROCESSED and data_base(True, None) == SYNTHETIC
    assert data_base(True, Path("/x/y")) == Path("/x/y") and data_base(False, "/x/y") == Path("/x/y")


def test_data_dir_redirects_loads(tmp_path: Path, syn_analytic, syn_splits, syn_subsamples):
    d, pros = _bundle(tmp_path, syn_analytic, syn_splits, syn_subsamples)
    df = load_analytic(False, d)                      # synthetic=False: no fallback to the generator, the directory is authoritative
    assert df.height == syn_analytic.height and set(df.filter(pl.col("split") == "prospective")["tx_id"]) == set(pros)
    sample = select_subsample(df, "prospective", 4, 1, False, d)
    assert sample["tx_id"].to_list() == pros[:4] and set(sample["split"]) == {"prospective"}
    fb = feature_backend(False, d)
    assert fb.geo_cache is not None and fb.geo_cache.height == 6
    with pytest.raises(FileNotFoundError):
        load_analytic(False, tmp_path / "missing")


def test_run_arm_with_data_dir(tmp_path: Path, monkeypatch, syn_analytic, syn_splits, syn_subsamples):
    monkeypatch.setenv("AV_PROVIDER_MODE", "mock")
    d, pros = _bundle(tmp_path, syn_analytic, syn_splits, syn_subsamples)
    spec = RunSpec(arm="E_single_llm", committee="mixed_random", split="prospective", n=3, synthetic=False, out_dir=tmp_path / "runs", run_id="t_datadir", data_dir=d)
    run_dir = run_arm(spec)
    preds = pl.read_parquet(run_dir / "predictions.parquet")
    assert sorted(preds["tx_id"].unique().to_list()) == pros[:3] and set(preds["split"]) == {"prospective"}
    m = json.loads((run_dir / "manifest.json").read_text())
    assert m["data_dir"] == str(d) and m["data_source"] == str(d / "analytic.parquet") and m["geo_cache"] == str(d / "geo_features.parquet")
    assert m["split"] == "prospective" and m["n_sampled"] == 3 and m["n_failed_this_session"] == 0 and m["spec"]["data_dir"] == str(d)
    pk = json.loads((run_dir / "packets" / f"{pros[0]}.json").read_text())
    assert "price" not in pk["subject"]


def test_run_arm_default_manifest_unchanged(tmp_path: Path, monkeypatch):
    monkeypatch.setenv("AV_PROVIDER_MODE", "mock")
    spec = RunSpec(arm="E_single_llm", committee="all_anthropic", split="validation", n=2, synthetic=True, out_dir=tmp_path, run_id="t_default")
    m = json.loads((run_arm(spec) / "manifest.json").read_text())
    assert m["data_dir"] is None and m["data_source"].startswith("data/synthetic/") and m["spec"]["data_dir"] is None