from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from env import TrafficEnv
from tasks import get_config
from baseline_agent import RuleBasedAgent
import os

app = FastAPI()
env = TrafficEnv(get_config("medium"))
agent = RuleBasedAgent()

class Action(BaseModel):
    action: int

@app.get("/", response_class=HTMLResponse)
def root():
    with open("index.html", "r", encoding="utf-8") as f:
        return f.read()

@app.post("/reset")
def reset():
    state = env.reset()
    try:
        state = state.tolist()
    except:
        pass
    agent.reset()
    return {"state":state}

@app.post("/step")
def step(data:Action):
    state,reward,done,info = env.step(data.action)
    try:
        state = state.tolist()
    except:
        pass
    return {
        "state":state,
        "reward":reward,
        "done":done,
        "info":info
    }

@app.post("/auto_step")
def auto_step():
    state_dict = env.get_state()
    action = agent.select_action(state_dict)
    state,reward,done,info = env.step(action)
    try:
        state = state.tolist()
    except:
        pass
    return {
        "state":state,
        "reward":reward,
        "done":done,
        "info":info,
        "action_taken": action
    }