app/session_store.py

from datetime import datetime, timedelta
from typing import Dict, List, Optional

from app.models import SessionState, WorkerState, DevicePolicy
from app.security import generate_session_id, generate_session_secret
from app.storage import save_session_event

_sessions_by_id: Dict[str, SessionState] = {}
_workers_by_session: Dict[str, Dict[str, WorkerState]] = {}


def create_session(owner_label: Optional[str] = None) -> SessionState:
    from app.config import Settings
    session_id = generate_session_id()
    secret = generate_session_secret()
    ttl = Settings.get_session_ttl_seconds()
    now = datetime.utcnow()
    session = SessionState(
        session_id=session_id,
        owner_label=owner_label,
        secret=secret,
        created_at=now,
        expires_at=now + timedelta(seconds=ttl),
        policy=DevicePolicy(session_id=session_id),
    )
    _sessions_by_id[session_id] = session
    _workers_by_session[session_id] = {}
    save_session_event({
        "event": "created",
        "session_id": session_id,
        "owner_label": owner_label,
        "timestamp": now.isoformat(),
    })
    return session


def get_session(session_id: str) -> Optional[SessionState]:
    expire_old_sessions()
    return _sessions_by_id.get(session_id)


def delete_session(session_id: str) -> bool:
    if session_id in _sessions_by_id:
        del _sessions_by_id[session_id]
        _workers_by_session.pop(session_id, None)
        return True
    return False


def expire_old_sessions() -> None:
    now = datetime.utcnow()
    expired = [sid for sid, s in _sessions_by_id.items() if s.expires_at < now]
    for sid in expired:
        delete_session(sid)


def attach_worker(session_id: str, worker: WorkerState) -> bool:
    session = get_session(session_id)
    if not session:
        return False
    if session_id not in _workers_by_session:
        _workers_by_session[session_id] = {}
    _workers_by_session[session_id][worker.worker_id] = worker
    session.workers[worker.worker_id] = worker
    return True


def detach_worker(session_id: str, worker_id: str) -> bool:
    session = get_session(session_id)
    if session and worker_id in session.workers:
        session.workers.pop(worker_id, None)
    ws = _workers_by_session.get(session_id)
    if ws and worker_id in ws:
        del ws[worker_id]
        return True
    return False


def list_workers(session_id: str) -> List[WorkerState]:
    return list(_workers_by_session.get(session_id, {}).values())


def update_worker_heartbeat(session_id: str, worker_id: str) -> bool:
    ws = _workers_by_session.get(session_id, {})
    worker = ws.get(worker_id)
    if worker:
        worker.last_heartbeat = datetime.utcnow()
        return True
    return False


def set_session_policy(session_id: str, policy: DevicePolicy) -> bool:
    session = get_session(session_id)
    if session:
        session.policy = policy
        return True
    return False


def get_session_policy(session_id: str) -> Optional[DevicePolicy]:
    session = get_session(session_id)
    return session.policy if session else None