from datetime import datetime
from random import choice

import polars as pl
from scipy.stats.distributions import lognorm
from sqlalchemy.orm import Session

from retailtwin.models import Cart


def cart(session: Session) -> Cart:
    db_uri = str(session.get_bind().url.render_as_string(hide_password=False))

    customer = pl.read_database("select random_customer()", db_uri, engine="adbc")[0, 0]
    location = pl.read_database("select random_location()", db_uri, engine="adbc")[0, 0]
    cart = Cart(customer=customer, location=location, checkout=datetime.now())
    session.add(cart)
    session.commit()

    # Budget can be a small or a large expenditure
    budget = choice([10 * lognorm.rvs(0.25), 100 * lognorm.rvs(0.25)])
    items = pl.read_database("select * from items", db_uri, engine="adbc")
    spent = 0
    while spent > budget:
        quantity = choice([1, 2, 3, 4])
        item = items.sample(1, shuffle=True)
        price = pl.read_database(
            "select price_on_location({id}, {qty})", db_uri, engine="adbc"
        )[0, 0]
        session.execute(
            f"add_item_to_cart({cart.id}, {item.sku}, {quantity}, {location})"
        )
        session.commit()
        spent += price


if __name__ == "__main__":
    from sqlalchemy import create_engine
    from sqlalchemy.orm import Session

    with Session(
        create_engine("postgresql://postgres:postgres@localhost/retail", echo=True)
    ) as session:
        cart(session)
