import os
import json
import requests
import re
import unicodedata
from collections import defaultdict
from bs4 import BeautifulSoup
from datetime import datetime, timezone
from htb_client import HTBClient
from trilium_py.client import ETAPI
 
# Get the absolute path of the script's directory
script_dir = os.path.dirname(os.path.abspath(__file__))
 
# Construct the full path to the config.json file
config_path = os.path.join(script_dir, 'config.json')
 
# Load configuration from the JSON file
with open(config_path, 'r') as f:
    config = json.load(f)
 
# Accessing config values
htb_code = config['htb_code']
trilium_server_url = config['trilium_server_url']
trilium_token = config['trilium_token']
trilium_challenges_folder = config['trilium_challenges_folder']
trilium_challenges_template_id = config['trilium_challenges_template_id']
 
def get_timestamp(machine):
    # Parse the release date string into a datetime object
    release_str = machine['release_date']
    dt = datetime.strptime(release_str, "%Y-%m-%dT%H:%M:%S.%fZ")
    
    # Set the timezone to UTC
    dt = dt.replace(tzinfo=timezone.utc)
    
    # Get the timestamp
    timestamp = dt.timestamp()
    return timestamp
 
print("[+] connecting to HTB")
client = HTBClient(password=htb_code)
print("[+] connecting to trilium")
ea = ETAPI(trilium_server_url, trilium_token)
print("[i] version: ", ea.app_info()['appVersion'])
 
print("[i] HTB User:", client.user['id'], "-", client.user['name'])
 
 
print("[+] gathering challenges info")
categories = defaultdict(list)
challenges = client.get_all_challenges()
challenges.sort(key=get_timestamp)
total_completed = 0  # Variable to track total completed challenges
 
print(f"[i] Retrieved {len(challenges)} challenges")
 
def normalise_title(title):
    # Strip accents, smart quotes, etc.
    title = unicodedata.normalize("NFKD", title)
    # Replace fancy punctuation
    title = title.replace("’", "'").replace("`", "'").replace("–", "-").replace("“", '"').replace("”", '"')
    # Remove excess whitespace
    title = re.sub(r'\s+', ' ', title)
    return title.strip().lower()
 
# Group challenges by their categories
for challenge in challenges:
    categories[challenge['category_name']].append(challenge)
    if challenge['is_owned']:
        total_completed += 1  # Increment total completed if challenge is owned
 
# Print out the grouped challenges with the number of completed and total challenges in each category
for category, grouped_challenges in categories.items():
    total = len(grouped_challenges)
    completed = sum(1 for challenge in grouped_challenges if challenge['is_owned'])  # Count completed challenges
    
    res = ea.search_note(
        search=f"note.title %= '^{category}*'",
        ancestorNoteId=trilium_challenges_folder,
        ancestorDepth='eq1',
        limit=1,
        fastSearch=True,
    )
    
    catId = ""
    if res['results']:
        matched_note = res['results'][0]
        existing_category = " - ".join(matched_note['title'].split(" - ")[:-1])  # Extract category portion
        if existing_category.strip().lower() == category.lower():
            ea.patch_note(
                noteId=matched_note['noteId'], 
                title=f"{category} - {completed} / {total}"
            )
            catId = matched_note['noteId']
            print(f"[i] Updated category: {category} - ({completed}/{total})")
    else:
        new_note = ea.create_note(
            parentNoteId=trilium_challenges_folder,
            type="text",
            title=f"{category} - {completed} / {total}",
            content=" ",
        )
        catId = new_note['note']['noteId']
        print(f"[+] Created category: {catId} {category} - ({completed}/{total})")
    
    for challenge in grouped_challenges:
        #print(f" - ID: {challenge['id']}, Name: {challenge['name']}, Difficulty: {challenge['difficulty']}")
        escaped_name = challenge['name'].replace("'", "\\'").replace(",", "\\,")
        res2 = ea.search_note(
            search=f"note.title = '{escaped_name}'",
            ancestorNoteId=catId,
            ancestorDepth='eq1',
            orderBy=["title"],
            limit=1,
            fastSearch=True,
        )
 
        if res2['results']:
            note_title = normalise_title(res2['results'][0]['title'])
            challenge_title = normalise_title(challenge['name'])
 
            if note_title == challenge_title:
 
                for attribute in res2['results'][0]['attributes']:
                    if attribute['name'] == "Difficulty":
                        ea.patch_attribute(attributeId=attribute['attributeId'], value=challenge['difficulty'])
                    if attribute['name'] == "Released":
                        release_str = challenge['release_date']
                        release_date = datetime.strptime(release_str, "%Y-%m-%dT%H:%M:%S.%fZ")
                        formatted_release_date = release_date.strftime("%d %B %Y")
                        ea.patch_attribute(attributeId=attribute['attributeId'], value=formatted_release_date)
                    if attribute['name'] == "Solved":
                        if challenge['is_owned']:
                            ea.patch_attribute(attributeId=attribute['attributeId'], value="done")
                        else:
                            ea.patch_attribute(attributeId=attribute['attributeId'], value=" ")
                    if attribute['name'] == "cssClass":
                        if challenge['is_owned']:
                            ea.patch_attribute(attributeId=attribute['attributeId'], value="done")
                        else:
                            ea.patch_attribute(attributeId=attribute['attributeId'], value="todo")
 
        else: # doesnt already exist, create page
            release_str = challenge['release_date']
            release_date = datetime.strptime(release_str, "%Y-%m-%dT%H:%M:%S.%fZ")
            formatted_release_date = release_date.strftime("%d %B %Y")
            new_note = ea.create_note(
                parentNoteId=catId,
                type="text",
                title=challenge['name'],
                content=" ",
            )
            ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="relation", name="template", value=trilium_challenges_template_id)
            ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Difficulty", value=challenge['difficulty'])
            ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Released", value=formatted_release_date)
            if challenge['is_owned']:
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="cssClass", value="done")
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Solved", value="done")
            else:
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="cssClass", value="todo")
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Solved", value=" ")
 
            print(f"[+] created ID: {challenge['id']}, Name: {challenge['name']}, Difficulty: {challenge['difficulty']}")
 
print("[+] updating folder name                                          ")
ea.patch_note(noteId=trilium_challenges_folder,title="Challenges - "+str(total_completed)+" / "+str(len(challenges)))