import requests
import time
import tomllib
from os import path
import json

if path.exists("config.toml"):
    with open("config.toml", "rb") as f:
        config = tomllib.load(f)
else:
    raise RuntimeError("config.toml not found")

TTL = config["anilist"].get("TTL", 2 * 60 * 60) # 2 hour default if not specificed

def getMediaListCollection(
    mediaType=config["anilist"]["mediaType"],
    userId=config["anilist"]["userId"],
    *,
    cachePath="listdata.json"
):
    """
    Returns a tuple of (entries, cached)
    
    Entries is a dictionary of the values to be written to the cache file, such as:
        - timestamp
        - ttl
        - params supplied
        - response
        
    cached is a boolean indicating whether the cache file was used or not
    
    Takes the following arguments:
    - mediaType: the media type to use for the request (if not supplied, defaults to value supplied from config.toml)
    - userId: the user id to use for the request (if not supplied, defaults to value supplied from config.toml)
    - cachePath: the path to the cache file (if not supplied, defaults to "data.json")
    """
    
    if not mediaType or not userId:
        raise RuntimeError("MediaType and UserId must be set")

    if path.exists(cachePath):
        with open(cachePath, "r") as f:
            entries = json.load(f)

        if (
            "timestamp" in entries
            and "ttl" in entries
            and "response" in entries
            and "params" in entries
        ):
            timestamp = float(entries["timestamp"])
            ttl = float(entries["ttl"])

            cacheUserId = entries["params"].get("userId")
            cacheMediaType = entries["params"].get("mediaType")

            cacheExpired = time.time() - timestamp > ttl
            parametersChanged = (
                cacheUserId != userId
                or cacheMediaType != mediaType
            )

            if not cacheExpired and not parametersChanged:
                print("Cache is fresh, using cached data...")
                return (entries, True)

            if cacheExpired:
                print("Cache has gone stale, refreshing...")
            elif parametersChanged:
                print("Cache parameters changed, refreshing...")
        else:
            print("Cache is invalid, refreshing...")

    # Fetch fresh data from AniList
    query = """
    query ($type: MediaType!, $userId: Int!, $chunk: Int, $perChunk: Int) {
        MediaListCollection(
            type: $type,
            userId: $userId,
            chunk: $chunk,
            perChunk: $perChunk
        ) {
            hasNextChunk
            lists {
                name
                entries {
                    id
                    media {
                        id
                        title {
                            english
                            romaji
                        }
                    }
                }
            }
        }
    }
    """

    entries = {}
    chunk = 1

    while True:
        variables = {
            "type": mediaType,
            "userId": int(userId),
            "chunk": chunk,
            "perChunk": 500,
        }

        resp = requests.post(
            "https://graphql.anilist.co",
            json={
                "query": query,
                "variables": variables,
            },
            timeout=30,
        )

        resp.raise_for_status()

        data = resp.json()

        if "errors" in data:
            raise RuntimeError(data["errors"])

        collection = data["data"]["MediaListCollection"]

        for lst in collection["lists"]:
            entries.setdefault(lst["name"], []).extend(lst["entries"])

        if not collection["hasNextChunk"]:
            break

        chunk += 1

    form = {
        "timestamp": str(time.time()),
        "ttl": str(TTL),
        "params": {
            "mediaType": mediaType,
            "userId": userId,
        },
        "response": entries
    }

    return (form, False)

