SatisfactoryPlusCalculator/factorygame/data/fetch.py

#!/usr/bin/env python3

import click
from sqlalchemy import create_engine, select, delete
from sqlalchemy.orm import Session

from factorygame.data.common import resource_needs_update, chose_resource
from .models import Base, ResourceFlow, Recipe
from .sfp import SatisfactoryPlus, DEFAULT_IGNORE_FACTORIES
from ..helper import click_prompt


@click.command()
@click.option("--debug", is_flag=True)
@click.option("--refetch", is_flag=True)
@click.option("--ignore-factories", type=list[str], default=DEFAULT_IGNORE_FACTORIES)
@click.argument("search")
def main(debug: bool, refetch: bool, ignore_factories: list[str], search: str):
    engine = create_engine("sqlite:///file.db", echo=debug)
    Base.metadata.create_all(bind=engine)
    if not search:
        print("Empty search option. Exiting…")
        return

    with Session(engine) as session:
        resource = chose_resource(session=session, resource_label=search, prompt=click_prompt).result()
        exists_in_db = resource is not None

        with SatisfactoryPlus(ignore_factories=ignore_factories, debug=debug) as data_provider:
            if resource is None:
                ret = data_provider.search_for_resource(session=session, search=search, prompt=click_prompt).result()
                if ret is None:
                    return
                else:
                    resource, exists_in_db = ret

            refetch |= resource_needs_update(resource)
            if refetch:
                if exists_in_db:
                    print("Deleting recipes for", resource.label)
                    with session.begin_nested():
                        resource.recipes_populated_at = None
                        flow_ids_to_delete: list[int] = list()
                        for flow in session.scalars(
                            select(ResourceFlow).where(ResourceFlow.resource_id == resource.id)
                        ):
                            if flow.ingredient_in:
                                flow_ids_to_delete += map(lambda obj: obj.id, flow.ingredient_in.ingredients)
                                flow_ids_to_delete += map(lambda obj: obj.id, flow.ingredient_in.results)
                                session.delete(flow.ingredient_in)
                            if flow.result_of:
                                flow_ids_to_delete += map(lambda obj: obj.id, flow.result_of.ingredients)
                                flow_ids_to_delete += map(lambda obj: obj.id, flow.result_of.results)
                                session.delete(flow.result_of)
                        stmt = delete(ResourceFlow).where(ResourceFlow.id.in_(flow_ids_to_delete))
                        session.execute(stmt)
                    print("Refetching recipes for", resource.label)
                else:
                    print("Fetching recipes for new resource", resource.label)

                with session.begin_nested():
                    resource = data_provider.update_resource_recipes(session=session, resource=resource)

            session.refresh(resource)
            assert resource, "Resource must be set at this point"

            stmt = select(Recipe).distinct().join(Recipe.ingredients).filter(ResourceFlow.resource_id == resource.id)
            for recipe in session.scalars(stmt):
                print("Used in recipe", recipe.describe())

            stmt = select(Recipe).distinct().join(Recipe.results).filter(ResourceFlow.resource_id == resource.id)
            for recipe in session.scalars(stmt):
                print("Result of recipe", recipe.describe())
        session.commit()


if __name__ == "__main__":
    main()