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
from environment.model_inference import InferenceRequest, model_status, run_inference

_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,
)


@app.get("/model/status")
def get_model_status():
    return model_status()


@app.post("/model/infer")
def infer_with_trained_model(request: InferenceRequest):
    return run_inference(request.observation)

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