r"""
CivicAI OpenEnv Compliance Validation Script
Run: venv/Scripts/python.exe validate_openenv.py
"""
import sys

print("=== Import Check ===")

from civicai.models import Action, Observation, Reward, SocietyState, SubsidyPolicy
print("[OK] civicai.models: Action, Observation, Reward, SocietyState")

from civicai.environment import CivicAIEnv
print("[OK] civicai.environment: CivicAIEnv")

from openenv.env import Env
assert issubclass(CivicAIEnv, Env), "CivicAIEnv must inherit from openenv.env.Env"
print("[OK] CivicAIEnv inherits from openenv.env.Env")

from pydantic import BaseModel
assert issubclass(Action, BaseModel), "Action must be Pydantic"
assert issubclass(Observation, BaseModel), "Observation must be Pydantic"
assert issubclass(Reward, BaseModel), "Reward must be Pydantic"
print("[OK] Action, Observation, Reward are Pydantic BaseModels")

print()
print("=== OpenEnv API Compliance ===")
env = CivicAIEnv()

# reset() -> Observation
obs = env.reset()
assert isinstance(obs, Observation), f"reset() must return Observation, got {type(obs)}"
print(f"[OK] reset() -> Observation  (turn={obs.turn}, task={obs.task_id})")

# state() -> SocietyState  (must be callable method, NOT a property)
assert callable(getattr(env, "state")), "state must be a callable method, not a property"
st = env.state()
assert isinstance(st, SocietyState), f"state() must return SocietyState, got {type(st)}"
print(f"[OK] state() -> SocietyState  (callable method, turn={st.turn})")

# step(action) -> (Observation, float, bool, dict)
action = Action()
result = env.step(action)
assert len(result) == 4, f"step() must return 4-tuple, got {len(result)}"
obs2, reward, done, info = result
assert isinstance(obs2, Observation), f"step()[0] must be Observation, got {type(obs2)}"
assert isinstance(reward, float), f"step()[1] must be float, got {type(reward)}"
assert isinstance(done, bool), f"step()[2] must be bool, got {type(done)}"
assert isinstance(info, dict), f"step()[3] must be dict, got {type(info)}"
assert 0.0 <= reward <= 1.0, f"reward must be in [0,1], got {reward}"
print(f"[OK] step(action) -> (Observation, float, bool, dict)  reward={reward:.4f}")

print()
print("=== Task Tests ===")
for task_id in ["stabilize_economy", "manage_pandemic", "control_crisis"]:
    obs = env.reset(task_id=task_id)
    assert isinstance(obs, Observation)
    obs2, r, done, info_ = env.step(Action())
    assert 0.0 <= r <= 1.0
    print(f"[OK] task={task_id}  initial_reward={r:.4f}")

print()
print("=== Reward Model ===")
from civicai.reward import compute_reward
obs = env.reset()
env.step(Action())
st = env.state()
reward_obj = compute_reward(st, Action())
rd = reward_obj.model_dump()
assert "score" in rd and "rubrics" in rd and "penalties" in rd
rubric_keys = set(rd["rubrics"].keys())
assert rubric_keys == {"economic", "health", "social", "sustainability", "crime"}, \
    f"Unexpected rubric keys: {rubric_keys}"
print(f"[OK] Reward.score={reward_obj.score:.4f}  rubrics={sorted(rubric_keys)}")

print()
print("=== openenv.yaml Validation ===")
import yaml, os
yaml_path = "openenv.yaml"
assert os.path.exists(yaml_path), "openenv.yaml not found"
with open(yaml_path) as f:
    cfg = yaml.safe_load(f)

required_top_keys = ["name", "description", "observation_space", "action_space", "reward_range", "tasks"]
for k in required_top_keys:
    assert k in cfg, f"openenv.yaml missing required key: {k}"
    print(f"[OK] openenv.yaml has '{k}'")

assert len(cfg["tasks"]) >= 3, f"Need >= 3 tasks, found {len(cfg['tasks'])}"
print(f"[OK] openenv.yaml has {len(cfg['tasks'])} tasks (>= 3 required)")

for task in cfg["tasks"]:
    for field in ["id", "name", "description", "success_criteria", "max_steps"]:
        assert field in task, f"Task '{task.get('id', '?')}' missing field: {field}"
print(f"[OK] All task entries have required fields")

assert isinstance(cfg["reward_range"], list) and len(cfg["reward_range"]) == 2
print(f"[OK] reward_range={cfg['reward_range']}")

print()
print("=" * 55)
print("  ALL CHECKS PASSED")
print("  CivicAI is fully OpenEnv compliant.")
print("=" * 55)