PyConES24/src/retailtwin/cli/stock.py

# ruff: noqa: E501
import typer
import polars as pl
from sqlalchemy import create_engine, text as query
from retailtwin.cli.db import query_warehouse_stock, df_to_table, fetch_from_db
from prompt_toolkit import PromptSession
from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.shortcuts import clear
from rich.console import Console
from rich.markdown import Markdown

console = Console()

HELP = """
# Retail twin stock management CLI

This is a simple terminal to manage stock. Enter a single-letter command followed by <Enter>. The available commands are:

* **l**: Lists all the current items stocked in any location.
* **s**: Enters search mode. Search an item by name.
* **q**: Store query mode. Queries the stock of an item by UPC in the current location.
* **w**: Warehouse query mode. Queries the stock of an item by UPC in all warehouses. Requires connection to the database.
* **c**: Cancel mode. Retires a batch giving a UPC. Requires connection to the database.
* **b**: Batch mode. Requests a given quantity from an item to the warehouse. Requires connection to the database.
* **r**: Refresh data from the stock database.
* **h**: Print this help message.
* **x**: Exit this terminal.

"""


def handle_command(command: str, **kwargs):
    """_summary_

    Args:
        command (str): Single letter command
    """
    items = kwargs["items"]
    session = kwargs["session"]
    item_completer = kwargs["completer"]
    location = kwargs["location"]
    local_stock = kwargs["local_stock"]
    db_uri = kwargs["db_uri"]

    if command == "l":
        console.print(df_to_table(items))
    elif command == "s":
        # Command opens a new prompt
        completer = WordCompleter(
            items.select(pl.col("name")).to_series().str.to_lowercase().to_list()
        )
        text = session.prompt("s> ", completer=completer)
        console.print(
            df_to_table(
                items.select(pl.all()).filter(
                    pl.col("name").str.to_lowercase().str.contains(text)
                )
            )
        )
    elif command == "q":  # Query stock in current location
        upc = session.prompt("q> ", completer=item_completer)
        df = local_stock.select(
            [
                pl.col("upc"),
                pl.col("name"),
                pl.col("package"),
                pl.col("unitprice"),
                pl.col("best_until"),
                pl.col("quantity"),
            ]
        ).filter(pl.col("upc") == upc)
        console.print(
            df_to_table(
                df.with_columns(pl.col("unitprice").cast(str)),
                title=f"Item {upc} on location {location}",
            )
        )

    elif command == "w":  # Query stock in warehouses
        text = session.prompt("w> ", completer=item_completer)
        console.print(
            df_to_table(
                query_warehouse_stock(db_uri, text).with_columns(
                    pl.col("received").cast(str),
                    pl.col("best_until").cast(str),
                    pl.col("upc").cast(str),
                    pl.col("quantity").cast(str),
                ),
                title=f"Warehouse stock for item {text}",
            )
        )

    elif command == "c":  # Retire a batch from the shelves.
        batches = local_stock.select(pl.col("batch")).to_series().to_list()
        batchid = session.prompt("c [enter batch id]> ")

        if batchid in batches:
            console.print(
                df_to_table(
                    local_stock.select(pl.all()).filter(pl.col("batch") == batchid),
                    title=f"Batch {batchid} on location {location}",
                )
            )
            text = session.prompt(f"c [Remove batch {batchid}? (Y/N)]> ")
            if text.lower() == "y":
                engine = create_engine(db_uri)
                conn = engine.connect()
                conn.execute(
                    query(f"select retire_batch_from_shelves({batchid}, {location})")
                )
                conn.commit()
            else:
                console.print("Aborted")

        else:
            console.print(f"Error, batch {batchid} doesn't exist")
    else:
        console.print(f"Command {command} not supported")


def main(db_uri: str, location: int):
    # Fetch all data necessary from the database
    console.print("Fetching data...")
    items, local_stock = fetch_from_db(db_uri, location)

    clear()
    console.print(Markdown(HELP))

    # Default completer with barcodes
    completer = WordCompleter(
        items.select(pl.col("upc")).to_series().cast(str).to_list()
    )
    session = PromptSession()

    while True:
        try:
            text = session.prompt(
                "#> ",
            )
        except KeyboardInterrupt:
            continue
        except EOFError:
            break
        else:
            if text == "x":
                break
            elif text == "h":
                console.print(Markdown(HELP))
            elif text == "r":
                console.print("Refreshing...")
                items, local_stock = fetch_from_db(db_uri, location)
            else:
                handle_command(
                    text,
                    items=items,
                    location=location,
                    session=session,
                    completer=completer,
                    local_stock=local_stock,
                    db_uri=db_uri,
                )

    print("GoodBye!")


app = typer.run(main)