SmartEDT/backend/services/simulation_manager.py

from __future__ import annotations

import asyncio
import logging
import secrets
from dataclasses import dataclass
from typing import Any

from sqlalchemy import insert
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from backend.database.schema import Simulation, vehicle_signals
from backend.device.mock_vehicle import MockVehicleDevice
from backend.services.broadcaster import Broadcaster
from backend.utils import utc_now


logger = logging.getLogger("backend.simulation")


@dataclass
class SimulationRuntime:
    simulation_id: str
    status: str
    task: asyncio.Task | None = None


class SimulationManager:
    def __init__(
        self,
        session_factory: async_sessionmaker[AsyncSession],
        broadcaster: Broadcaster,
    ) -> None:
        self._session_factory = session_factory
        self._broadcaster = broadcaster
        self._runtime: SimulationRuntime | None = None
        self._device = MockVehicleDevice()
        self._seq = 0

    def current(self) -> SimulationRuntime | None:
        return self._runtime

    async def register_device(self, device: MockVehicleDevice) -> None:
        self._device = device

    async def start(self, scenario_config: dict[str, Any]) -> str:
        if self._runtime and self._runtime.status == "running":
            return self._runtime.simulation_id

        simulation_id = "SIM" + utc_now().strftime("%Y%m%d%H%M%S") + secrets.token_hex(2).upper()
        started_at = utc_now()
        scenario_name = scenario_config.get("scenario")
        operator = scenario_config.get("driver") or scenario_config.get("operator")
        config_created_at = started_at

        async with self._session_factory() as session:
            session.add(
                Simulation(
                    simulation_id=simulation_id,
                    status="running",
                    started_at=started_at,
                    ended_at=None,
                    scenario_name=scenario_name,
                    scenario_config=scenario_config,
                    config_created_at=config_created_at,
                    operator=operator,
                    archived=False,
                )
            )
            await session.commit()

        await self._device.connect()
        self._runtime = SimulationRuntime(simulation_id=simulation_id, status="running")
        self._runtime.task = asyncio.create_task(self._run_loop(simulation_id))
        await self._broadcaster.broadcast_json(
            {"type": "simulation.status", "ts": started_at.timestamp(), "simulation_id": simulation_id, "payload": {"status": "running"}}
        )
        return simulation_id

    async def stop(self, simulation_id: str) -> None:
        runtime = self._runtime
        if not runtime or runtime.simulation_id != simulation_id:
            return

        runtime.status = "stopping"
        if runtime.task:
            runtime.task.cancel()
            try:
                await runtime.task
            except asyncio.CancelledError:
                pass

        await self._device.disconnect()
        ended_at = utc_now()

        async with self._session_factory() as session:
            sim = await session.get(Simulation, simulation_id)
            if sim:
                sim.status = "stopped"
                sim.ended_at = ended_at
                await session.commit()

        await self._broadcaster.broadcast_json(
            {"type": "simulation.status", "ts": ended_at.timestamp(), "simulation_id": simulation_id, "payload": {"status": "stopped"}}
        )
        self._runtime = None

    async def _run_loop(self, simulation_id: str) -> None:
        try:
            while True:
                await asyncio.sleep(0.05)
                if not await self._device.is_connected():
                    continue

                self._seq += 1
                ts = utc_now()
                payload = self._device.sample().to_dict()
                message = {
                    "type": "vehicle.signal",
                    "ts": ts.timestamp(),
                    "simulation_id": simulation_id,
                    "device_id": self._device.device_id,
                    "seq": self._seq,
                    "payload": payload,
                }
                await self._broadcaster.broadcast_json(message)
                await self._persist_signal(ts, simulation_id, self._device.device_id, self._seq, payload)
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception("simulation loop crashed")

    async def _persist_signal(self, ts, simulation_id: str, device_id: str, seq: int, signals: dict[str, Any]) -> None:
        async with self._session_factory() as session:
            await session.execute(
                insert(vehicle_signals).values(
                    ts=ts,
                    simulation_id=simulation_id,
                    device_id=device_id,
                    seq=seq,
                    signals=signals,
                )
            )
            await session.commit()