spoolman2/spoolman/externaldb.py

"""Functions for syncing data from an external database of manufacturers, filaments, materials, etc."""

import datetime
import logging
import os
from collections.abc import Iterator
from enum import Enum
from pathlib import Path
from typing import Any, Optional
from urllib.parse import urljoin

import hishel
from pydantic import BaseModel, Field, RootModel
from scheduler.asyncio.scheduler import Scheduler

from spoolman import filecache
from spoolman.env import get_cache_dir

logger = logging.getLogger(__name__)


DEFAULT_EXTERNAL_DB_URL = "https://donkie.github.io/SpoolmanDB/"
DEFAULT_3DFP_URL = "https://raw.githubusercontent.com/jklewa/filament-profiles-data/main/"
DEFAULT_SYNC_INTERVAL = 3600


class ExternalSource(Enum):
    """External database sources."""

    SPOOLMANDB = "spoolmandb"
    FILAMENT_PROFILES_3D = "3dfilamentprofiles"

controller = hishel.Controller(allow_stale=True)
try:
    cache_path = get_cache_dir() / "hishel"
    cache_storage = hishel.AsyncFileStorage(base_path=cache_path)
except PermissionError:
    logger.warning(
        "Failed to setup disk-based cache due to permission error. Ensure the path %s is writable. "
        "Using in-memory cache instead as fallback.",
        str(cache_path.resolve()),
    )
    cache_storage = hishel.AsyncInMemoryStorage()


class SpoolType(Enum):
    PLASTIC = "plastic"
    CARDBOARD = "cardboard"
    METAL = "metal"


class Finish(Enum):
    MATTE = "matte"
    GLOSSY = "glossy"


class MultiColorDirection(Enum):
    COAXIAL = "coaxial"
    LONGITUDINAL = "longitudinal"


class Pattern(Enum):
    MARBLE = "marble"
    SPARKLE = "sparkle"


class ExternalFilament(BaseModel):
    id: str = Field(description="A unique ID for this filament.", examples=["polymaker_pla_polysonicblack_1000_175"])
    manufacturer: str = Field(description="Filament manufacturer.", examples=["Polymaker"])
    name: str = Field(description="Filament name.", examples=["Polysonic\u2122 Black"])
    material: str = Field(description="Filament material.", examples=["PLA"])
    density: float = Field(description="Density in g/cm3.", examples=[1.23])
    weight: Optional[float] = Field(default=None, description="Net weight of a single spool.", examples=[1000])
    spool_weight: Optional[float] = Field(default=None, description="Weight of an empty spool.", examples=[140])
    spool_type: Optional[SpoolType] = Field(None, description="Type of spool.", examples=[SpoolType.PLASTIC])
    diameter: float = Field(description="Filament in mm.", examples=[1.75])
    color_hex: Optional[str] = Field(
        default=None,
        description="Filament color code in hex format, for single-color filaments.",
        examples=["2c3232"],
    )
    color_hexes: Optional[list[str]] = Field(
        default=None,
        description="For multi-color filaments. List of hex color codes in hex format.",
        examples=[["2c3232", "5f5f5f"]],
    )
    extruder_temp: Optional[int] = Field(default=None, description="Extruder/nozzle temperature in °C.", examples=[210])
    extruder_temp_max: Optional[int] = Field(default=None, description="Max extruder temp in °C.", examples=[220])
    bed_temp: Optional[int] = Field(default=None, description="Bed temperature in °C.", examples=[50])
    bed_temp_max: Optional[int] = Field(default=None, description="Max bed temp in °C.", examples=[60])
    finish: Optional[Finish] = Field(default=None, description="Finish of the filament.", examples=[Finish.MATTE])
    multi_color_direction: Optional[MultiColorDirection] = Field(
        default=None,
        description="Direction of multi-color filaments.",
        examples=[MultiColorDirection.COAXIAL],
    )
    pattern: Optional[Pattern] = Field(default=None, description="Pattern of the filament.", examples=[Pattern.MARBLE])
    translucent: bool = Field(default=False, description="Whether the filament is translucent.")
    glow: bool = Field(default=False, description="Whether the filament is glow-in-the-dark.")
    source: ExternalSource = Field(
        default=ExternalSource.SPOOLMANDB,
        description="External database source this filament came from.",
    )


class ExternalFilamentsFile(RootModel):
    root: list[ExternalFilament]

    def __iter__(self) -> Iterator[ExternalFilament]:
        """Iterate over the filaments."""
        return iter(self.root)

    def __getitem__(self, index: int) -> ExternalFilament:
        """Get a specific filament by index."""
        return self.root[index]


class ExternalMaterial(BaseModel):
    material: str = Field(examples=["PLA"])
    density: float = Field(examples=[1.24])
    extruder_temp: Optional[int] = Field(default=None, description="Extruder/nozzle temperature in °C.", examples=[210])
    bed_temp: Optional[int] = Field(default=None, description="Bed temperature in °C.", examples=[50])


