import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from typing import Optional
from env.tasks import load_task

app = FastAPI()
state = {"current_task": None}

class Action(BaseModel):
    action_type: str
    content: Optional[str] = ""

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

@app.post("/reset")
async def reset(task: str = "easy_refund"):
    global state
    task_data = load_task(task)[0]
    state["current_task"] = task_data
    return {"observation": task_data["text"], "task_id": task}

@app.post("/step")
async def step(action: Action):
    global state
    t = state.get("current_task")
    if not t: return {"reward": 0.05, "done": True, "error": "No task"}
    
    # Logic: If agent mentions the keyword, give partial reward
    success = t["target"] in action.content.lower() or action.action_type == "solve"
    reward = t["reward_weight"] if success else 0.15
    return {"observation": "Processed", "reward": reward, "done": True, "error": None}

def main():
    uvicorn.run(app, host="0.0.0.0", port=7860)

if __name__ == "__main__":
    main()