SatisfactoryPlusCalculator/factorygame/data/models.py

import datetime
import re
from typing import Optional

from sqlalchemy import String, Select, select, Table, Column, ForeignKey
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship

amount_per_step_regex = re.compile(r"^\d*(\.\d+)?")
time_per_step_regex = re.compile(r"^(\d+(\.\d+)?) seconds?")


class Base(DeclarativeBase):
    pass


ingredients_table = Table(
    "recipe_ingredients",
    Base.metadata,
    Column("recipe_id", ForeignKey("recipes.id"), primary_key=True),
    Column("resource_flow_id", ForeignKey("resource_flows.id"), primary_key=True),
)
results_table = Table(
    "recipe_results",
    Base.metadata,
    Column("recipe_id", ForeignKey("recipes.id"), primary_key=True),
    Column("resource_flow_id", ForeignKey("resource_flows.id"), primary_key=True),
)
recipe_factories_table = Table(
    "recipe_factories",
    Base.metadata,
    Column("recipe_id", ForeignKey("recipes.id"), primary_key=True),
    Column("factory_id", ForeignKey("factories.id"), primary_key=True),
)


class Resource(Base):
    __tablename__ = "resources"

    id: Mapped[int] = mapped_column(primary_key=True)
    label: Mapped[str] = mapped_column(String(127))
    uri: Mapped[str]
    flows: Mapped[list["ResourceFlow"]] = relationship(back_populates="resource")
    recipes_populated_at: Mapped[Optional[datetime.datetime]]

    @classmethod
    def by_label(cls, search: str) -> Select[tuple["Resource"]]:
        return select(cls).where(cls.label.ilike(search))

    def __repr__(self):
        return (
            f"Resource(id={self.id}, "
            f"label={self.label}, "
            f"uri={self.uri}, "
            f"recipes_populated_at={self.recipes_populated_at})"
        )


class Factory(Base):
    __tablename__ = "factories"

    id: Mapped[int] = mapped_column(primary_key=True)
    label: Mapped[str] = mapped_column(String(127))
    uri: Mapped[str]
    used_in: Mapped[list["Recipe"]] = relationship(secondary=recipe_factories_table, viewonly=True)

    @classmethod
    def by_label(cls, search: str) -> Select[tuple["Factory"]]:
        return select(cls).where(cls.label.ilike(search))

    def __repr__(self):
        return f"Factory(id={self.id}, label={self.label}, uri={self.uri})"


class ResourceFlow(Base):
    __tablename__ = "resource_flows"

    id: Mapped[int] = mapped_column(primary_key=True)
    ingredient_in: Mapped[Optional["Recipe"]] = relationship(secondary=ingredients_table, back_populates="ingredients")
    result_of: Mapped[Optional["Recipe"]] = relationship(secondary=results_table, back_populates="results")
    resource_id: Mapped[int] = mapped_column(ForeignKey("resources.id"))
    resource: Mapped["Resource"] = relationship(back_populates="flows")
    amount: Mapped[str]
    time: Mapped[str]

    def amount_per_minute(self) -> float:
        amount_per_step_match = amount_per_step_regex.match(self.amount)
        assert amount_per_step_match, (
            f"Amount per crafting step of resource {self.resource.label} in crafting recipe "
            f"{self.result_of or self.ingredient_in} could not be parsed"
        )
        amount_per_step = float(amount_per_step_match.group(0))

        time_per_step_match = time_per_step_regex.match(self.time)
        assert time_per_step_match, (
            f"Time per crafting step of resource {self.resource.label} in crafting recipe "
            f"{self.result_of or self.ingredient_in} could not be parsed"
        )
        time_per_step = float(time_per_step_match.group(1))
        return amount_per_step * (60.0 / time_per_step)

    def describe(self) -> str:
        return f"{repr(str(self.resource.label))}x{self.amount}"

    def __repr__(self):
        return f"ResourceFlow(id={self.id}, resource_id={self.resource_id}, amount={self.amount}, time={self.time})"


class Recipe(Base):
    __tablename__ = "recipes"

    id: Mapped[int] = mapped_column(primary_key=True)
    factories: Mapped[list["Factory"]] = relationship(secondary=recipe_factories_table, back_populates="used_in")
    ingredients: Mapped[list["ResourceFlow"]] = relationship(
        secondary=ingredients_table, back_populates="ingredient_in"
    )
    results: Mapped[list["ResourceFlow"]] = relationship(secondary=results_table, back_populates="result_of")

    def join_factories(self):
        return ", ".join(map(lambda factory: factory.label, self.factories))

    def describe(self) -> str:
        def list_flows(flows: list["ResourceFlow"]) -> str:
            return ", ".join(map(ResourceFlow.describe, flows))

        return (
            f"in machine(s): {self.join_factories()}, "
            f"ingredient(s): {list_flows(self.ingredients)}, "
            f"result(s): {list_flows(self.results)}"
        )

    def __repr__(self):
        return (
            f"Recipe(id={self.id}, factories={self.factories}, ingredients={self.ingredients}, results={self.results})"
        )