class ExternalMaterialsFile(RootModel):
    root: list[ExternalMaterial]

    def __iter__(self) -> Iterator[ExternalMaterial]:
        """Iterate over the materials."""
        return iter(self.root)

    def __getitem__(self, index: int) -> ExternalMaterial:
        """Get a specific material by index."""
        return self.root[index]


# 3dfilamentprofiles.com data models
class FilamentProperties3DFP(BaseModel):
    """Properties object from 3dfilamentprofiles."""

    temp_min: Optional[int] = None
    temp_max: Optional[int] = None
    bed_temp_min: Optional[int] = None
    bed_temp_max: Optional[int] = None
    spool_weight: Optional[float] = None
    flow_ratio: Optional[float] = None


class Filament3DFP(BaseModel):
    """Raw filament data from 3dfilamentprofiles.com."""

    id: int
    brand_name: str
    material: str
    material_type: Optional[str] = None
    color: str
    rgb: Optional[str] = None
    properties: Optional[dict[str, Any]] = None


class FilamentsFile3DFP(BaseModel):
    """The filaments.json file from 3dfilamentprofiles."""

    filaments: list[Filament3DFP]


# Material density defaults (used when 3dfilamentprofiles doesn't provide density)
MATERIAL_DENSITIES: dict[str, float] = {
    "PLA": 1.24,
    "ABS": 1.04,
    "PETG": 1.27,
    "NYLON": 1.52,
    "TPU": 1.21,
    "PC": 1.3,
    "WOOD": 1.28,
    "CF": 1.3,
    "PC/ABS": 1.19,
    "HIPS": 1.03,
    "PVA": 1.23,
    "ASA": 1.05,
    "PP": 0.9,
    "POM": 1.4,
    "PMMA": 1.18,
    "FPE": 2.16,
}


def _get_density_for_material(material: str) -> float:
    """Get density for a material, defaulting to PLA's density if unknown."""
    material_upper = material.upper()
    for key, density in MATERIAL_DENSITIES.items():
        if key in material_upper:
            return density
    return 1.24  # Default to PLA density


def _transform_3dfp_to_external(filament: Filament3DFP) -> ExternalFilament:
    """Transform a 3dfilamentprofiles filament to ExternalFilament format."""
    # Build name from material_type and color
    name_parts = []
    if filament.material_type:
        name_parts.append(filament.material_type)
    name_parts.append(filament.color)
    name = " ".join(name_parts)

    # Handle color hex (3dfp includes # prefix, SpoolmanDB doesn't)
    color_hex = None
    if filament.rgb:
        color_hex = filament.rgb.lstrip("#").lower()

    # Extract properties
    props = filament.properties or {}
    extruder_temp = props.get("temp_min")
    extruder_temp_max = props.get("temp_max")
    bed_temp = props.get("bed_temp_min")
    bed_temp_max = props.get("bed_temp_max")
    spool_weight = props.get("spool_weight")

    # Determine finish from material_type
    finish = None
    if filament.material_type:
        material_type_lower = filament.material_type.lower()
        if "matte" in material_type_lower:
            finish = Finish.MATTE
        elif "glossy" in material_type_lower or "silk" in material_type_lower:
            finish = Finish.GLOSSY

    # Determine pattern from material_type
    pattern = None
    if filament.material_type:
        material_type_lower = filament.material_type.lower()
        if "marble" in material_type_lower:
            pattern = Pattern.MARBLE
        elif "sparkle" in material_type_lower or "glitter" in material_type_lower:
            pattern = Pattern.SPARKLE

    return ExternalFilament(
        id=f"3dfp_{filament.id}",
        manufacturer=filament.brand_name,
        name=name,
        material=filament.material,
        density=_get_density_for_material(filament.material),
        weight=None,  # 3dfp doesn't provide filament weight
        spool_weight=spool_weight,
        spool_type=None,
        diameter=1.75,  # Default, 3dfp doesn't provide diameter
        color_hex=color_hex,
        color_hexes=None,
        extruder_temp=extruder_temp,
        extruder_temp_max=extruder_temp_max,
        bed_temp=bed_temp,
        bed_temp_max=bed_temp_max,
        finish=finish,
        multi_color_direction=None,
        pattern=pattern,
        translucent=False,
        glow=False,
        source=ExternalSource.FILAMENT_PROFILES_3D,
    )


def get_external_db_url() -> str:
    """Get the external database URL from environment variables. Defaults to DEFAULT_EXTERNAL_DB_URL."""
    return os.getenv("EXTERNAL_DB_URL", DEFAULT_EXTERNAL_DB_URL)


def get_3dfp_url() -> str:
    """Get the 3dfilamentprofiles URL from environment variables."""
    return os.getenv("EXTERNAL_DB_3DFP_URL", DEFAULT_3DFP_URL)


