import os
from dotenv import load_dotenv
from openenv.core import create_app
import uvicorn

load_dotenv()

from environment.actions import ContextCorruptionAction, EpisodeObservation
from environment.env import ContextCorruptionEnv

_difficulty_env = os.getenv("DIFFICULTY")
_difficulty = int(_difficulty_env) if _difficulty_env else None
_max_sessions = int(os.getenv("MAX_CONCURRENT_ENVS", "64"))

app = create_app(
    env=lambda: ContextCorruptionEnv(difficulty=_difficulty),
    action_cls=ContextCorruptionAction,
    observation_cls=EpisodeObservation,
    env_name="ContextCorruption-Env",
    max_concurrent_envs=_max_sessions,
)

if __name__ == "__main__":
    uvicorn.run("environment.server:app", host="0.0.0.0", port=7860, reload=False)