SatisfactoryPlusCalculator/factorygame/data/fetch.py

#!/usr/bin/env python3

import datetime
from typing import Optional

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

from .models import Base, Resource, ResourceFlow, Recipe
from .sfp import SatisfactoryPlus
from ..helper import prompt

__recipe_info_timeout = datetime.timedelta(days=30)


@click.command()
@click.option("--result", is_flag=True)
@click.option("--debug", is_flag=True)
@click.option("--refetch", is_flag=True)
@click.argument("search")
def main(result: bool, debug: bool, refetch: bool, search: str):
    engine = create_engine("sqlite:///file.db", echo=debug)
    Base.metadata.create_all(bind=engine)
    if result and search:
        do_provider_search = True
        resource: Optional[Resource] = None
        with Session(engine) as session:
            matching_resources = session.scalars(Resource.by_label(search)).all()
            if len(matching_resources) == 0:
                print("Could not find existing resources matching the search string.. starting wiki search")
            else:
                options = {(idx + 1): str(matching_resources[idx].label) for idx in range(len(matching_resources))}
                options[0] = ""
                selected_res = prompt(
                    options=options, text="Chose a resource to continue or 0 to continue with a wiki search", default=1
                )
                if selected_res is None:
                    return
                if selected_res != 0:
                    resource = matching_resources[selected_res - 1]
                    do_provider_search = False
                    exists_in_db = True

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

                refetch = (
                    refetch
                    or resource.recipes_populated_at is None
                    or datetime.datetime.utcnow() - resource.recipes_populated_at > __recipe_info_timeout
                )
                if refetch and exists_in_db:
                    print("Deleting recipes for", resource.label)
                    with session.begin_nested():
                        for flow in session.scalars(
                            select(ResourceFlow).where(ResourceFlow.resource_id == resource.id)
                        ):
                            if flow.result_of:
                                for flow2 in flow.result_of.ingredients:
                                    session.delete(flow2)
                                for flow2 in flow.result_of.results:
                                    session.delete(flow2)
                                session.delete(flow.result_of)

                if refetch:
                    if exists_in_db:
                        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).join(Recipe.results).filter(ResourceFlow.resource_id == resource.id)
                for recipe in session.scalars(stmt):
                    print("Recipe:", recipe.describe())


if __name__ == "__main__":
    main()