def is_3dfp_enabled() -> bool:
    """Check if 3dfilamentprofiles sync is enabled."""
    return os.getenv("EXTERNAL_DB_3DFP_ENABLED", "true").lower() in ("true", "1", "yes")


def get_external_db_sync_interval() -> int:
    """Get the external database sync interval from environment variables. Defaults to DEFAULT_SYNC_INTERVAL."""
    return int(os.getenv("EXTERNAL_DB_SYNC_INTERVAL", DEFAULT_SYNC_INTERVAL))


async def _download_file(url: str) -> bytes:
    """Download a file from a URL and return the contents as a string.

    Uses a file-based cache.
    """
    async with hishel.AsyncCacheClient(storage=cache_storage, controller=controller) as client:
        response = await client.get(url)
        response.raise_for_status()
        return response.read()


def _parse_filaments_from_bytes(data: bytes) -> ExternalFilamentsFile:
    """Parse a bytes string into a list of ExternalFilament objects."""
    return ExternalFilamentsFile.parse_raw(data)


def _parse_materials_from_bytes(data: bytes) -> ExternalMaterialsFile:
    """Parse a bytes string into a list of ExternalMaterial objects."""
    return ExternalMaterialsFile.parse_raw(data)


def _write_to_local_cache(filename: str, data: bytes) -> None:
    """Write data to the local cache."""
    filecache.update_file(filename, data)


def get_filaments_file() -> Path:
    """Get the path to the filaments file."""
    return filecache.get_file("filaments.json")


def get_materials_file() -> Path:
    """Get the path to the materials file."""
    return filecache.get_file("materials.json")


def _parse_3dfp_filaments(data: bytes) -> list[Filament3DFP]:
    """Parse 3dfilamentprofiles JSON data."""
    try:
        parsed = FilamentsFile3DFP.model_validate_json(data)
        return parsed.filaments
    except Exception:
        # Fallback: try parsing as raw list
        import json

        raw = json.loads(data)
        if isinstance(raw, dict) and "filaments" in raw:
            return [Filament3DFP.model_validate(f) for f in raw["filaments"]]
        return []


async def _sync() -> None:
    logger.info("Syncing external DB.")

    all_filaments: list[ExternalFilament] = []

    # Sync SpoolmanDB
    spoolmandb_url = get_external_db_url()
    if spoolmandb_url.strip():
        try:
            spoolmandb_filaments = _parse_filaments_from_bytes(
                await _download_file(urljoin(spoolmandb_url, "filaments.json"))
            )
            # Mark all SpoolmanDB filaments with source
            for filament in spoolmandb_filaments.root:
                filament.source = ExternalSource.SPOOLMANDB
            all_filaments.extend(spoolmandb_filaments.root)
            logger.info("SpoolmanDB synced: %d filaments", len(spoolmandb_filaments.root))
        except Exception as e:
            logger.warning("Failed to sync SpoolmanDB: %s", e)

    # Sync 3dfilamentprofiles
    if is_3dfp_enabled():
        url_3dfp = get_3dfp_url()
        if url_3dfp.strip():
            try:
                raw_data = await _download_file(urljoin(url_3dfp, "sample-filaments.json"))
                filaments_3dfp = _parse_3dfp_filaments(raw_data)
                transformed = [_transform_3dfp_to_external(f) for f in filaments_3dfp]
                all_filaments.extend(transformed)
                logger.info("3dfilamentprofiles synced: %d filaments", len(transformed))
            except Exception as e:
                logger.warning("Failed to sync 3dfilamentprofiles: %s", e)

    # Write combined filaments
    combined = ExternalFilamentsFile(root=all_filaments)
    _write_to_local_cache("filaments.json", combined.model_dump_json().encode())

    # Materials only from SpoolmanDB
    if spoolmandb_url.strip():
        try:
            materials = _parse_materials_from_bytes(
                await _download_file(urljoin(spoolmandb_url, "materials.json"))
            )
            _write_to_local_cache("materials.json", materials.model_dump_json().encode())
            logger.info("Materials synced: %d", len(materials.root))
        except Exception as e:
            logger.warning("Failed to sync materials: %s", e)

    logger.info("External DB sync complete. Total filaments: %d", len(all_filaments))


def schedule_tasks(scheduler: Scheduler) -> None:
    """Schedule tasks to be executed by the provided scheduler.

    Args:
        scheduler: The scheduler to use for scheduling tasks.

    """
    if len(get_external_db_url().strip()) == 0:
        logger.info("External DB URL is empty. Skipping sync.")
        return

    logger.info("Scheduling external DB sync.")

    # Run once on startup
    scheduler.once(datetime.timedelta(seconds=0), _sync)  # type: ignore[arg-type]

    sync_interval = get_external_db_sync_interval()
    if sync_interval > 0:
        scheduler.cyclic(datetime.timedelta(seconds=DEFAULT_SYNC_INTERVAL), _sync)  # type: ignore[arg-type]
    else:
        logger.info("Sync interval is 0, skipping periodic sync of external db.")