from __future__ import annotations

import sys
from pathlib import Path
from pprint import pprint

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

from bdo_ai_env.baseline_agent import greedy_policy
from models import BDOAction
from server.bdo_environment import BDOEnvironment


def main() -> None:
    env = BDOEnvironment(scenario="black_swan")
    observation = env.reset()
    print("Initial observation")
    pprint(observation.model_dump(mode="json", exclude_none=True))
    print("State:", env.state.model_dump(mode="json", exclude_none=True))
    print()

    month = 1
    while not observation.done:
        action_batch = greedy_policy(observation.model_dump(mode="json", exclude_none=True))
        observation = env.step(BDOAction.model_validate(action_batch))
        print(f"Month {month} reward={float(observation.reward or 0.0):.4f}")
        pprint(observation.info or observation.metadata)
        print("-" * 60)
        month += 1


if __name__ == "__main__":
    main()