from __future__ import annotations

import sys
from pathlib import Path

for p in (Path("/kaggle/working/src"), Path(__file__).resolve().parent):
    if (p / "weather_forecast_env.py").exists():
        sys.path.insert(0, str(p))
        break

from zone_observation import ForecastConfig
from weather_forecast_env import WeatherForecastEnv


def _cfg(**overrides) -> ForecastConfig:
    base = dict(
        n_zones=2,
        horizon_days=15,
        max_steps=6,
        prior_belief=0.12,
        clean_episode_ratio=0.0,
        event_spatial_correlation=0.85,
        seed=7,
        real_data_ratio=0.0,
        real_data_pkl_path=None,
        shuffle_zone_order=False,
    )
    base.update(overrides)
    allowed = set(getattr(ForecastConfig, "__dataclass_fields__", {})) or set(base)
    return ForecastConfig(**{k: v for k, v in base.items() if k in allowed})


def _product_on_terminate(env, seed: int, inspect_all: bool) -> bool:
    _, info = env.reset(seed=seed)
    n = int(env.max_zones)
    if inspect_all:
        for a in range(n):
            _, _, terminated, _, info = env.step(a)
            assert not terminated, "budget exhausted before terminate; raise max_steps"
    _, _, terminated, _, info = env.step(env.terminate_action)
    assert terminated
    assert "product_actionable" in info
    return bool(info["product_actionable"])


def test_product_actionable_independent_of_visits() -> None:
    env = WeatherForecastEnv(_cfg())
    seed = 11
    flags_zero = _product_on_terminate(env, seed=seed, inspect_all=False)
    flags_full = _product_on_terminate(env, seed=seed, inspect_all=True)
    assert flags_zero == flags_full, (
        f"product_actionable changed after inspection "
        f"(zero-inspect={flags_zero} full-inspect={flags_full})"
    )


def test_product_flag_absent_on_inspect_step() -> None:
    env = WeatherForecastEnv(_cfg())
    _, info = env.reset(seed=3)
    assert "product_actionable" not in info
    _, _, terminated, _, info = env.step(0)
    assert not terminated
    assert "product_actionable" not in info


if __name__ == "__main__":
    test_product_flag_absent_on_inspect_step()
    print("  product_actionable absent on inspect step OK")
    test_product_actionable_independent_of_visits()
    print("  product_actionable independent of visits OK")
    print("All product-inspection contract tests passed.")