#!/usr/bin/env python
"""Download full game feeds."""

import argparse
import asyncio
from pathlib import Path

import aiofiles
import aiofiles.os
import backoff
import httpx
import msgpack
import toolz
from tqdm.asyncio import tqdm


api_url = "https://statsapi.web.nhl.com/api/v1{}".format


def fetch(url: str):
    @backoff.on_exception(
        backoff.expo,
        (httpx.ConnectTimeout, httpx.ConnectError),
        max_time=30,
    )
    async def _(c: httpx.Client, **kwargs):
        rsp = await c.get(api_url(url.format(**kwargs)))
        return rsp.json()

    return _


fetch_teams = fetch("/teams")
fetch_roster = fetch("/teams/{team_id}/roster?season={season}")
fetch_player_info = fetch("/people/{player_id}")
fetch_player_stats = fetch("/people/{player_id}/stats?stats={stats_type}")


def player_datadir(player_id: int) -> Path:
    return Path(f"data/players/{player_id}")


def player_data(player_id: int, what: str) -> Path:
    return player_datadir(player_id) / f"{what}.msgpack"


async def save_player_data(player_id: int, what: str, data):
    datadir = player_datadir(player_id)
    if not datadir.exists():
        await aiofiles.os.mkdir(datadir)
    async with aiofiles.open(datadir / f"{what}.msgpack", "wb") as f:
        await f.write(msgpack.dumps(data))


async def load_player_data(player_id: int, what: str):
    async with aiofiles.open(player_data(player_id, what), "rb") as f:
        return msgpack.unpackb(await f.read())


async def get_player_stats(c, player_id, stats_type):

    if player_data(player_id, stats_type).exists():
        stats = await load_player_data(player_id, stats_type)
        return player_id, stats

    rsp = await fetch_player_stats(c, player_id=player_id, stats_type=stats_type)
    stats = toolz.get_in(["stats", 0, "splits"], rsp)
    await save_player_data(player_id, stats_type, stats)
    return player_id, stats


def seasons_in_nhl(yoy_stats) -> list[str]:
    return [
        y["season"]
        for y in yoy_stats
        if y["league"]["name"] == "National Hockey League"
    ]


async def get_player_gamelog(c, player_id, yoy_stats):

    if player_data(player_id, "gameLog").exists():
        total_game_log = await load_player_data(player_id, "gameLog")
        return player_id, total_game_log

    game_logs = await asyncio.gather(
        *[
            fetch_player_stats(
                c, player_id=player_id, stats_type=f"gameLog&season={season}"
            )
            for season in seasons_in_nhl(yoy_stats)
        ]
    )
    total_game_log = list(
        toolz.mapcat(
            toolz.curried.get_in(["stats", 0, "splits"]),
            game_logs,
        )
    )
    await save_player_data(player_id, "gameLog", total_game_log)
    return player_id, total_game_log


async def get_player_info(c, player_id):

    if player_data(player_id, "info").exists():
        return await load_player_data(player_id, "info")

    rsp = await fetch_player_info(c, player_id=player_id)
    info = toolz.get_in(["people", 0], rsp)
    await save_player_data(player_id, "info", info)
    return player_id, info


async def main(season: int, force: bool = False):
    season_id = f"{season}{season + 1}"

    c = httpx.AsyncClient(
        limits=httpx.Limits(max_keepalive_connections=5, max_connections=10)
    )

    async with c:
        teams = await fetch_teams(c)
        teams = teams["teams"]

        # -- Team rosters
        futs = tqdm.as_completed(
            [fetch_roster(c, team_id=team["id"], season=season_id) for team in teams],
            desc="Rosters",
        )
        rosters = []
        for r in futs:
            r = await r
            rosters.append(r["roster"])

        player_ids = [p["person"]["id"] for p in toolz.concat(rosters)]

        # -- Player info
        futs = tqdm.as_completed(
            [get_player_info(c, player_id=pid) for pid in player_ids],
            desc="Player info",
        )
        for r in futs:
            await r

        # -- Year-on-year stats
        futs = tqdm.as_completed(
            [
                get_player_stats(c, player_id=pid, stats_type="yearByYear")
                for pid in player_ids
            ],
            desc="Year-on-Year stats",
        )
        yoy_stats = []
        for r in futs:
            yoy_stats.append(await r)

        # -- Game-by-game stats
        futs = tqdm.as_completed(
            [
                get_player_gamelog(c, player_id=pid, yoy_stats=stats)
                for pid, stats in yoy_stats
            ],
            desc="Game-by-Game stats",
        )
        for r in futs:
            await r


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Download stats for all NHL players")
    parser.add_argument("season", type=int, help="The season's roster to include")
    args = parser.parse_args()
    asyncio.run(main(args.season))