soleprint/hub/dataloader/__init__.py

"""
Soleprint Data Loader

Loads JSON data files and provides typed access via Pydantic models.
JSON files live in data/ directory (content only, no code).
"""

import json
import os
import sys
from pathlib import Path
from typing import List, Optional

# When symlinked, __file__ resolves to actual location (hub/dataloader)
# but we need to find models/ in the runtime directory (gen/)
# Use cwd as the base since we always run from gen/
_runtime_dir = Path.cwd()
_file_parent = Path(__file__).resolve().parent.parent

# Try runtime dir first (gen/), then fall back to file's parent (hub/)
if (_runtime_dir / "models").exists():
    sys.path.insert(0, str(_runtime_dir))
else:
    sys.path.insert(0, str(_file_parent))

from models.pydantic import (
    Book,
    BookCollection,
    Depot,
    DepotCollection,
    Desk,
    DeskCollection,
    Pulse,
    PulseCollection,
    Room,
    RoomCollection,
    Status,
    Template,
    TemplateCollection,
    Tool,
    ToolCollection,
    Vein,
    VeinCollection,
)

# Data directory - try runtime dir first, then file's parent
_default_data = (
    _runtime_dir / "data" if (_runtime_dir / "data").exists() else _file_parent / "data"
)
DATA_DIR = Path(os.getenv("SOLEPRINT_DATA_DIR", _default_data)).resolve()


def _load_json(filename: str) -> dict:
    """Load a JSON file from the data directory."""
    filepath = DATA_DIR / filename
    if filepath.exists():
        with open(filepath) as f:
            return json.load(f)
    return {"items": []}


def _save_json(filename: str, data: dict):
    """Save data to a JSON file in the data directory."""
    filepath = DATA_DIR / filename
    with open(filepath, "w") as f:
        json.dump(data, f, indent=2)


# === Collection Loaders ===


def get_veins() -> List[Vein]:
    data = _load_json("veins.json")
    return VeinCollection(**data).items


def get_rooms() -> List[Room]:
    data = _load_json("rooms.json")
    return RoomCollection(**data).items


def get_depots() -> List[Depot]:
    data = _load_json("depots.json")
    return DepotCollection(**data).items


def get_templates() -> List[Template]:
    data = _load_json("templates.json")
    return TemplateCollection(**data).items


def get_tools() -> List[Tool]:
    data = _load_json("tools.json")
    return ToolCollection(**data).items


def get_cabinets() -> list:
    """Load cabinets (simple dict for now)."""
    data = _load_json("cabinets.json")
    return data.get("items", [])


def get_monitors() -> list:
    """Load monitors (simple dict for now)."""
    data = _load_json("monitors.json")
    return data.get("items", [])


def get_pulses() -> List[Pulse]:
    data = _load_json("pulses.json")
    return PulseCollection(**data).items


def get_books() -> List[Book]:
    data = _load_json("books.json")
    return BookCollection(**data).items


def get_desks() -> List[Desk]:
    data = _load_json("desks.json")
    return DeskCollection(**data).items


# === Single Item Helpers ===


def get_vein(name: str) -> Optional[Vein]:
    for v in get_veins():
        if v.name == name:
            return v
    return None


def get_room(name: str) -> Optional[Room]:
    for r in get_rooms():
        if r.name == name:
            return r
    return None


def get_depot(name: str) -> Optional[Depot]:
    for d in get_depots():
        if d.name == name:
            return d
    return None


def get_tool(name: str) -> Optional[Tool]:
    for t in get_tools():
        if t.name == name:
            return t
    return None


# === System Data (for frontend rendering) ===


def get_artery_data() -> dict:
    """Data for artery frontend."""
    return {
        "veins": [v.model_dump() for v in get_veins()],
        "rooms": [r.model_dump() for r in get_rooms()],
        "depots": [d.model_dump() for d in get_depots()],
        "pulses": [p.model_dump() for p in get_pulses()],
    }


def get_atlas_data() -> dict:
    """Data for atlas frontend."""
    return {
        "templates": [t.model_dump() for t in get_templates()],
        "depots": [d.model_dump() for d in get_depots()],
        "books": [b.model_dump() for b in get_books()],
    }


def get_station_data() -> dict:
    """Data for station frontend."""
    return {
        "tools": [t.model_dump() for t in get_tools()],
        "monitors": get_monitors(),
        "cabinets": get_cabinets(),
        "rooms": [r.model_dump() for r in get_rooms()],
        "depots": [d.model_dump() for d in get_depots()],
        "desks": [d.model_dump() for d in get_desks()],
    }