import argparse
import logging
from typing import Callable

from unwind import db, imdb, models

log = logging.getLogger(__name__)

name = "load-imdb-charts"
help = "Load and import charts from imdb.com."


def add_args(cmd: argparse.ArgumentParser) -> None:
    cmd.add_argument(
        "--select",
        action="append",
        dest="charts",
        default=[],
        choices={"top250", "bottom100", "pop100"},
        help="Select which charts to refresh.",
    )


async def remove_all_awards(
    conn: db.Connection, category: models.AwardCategory
) -> None:
    stmt = models.awards.delete().where(models.awards.c.category == category)
    await conn.execute(stmt)


_award_handlers: dict[models.AwardCategory, Callable] = {
    "imdb-pop-100": imdb.load_most_popular_100,
    "imdb-top-250": imdb.load_top_250,
    "imdb-bottom-100": imdb.load_bottom_100,
}


async def update_awards(conn: db.Connection, category: models.AwardCategory) -> None:
    load_imdb_ids = _award_handlers[category]
    imdb_ids = await load_imdb_ids()

    available = await db.get_movie_ids(conn, imdb_ids)
    if missing := set(imdb_ids).difference(available):
        log.warning(
            "⚠️ Charts for category (%a) contained %i unknown movies: %a",
            category,
            len(missing),
            missing,
        )

    await remove_all_awards(conn, category=category)

    for pos, imdb_id in enumerate(imdb_ids, 1):
        if (movie_id := available.get(imdb_id)) is None:
            continue

        award = models.Award(
            movie_id=movie_id,
            category=category,
        )
        award.position = pos
        await db.add(conn, award)


async def main(args: argparse.Namespace) -> None:
    await db.open_connection_pool()

    if not args.charts:
        args.charts = {"top250", "bottom100", "pop100"}

    if "pop100" in args.charts:
        async with db.transaction() as conn:
            await update_awards(conn, "imdb-pop-100")
        log.info("✨ Updated most popular 100 movies.")

    if "bottom100" in args.charts:
        async with db.transaction() as conn:
            await update_awards(conn, "imdb-bottom-100")
        log.info("✨ Updated bottom 100 movies.")

    if "top250" in args.charts:
        async with db.transaction() as conn:
            await update_awards(conn, "imdb-top-250")
        log.info("✨ Updated top 250 rated movies.")

    await db.close_connection_pool()
