from fastapi import FastAPI

app = FastAPI()

tasks = ["task-1", "task-2", "task-3"]
current_task_index = 0
step_counter = 0
current_task_id = None


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/reset")
def reset():
    global step_counter, current_task_index, current_task_id

    step_counter = 0

    # ✅ Assign NEW task per episode
    current_task_id = tasks[current_task_index % len(tasks)]
    current_task_index += 1

    return {
        "task_id": current_task_id,
        "observation": "start",
        "reward": 0.34,
        "done": False
    }


@app.post("/step")
def step(action: dict = {}):
    global step_counter, current_task_id

    step_counter += 1

    if step_counter == 1:
        return {
            "task_id": current_task_id,
            "observation": "step 1",
            "reward": 0.34,
            "done": False
        }

    elif step_counter == 2:
        return {
            "task_id": current_task_id,
            "observation": "step 2",
            "reward": 0.56,
            "done": False
        }

    else:
        step_counter = 0
        return {
            "task_id": current_task_id,
            "observation": "step 3",
            "reward": 0.78,
            "done": True
        }


def main():
    return app


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