setting tasks for ships

5 files changed, 56 insertions, 12 deletions

diff --git a/__main__.py b/__main__.py
index 2434422..36cbc28 100644
--- a/__main__.py
+++ b/__main__.py
@@ -3,7 +3,7 @@ from typing import Annotated
 
 from fastapi import FastAPI, Depends
 
-from .apirouters import agents
+from .apirouters import agents, tasks
 from .modules.database import cursor, conn
 
 
@@ -20,8 +20,5 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 app.include_router(agents.router)
+app.include_router(tasks.router)
 
-
-@app.get("/tasks")
-async def get_tasks(agent: Annotated[agents.Agent, Depends(agents.auth_agent)]):
-    return {"callsign": agent.agent_symbol}
diff --git a/apirouters/agents.py b/apirouters/agents.py
index 55825e6..a1e015b 100644
--- a/apirouters/agents.py
+++ b/apirouters/agents.py
@@ -1,13 +1,13 @@
 from enum import Enum
-from typing import Annotated, Dict, List
+from typing import Annotated, Dict
 
 import openapi_client
 from fastapi import APIRouter, Depends, Response, HTTPException
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
 
-from ..modules import tasks
-from ..modules.ships import Ship
+from ..modules import task_type
 from ..modules.database import cursor, conn
+from ..modules.ships import Ship
 
 router = APIRouter()
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
diff --git a/apirouters/tasks.py b/apirouters/tasks.py
new file mode 100644
index 0000000..806de91
--- /dev/null
+++ b/apirouters/tasks.py
@@ -0,0 +1,42 @@
+from typing import Annotated
+from fastapi import APIRouter, Depends
+
+from ..modules.database import cursor, conn
+from ..modules.ships import Ship
+from ..modules import task_type
+from .agents import auth_agent, Agent
+
+router = APIRouter()
+
+
+@router.get("/tasks")
+async def get_tasks(agent: Annotated[Agent, Depends(auth_agent)]):
+    ships = []
+    for ship in agent.ships.values():
+        ship.load_task()
+        ships.append({"symbol": ship.symbol, "task": ship.task})
+
+    return {"ships": ships}
+
+
+@router.get("/task/{ship_symbol}")
+async def get_tasks(ship_symbol: str, agent: Annotated[Agent, Depends(auth_agent)]):
+    for current_ship in agent.ships.values():
+        if current_ship.symbol == ship_symbol:
+            current_ship.load_task()
+            return {"symbol": current_ship.symbol, "task": current_ship.task}
+
+    return {"error": "Unknown ship symbol."}
+
+
+@router.post("/task/{ship_symbol}/set/{task_type}")
+async def get_tasks(ship_symbol: str, task: str, agent: Annotated[Agent, Depends(auth_agent)]):
+    if task not in task_type.task_types:
+        return {"error": "Invalid task."}
+
+    for current_ship in agent.ships.values():
+        if current_ship.symbol == ship_symbol:
+            current_ship.set_task(task)
+            return {"symbol": current_ship.symbol, "task": current_ship.task}
+
+    return {"error": "Unknown ship symbol."}
diff --git a/modules/ships.py b/modules/ships.py
index 0a296ca..8dfe41a 100644
--- a/modules/ships.py
+++ b/modules/ships.py
@@ -1,4 +1,4 @@
-from . import tasks
+from . import task_type
 from .database import cursor, conn
 
 
@@ -13,15 +13,15 @@ class Ship:
         cursor.execute("SELECT task, params FROM ships WHERE symbol = ?", (self.symbol,))
         row = cursor.fetchone()
         if row is None:
-            cursor.execute("INSERT INTO ships (symbol, task, params) VALUES (?, ?, ?)", (self.symbol, tasks.IDLE, None))
+            cursor.execute("INSERT INTO ships (symbol, task, params) VALUES (?, ?, ?)", (self.symbol, task_type.IDLE, None))
             conn.commit()
 
-            self.task = tasks.IDLE
+            self.task = task_type.IDLE
         else:
             self.task = row[0]
 
     def set_task(self, task):
         self.task = task
-        if task != tasks.ERROR:
+        if task != task_type.ERROR:
             cursor.execute("UPDATE ships SET task = ? WHERE symbol = ?", (task, self.symbol))
             conn.commit()
diff --git a/modules/tasks.py b/modules/task_type.py
index 8479134..7579ad0 100644
--- a/modules/tasks.py
+++ b/modules/task_type.py
@@ -3,3 +3,8 @@ MINING = 'MINING'
 
 MIA = 'MIA'
 ERROR = 'ERROR'
+
+task_types = [
+    IDLE,
+    MINING,
+]