Initial commit

7 files changed, 245 insertions, 0 deletions

diff --git a/sbdata/__init__.py b/sbdata/__init__.py
new file mode 100644
index 0000000..572b454
--- /dev/null
+++ b/sbdata/__init__.py
@@ -0,0 +1 @@
+from sbdata import tasks as _
diff --git a/sbdata/__main__.py b/sbdata/__main__.py
new file mode 100644
index 0000000..596a51f
--- /dev/null
+++ b/sbdata/__main__.py
@@ -0,0 +1,36 @@
+import dataclasses
+import json
+import sys
+from typing import Any
+
+import questionary
+
+from sbdata.repo import Item
+from sbdata.task import Arguments, tasks
+
+
+class ObjectEncoder(json.JSONEncoder):
+
+    def default(self, o: Any) -> Any:
+        if isinstance(o, Item):
+            return o.internalname
+        if dataclasses.is_dataclass(o):
+            return o.__dict__
+        return super().default(o)
+
+
+def main():
+    args = Arguments(sys.argv)
+    task = args.get_value(
+        'Task', tasks.get(args.task),
+        questionary.select('Which task do you want to execute?', choices=[
+            questionary.Choice(task.label, task) for task in tasks.values()
+        ]))
+    print("Selected task: " + task.label)
+    data = task.run(args)
+    if args.has_flag('json'):
+        print(json.dumps(data, cls=ObjectEncoder))
+
+
+if __name__ == '__main__':
+    main()
diff --git a/sbdata/repo.py b/sbdata/repo.py
new file mode 100644
index 0000000..475f2a6
--- /dev/null
+++ b/sbdata/repo.py
@@ -0,0 +1,61 @@
+import dataclasses
+import json
+import os
+import pathlib
+import re
+import typing
+
+item_list = {}
+repo_dir = pathlib.Path(os.environ.get('REPO_DIR', 'NotEnoughUpdates-REPO'))
+
+
+@dataclasses.dataclass
+class Item:
+    displayname: str
+    itemid: str
+    internalname: str
+    lore: list[str]
+
+
+def unformat_name(name: str) -> str:
+    return re.sub('§.', '', name)
+
+
+def bare_name(name: str) -> str:
+    name = unformat_name(name).replace("'", '').lower()
+    if name.startswith('ultimate'):
+        name = name[8:]
+    return name.strip()
+
+
+def load_item(item: pathlib.Path):
+    with item.open('r') as fp:
+        data = json.load(fp)
+        item_list[data['internalname']] = Item(data['displayname'], data['itemid'], data['internalname'], data['lore'])
+
+
+def find_item_by_name(name: str) -> typing.Optional[Item]:
+    name = bare_name(name)
+    pot = [item for item in item_list.values()
+           if item.internalname.casefold() == name
+           or bare_name(item.displayname) in name
+           or (item.itemid == 'minecraft:enchanted_book'
+               and name in bare_name(item.lore[0]))]
+    if pot:
+        return pot[0]
+    return None
+
+
+def load_items():
+    item_dir = repo_dir / 'items'
+    for item in item_dir.iterdir():
+        if item.name.endswith('.json'):
+            load_item(item)
+
+
+def load_repo_data():
+    item_list.clear()
+    load_items()
+
+
+load_repo_data()
diff --git a/sbdata/task.py b/sbdata/task.py
new file mode 100644
index 0000000..7011ca2
--- /dev/null
+++ b/sbdata/task.py
@@ -0,0 +1,64 @@
+import dataclasses
+import os
+import sys
+import typing
+
+import questionary
+
+_T = typing.TypeVar('_T')
+
+
+class Arguments:
+
+    def __init__(self, args: list[str]):
+        self.prog = args[0]
+        self.args: typing.Dict[str, str] = {}
+        self.flags: list[str] = []
+        self.no_prompt = os.environ.get('PROMPT') == 'NO_PROMPT'
+        self.task: typing.Optional[str]
+        self.task = None
+        last_arg = None
+        for arg in args:
+            if last_arg is None:
+                if arg.startswith('--'):
+                    last_arg = arg[2:]
+                elif arg.startswith('-'):
+                    self.flags.append(arg[1:])
+                elif arg.startswith(':'):
+                    self.task = arg[1:]
+                else:
+                    print("Unknown arg: " + arg)
+            else:
+                self.args[last_arg] = arg
+                last_arg = None
+
+    def get_value(self, label: str, value: _T, question: questionary.Question) -> _T:
+        if value is None:
+            if self.no_prompt:
+                print('No argument present for ' + label)
+                sys.exit(1)
+            return question.ask()
+        return value
+
+    def has_flag(self, param: str) -> bool:
+        return param in self.flags
+
+
+@dataclasses.dataclass
+class Task:
+    label: str
+    name: str
+    run: typing.Callable
+
+
+tasks = {}
+
+TASK_TYPE = typing.Callable[[Arguments], None]
+
+
+def register_task(label: str) -> typing.Callable[[TASK_TYPE], TASK_TYPE]:
+    def d(func: TASK_TYPE) -> TASK_TYPE:
+        tasks[func.__name__] = Task(label, func.__name__, func)
+        return func
+
+    return d
diff --git a/sbdata/tasks.py b/sbdata/tasks.py
new file mode 100644
index 0000000..8ac1f8c
--- /dev/null
+++ b/sbdata/tasks.py
@@ -0,0 +1,55 @@
+import ast
+import dataclasses
+import json
+import re
+import sys
+import typing
+
+from sbdata.repo import find_item_by_name, Item
+from sbdata.task import register_task, Arguments
+from sbdata.wiki import get_wiki_sources_by_title
+
+
+@dataclasses.dataclass
+class DungeonDrop:
+    item: Item
+    floor: int
+    chest: str
+    drop_chances: dict[str, str]
+
+    def get_drop_chance(self, has_s_plus: bool, talisman_level: int, boss_luck: int):
+        drop_identifier = "S" + ('+' if has_s_plus else '') + 'ABCD'[talisman_level] + str(len([i for i in [0, 1, 3, 5, 10] if i >= boss_luck]))
+        return self.drop_chances.get(drop_identifier)
+
+
+@register_task("Fetch Dungeon Loot")
+def fetch_dungeon_loot(args: Arguments):
+    items = []
+    for floor in get_wiki_sources_by_title(*[f'Template:Catacombs Floor {f} Loot Master' for f in ['I', 'II', 'III', 'IV', 'V', 'VI', 'VII']]).values():
+        for template in floor.filter_templates():
+            if template.name.strip() == 'Dungeon Chest Table/Row':
+                item = None
+                ifloor = None
+                chest = None
+                drop_chances = {}
+
+                for param in template.params:
+                    attr_name = param.name.nodes[0].strip()
+                    attr_value = param.value.nodes[0].strip()
+                    if attr_name == 'item':
+                        if item is None:
+                            item = find_item_by_name(attr_value)
+                    elif attr_name == 'customlink':
+                        if item is None:
+                            item = find_item_by_name(attr_value.split('#')[-1])
+                    elif attr_name == 'chest':
+                        chest = attr_value
+                    elif attr_name == 'floor':
+                        ifloor = int(attr_value)
+                    elif attr_name.startswith("S"):
+                        drop_chances[attr_name] = attr_value
+                if item is None or ifloor is None or chest is None:
+                    print('WARNING: Missing data for item: ' + str(template))
+                else:
+                    items.append(DungeonDrop(item, ifloor, chest, drop_chances))
+    return items
diff --git a/sbdata/util.py b/sbdata/util.py
new file mode 100644
index 0000000..87a3098
--- /dev/null
+++ b/sbdata/util.py
@@ -0,0 +1,13 @@
+import functools
+from typing import TypeVar, Generator, ParamSpec, Callable
+
+_Param = ParamSpec('_Param')
+_RetType = TypeVar('_RetType')
+
+
+def no_generator(func: Callable[_Param, Generator[_RetType, None, None]]) -> Callable[_Param, list[_RetType]]:
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs) -> list[_RetType]:
+        return list(func(*args, **kwargs))
+
+    return wrapper
diff --git a/sbdata/wiki.py b/sbdata/wiki.py
new file mode 100644
index 0000000..7e029a5
--- /dev/null
+++ b/sbdata/wiki.py
@@ -0,0 +1,15 @@
+import urllib.parse
+import requests
+import mwparserfromhell
+
+
+def get_wiki_sources_by_title(*page_titles: str, wiki_host: str = 'wiki.hypixel.net') -> dict[str, mwparserfromhell.wikicode.Wikicode]:
+    prepared_titles = "|".join(map(urllib.parse.quote, page_titles))
+    api_data = requests.get(f'https://{wiki_host}/api.php?action=query&prop=revisions&titles={prepared_titles}&rvprop=content&format=json&rvslots=main').json()
+    if "batchcomplete" not in api_data:
+        print(f'Batch data not present in wiki response for: {page_titles}')
+
+    return {
+        page['title']: mwparserfromhell.parse(page['revisions'][0]['slots']['main']['*'])
+        for _, page in api_data["query"]["pages"].items()
+    }