def getAiringScheduleBatch(
    animeIds,
    *,
    ttl=TTL,
    cachePath="airingdata.json",
    maxRetries=5
):
    """
    Returns a tuple of (entries, cached)

    Entries is a dictionary of the values to be written to the cache file, such as:
        - timestamp
        - ttl
        - params supplied
        - response

    response is a dictionary of mediaId (as a string) -> {"title": {"romaji", "english"}, "schedule": [upcoming airing schedule items]}

    cached is a boolean indicating whether the cache file was used or not

    Takes the following arguments:
    - animeIds: list of AniList media ids to fetch airing schedules for
    - ttl: seconds the cache stays fresh for. Only applied when the cache is refreshed;
      an existing cache keeps the ttl it was written with (defaults to value from config.toml)
    - cachePath: the path to the cache file (if not supplied, defaults to "airingdata.json")
    - maxRetries: how many times in a row to retry after a 429 before giving up (defaults to 5)
    """

    if not animeIds:
        raise RuntimeError("animeIds must not be empty")

    animeIds = sorted(set(int(i) for i in animeIds))

    if path.exists(cachePath):
        with open(cachePath, "r") as f:
            entries = json.load(f)

        if (
            "timestamp" in entries
            and "ttl" in entries
            and "response" in entries
            and "params" in entries
        ):
            timestamp = float(entries["timestamp"])
            cacheTtl = float(entries["ttl"])

            cacheIds = entries["params"].get("animeIds")

            cacheExpired = time.time() - timestamp > cacheTtl
            parametersChanged = cacheIds != animeIds

            if not cacheExpired and not parametersChanged:
                print("Airing cache is fresh, using cached data...")
                return (entries, True)

            if cacheExpired:
                print("Airing cache has gone stale, refreshing...")
            elif parametersChanged:
                print("Airing cache parameters changed, refreshing...")
        else:
            print("Airing cache is invalid, refreshing...")

    query = """
    query ($ids: [Int], $page: Int) {
        Page(page: $page, perPage: 50) {
            pageInfo {
                hasNextPage
                currentPage
            }
            airingSchedules(mediaId_in: $ids, notYetAired: true, sort: TIME) {
                id
                airingAt
                timeUntilAiring
                episode
                mediaId
                media {
                    title {
                        romaji
                        english
                    }
                }
            }
        }
    }
    """

    schedules = {}
    page = 1
    retries = 0

    while True:
        variables = {
            "ids": animeIds,
            "page": page,
        }

        resp = requests.post(
            "https://graphql.anilist.co",
            json={
                "query": query,
                "variables": variables,
            },
            timeout=30,
        )

        # retry if rate limited
        if resp.status_code == 429:
            retries += 1
            if retries > maxRetries:
                raise RuntimeError("Rate limited too many times, giving up")
            wait = int(resp.headers.get("Retry-After", 60))
            print(f"Rate limited, waiting {wait}s...")
            time.sleep(wait)
            continue

        retries = 0
        resp.raise_for_status()

        data = resp.json()

        if "errors" in data:
            raise RuntimeError(data["errors"])

        pageData = data["data"]["Page"]

        for item in pageData["airingSchedules"]:
            # media can be null, fall back to empty titles
            media = item.pop("media", None) or {}
            title = media.get("title") or {}

            entry = schedules.setdefault(
                str(item["mediaId"]),
                {
                    "title": {
                        "romaji": title.get("romaji"),
                        "english": title.get("english"),
                    },
                    "schedule": [],
                },
            )
            entry["schedule"].append(item)

        if not pageData["pageInfo"]["hasNextPage"]:
            break

        page += 1

    form = {
        "timestamp": str(time.time()),
        "ttl": str(ttl),
        "params": {
            "animeIds": animeIds,
        },
        "response": schedules
    }

    return (form, False)

if __name__ == "__main__":
    entries, cached = getMediaListCollection()

    if not cached:
        with open("listdata.json", "w") as f:
            f.write(f"{json.dumps(entries, indent=2)}")

    lists = entries["response"]

    print("\nList Names updating for:")
    for li in config["anical"]["listNames"]:
        print(li)

    # get a list of all the animes ids in these lists:
    animeIds = []
    for li in config["anical"]["listNames"]:
        if li not in lists:
            print(f"Warning: list '{li}' not found, skipping")
            continue
        for entry in lists[li]:
            animeIds.append(entry["media"]["id"])

    print(f"\nFound {len(animeIds)} anime ids")

    airing, airingCached = getAiringScheduleBatch(animeIds)

    if not airingCached:
        with open("airingdata.json", "w") as f:
            f.write(f"{json.dumps(airing, indent=2)}")

    schedules = airing["response"]
    print(f"\nGot upcoming episodes for {len(schedules)} anime")