import orjson as json
from typing import Callable
from datetime import datetime
from decimal import Decimal

from helpers import dx_rating, find_level, query_music_db, salted_hash_userid


def clean_b50(b50: dict[str, str | dict]):
    urating: dict[str, list[dict[str, int]]] = b50["userRating"]

    def add_rating(entry: dict[str, int]):
        """
        ```
        {
            "musicId": 11638,
            "level": 2,
            "romVersion": 24005,
            "achievement": 988145
        }
        ```
        - level: EXPERT
        - ver: DX, 1.40.05
        - ach: 98.8145%
        """

        entry["musicTitle"] = None
        entry["difficulty"] = None
        entry["dxRating"] = 0

        music_info = query_music_db(entry["musicId"])

        if music_info is None:
            return

        entry["musicTitle"] = music_info["name"]
        levels = find_level(music_info, entry["level"])

        if not levels:
            return

        level: dict[str, str | int] = levels.pop()
        difficulty = level["difficulty"]

        entry["difficulty"] = difficulty
        entry["dxRating"] = dx_rating(
            difficulty=Decimal(difficulty),
            achievement=entry["achievement"],
        )

    for b35 in urating["ratingList"]:
        add_rating(b35)
    for b15 in urating["newRatingList"]:
        add_rating(b15)

    urating["rating"] = sum(
        map(
            lambda lst: sum(map(lambda entry: entry["dxRating"], urating[lst])),
            ["ratingList", "newRatingList"],
        )
    )


def record_time(*, _: list[datetime] = []):
    last_time = _
    if not last_time:
        last_time.append(datetime.now())
    else:
        new = datetime.now()
        diff = (new - last_time.pop()).total_seconds()
        last_time.append(new)
        return diff


def process(
    clean_fields: Callable[[dict], None],
    input_file: str,
    output_file: str,
):
    record_time()
    with open(input_file, "rb") as f:
        data = json.loads(f.read())
    print(f"loaded, cost {record_time():.2f}s")

    for entry in data:
        entry["userId"] = salted_hash_userid(entry["userId"])
        clean_fields(entry)
    print(f"processed, cost {record_time():.2f}s")

    with open(output_file, "wb") as f:
        f.write(json.dumps(data))
    print(f"written out, cost {record_time():.2f}s")

    return data


def main():
    process(
        clean_b50,
        "b50.json",
        "b50_pub.json",
    )


if __name__ == "__main__":
    main()