ship data initialization on startup

4 files changed, 65 insertions, 2 deletions

diff --git a/__main__.py b/__main__.py
index 9b23484..03fb366 100644
--- a/__main__.py
+++ b/__main__.py
@@ -11,6 +11,8 @@ from .modules.database import cursor, conn
 async def lifespan(app: FastAPI):
     cursor.execute(
         "CREATE TABLE IF NOT EXISTS agents(primary_key INTEGER PRIMARY KEY, callsign TEXT NOT NULL UNIQUE, token TEXT NOT NULL)")
+    cursor.execute(
+        "CREATE TABLE IF NOT EXISTS ships(primary_key INTEGER PRIMARY KEY, symbol TEXT NOT NULL UNIQUE, task TEXT NOT NULL, params TEXT)")
     conn.commit()
     agents.load_agents_from_database()
     yield
diff --git a/apirouters/agents.py b/apirouters/agents.py
index 89f736f..4ac2167 100644
--- a/apirouters/agents.py
+++ b/apirouters/agents.py
@@ -1,15 +1,17 @@
 from enum import Enum
-from typing import Annotated
+from typing import Annotated, Dict, List
 
 import openapi_client
 from fastapi import APIRouter, Depends, Response
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
 
+from ..modules import tasks
+from ..modules.ships import Ship
 from ..modules.database import cursor, conn
 
 router = APIRouter()
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
-agents = {}
+agents: Dict[str, 'Agent'] = {}
 
 
 class Agent:
@@ -22,6 +24,7 @@ class Agent:
         self.systems_api = openapi_client.SystemsApi(self.api_client)
 
         self.agent_symbol = agent_symbol
+        self.ships: Dict[str, Ship] = {}
 
 
 @router.post("/token")
@@ -65,6 +68,32 @@ def load_agents_from_database() -> None:
             cursor.execute("DELETE from agents WHERE callsign = ?", (invalid,))
         conn.commit()
 
+    for agent in agents.values():
+        try:
+            response = agent.fleet_api.get_my_ships()
+            for ship_data in response.data:
+                agent.ships[ship_data.symbol] = Ship(ship_data.symbol)
+
+            for ship in agent.ships.values():
+                ship.load_task()
+        except:
+            for ship in agent.ships.values():
+                ship.set_task(tasks.ERROR)
+
+        cursor.execute("SELECT symbol FROM ships WHERE symbol LIKE ?", (f"{agent.agent_symbol}-%",))
+        ship_names = agent.ships.keys()
+        for row in cursor.fetchall():
+            if row[0] in ship_names:
+                continue
+
+            missing_ship = Ship(row[0])
+            missing_ship.set_task(tasks.MIA)
+            agent.ships[row[0]] = missing_ship
+
+
+
+
+
 
 @router.post("/{callsign}/init", status_code=201)
 async def init_agent(callsign: str, token: Annotated[str, Depends(oauth2_scheme)], response: Response):
diff --git a/modules/ships.py b/modules/ships.py
new file mode 100644
index 0000000..0a296ca
--- /dev/null
+++ b/modules/ships.py
@@ -0,0 +1,27 @@
+from . import tasks
+from .database import cursor, conn
+
+
+class Ship:
+    def __init__(self, symbol: str):
+        self.symbol = symbol
+        self.callsign = "-".split(symbol)[0]
+
+        self.task = None
+
+    def load_task(self):
+        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))
+            conn.commit()
+
+            self.task = tasks.IDLE
+        else:
+            self.task = row[0]
+
+    def set_task(self, task):
+        self.task = task
+        if task != tasks.ERROR:
+            cursor.execute("UPDATE ships SET task = ? WHERE symbol = ?", (task, self.symbol))
+            conn.commit()
diff --git a/modules/tasks.py b/modules/tasks.py
new file mode 100644
index 0000000..8479134
--- /dev/null
+++ b/modules/tasks.py
@@ -0,0 +1,5 @@
+IDLE = 'IDLE'
+MINING = 'MINING'
+
+MIA = 'MIA'
+ERROR = 'ERROR'