from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any, Dict

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from models import EraseOpsAction
from server.environment import EraseOpsEnvironment
from server.graders import grade


def _assert(condition: bool, message: str) -> None:
    if not condition:
        raise AssertionError(message)


def validate_task(env: EraseOpsEnvironment, task_name: str) -> Dict[str, Any]:
    observation = env.reset(task_name=task_name, seed=42)
    _assert(observation.step == 0, f"{task_name}: reset step must be 0")
    _assert(len(observation.records) > 0, f"{task_name}: expected records")
    _assert(observation.step_budget > 0, f"{task_name}: invalid step budget")

    seen_any = False
    for record in observation.records:
        action = EraseOpsAction(action_type="inspect", record_id=record.record_id, justification="validate")
        observation, reward, done, info = env.step(action)
        seen_any = True
        _assert(isinstance(reward.total, float), f"{task_name}: reward must be numeric")
        _assert("task_name" in info, f"{task_name}: info should include task_name")
        if done:
            break

    _assert(seen_any, f"{task_name}: no records processed")
    final_state = env.state()
    score = grade(task_name, final_state)
    _assert(0.01 <= score <= 0.99, f"{task_name}: grade outside expected range")
    return {"task": task_name, "score": score, "final_step": final_state.get("step")}


def main() -> None:
    env = EraseOpsEnvironment()
    summary = [
        validate_task(env, task_name)
        for task_name in ["clean_sweep", "tangled_lake", "legal_minefield", "cascade_collapse"]
    ]
    print(json.dumps({"status": "ok", "summary": summary}, indent=2))


if __name__ == "__main__":
    main()