import os

from sqlmodel import create_engine, Session, select, update
from functools import lru_cache
from typing import Union
from sqlalchemy.exc import NoResultFound

engine = create_engine(os.environ.get('DB_CONN'))

# Grim hack to get the imports working with crawler and main.
# TODO: Split poke models and other common functions out into a separate package the api+crawler can share.
# TODO: After split crawler code out into a separate part of the repo and create an individual Docker image for it.

try:
    from poke.poke_model import pokemon as pokemon_model
except:
    from poke_model import pokemon as pokemon_model


@lru_cache(maxsize=16)
def get_pokemon(poke_id: int) -> pokemon_model:
    """ Get a pokemon's data from the database from its ID.

    Args:
        poke_id: ID of the pokemon you want the data for.

    Returns:
        pokemon_model object containing the data for the pokemon found in the DB.

    Raises:
        NoResultFound: If there isn't a pokemon in the DB with the passed in ID.
    """
    with Session(engine) as session:
        poke = session.exec(