import os
import json
import requests
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_sherlocks_folder = config['trilium_sherlocks_folder']
trilium_sherlocks_template_id = config['trilium_sherlocks_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 sherlocks info")
categories = defaultdict(list)
sherlocks = client.get_all_sherlocks()
sherlocks.sort(key=get_timestamp)
total_completed = 0  # Variable to track total completed sherlocks

print(f"[i] Retrieved {len(sherlocks)} sherlocks")

# Group sherlocks by their categories
for sherlock in sherlocks:
    categories[sherlock['category_name']].append(sherlock)
    if sherlock['progress'] == 100:
        total_completed += 1  # Increment total completed if challenge is owned

# Print out the grouped sherlocks with the number of completed and total sherlocks in each category
for category, grouped_sherlocks in categories.items():
    total = len(grouped_sherlocks)
    completed = sum(1 for sherlock in grouped_sherlocks if sherlock['progress'] == 100)  # Count completed challenges
    
    res = ea.search_note(
        search=f"note.title %= '{category}*'",
        ancestorNoteId=trilium_sherlocks_folder,
        ancestorDepth='eq1',
        limit=1,
        fastSearch=True,
    )
    
    catId = ""
    if res['results'] and res['results'][0]['title'].split(' - ')[0].strip().lower() == category.lower():
	# page exists - lets check if the details have changed
        ea.patch_note(noteId=res['results'][0]['noteId'], title=category+" - "+str(completed)+" / "+str(total))
        catId = res['results'][0]['noteId']
        print(f"[i] updated category: {category} - ({completed}/{total})")
    else:
        new_note = ea.create_note(
            parentNoteId=trilium_sherlocks_folder,
            type="text",
            title=category+" - "+str(completed)+" / "+str(total),
            content=" ",
        )
        catId = new_note['note']['noteId']
        print(f"[+] created category: {catId} {category} - ({completed}/{total})")
    
    for sherlock in grouped_sherlocks:
        #print(f" - ID: {challenge['id']}, Name: {challenge['name']}, Difficulty: {challenge['difficulty']}")
        res2 = ea.search_note(
            search=f"{sherlock['name']}",
            ancestorNoteId=catId,
            ancestorDepth='eq1',
            limit=1,
            fastSearch=True,
        )
        # already exists update the values
        if res2['results'] and res2['results'][0]['title'].lower() == sherlock['name'].lower():
            print(f"[i] found ID: {sherlock['id']}, Name: {sherlock['name']}, Difficulty: {sherlock['difficulty']}")
    
            for attribute in res2['results'][0]['attributes']:
                if attribute['name'] == "Difficulty":
                    ea.patch_attribute(attributeId=attribute['attributeId'], value=sherlock['difficulty'])
        
                if attribute['name'] == "Released":
                    release_str = sherlock['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":
                    ea.patch_attribute(attributeId=attribute['attributeId'], value=str(sherlock['progress']))
        
                if attribute['name'] == "cssClass":
                    # Determine the cssClass value based on the progress
                    if sherlock['progress'] == 0:
                        ea.patch_attribute(attributeId=attribute['attributeId'], value="todo")
                    elif sherlock['progress'] == 100:
                        ea.patch_attribute(attributeId=attribute['attributeId'], value="done")
                    elif 0 < sherlock['progress'] < 100:
                        ea.patch_attribute(attributeId=attribute['attributeId'], value="inprogress")


        else: # doesnt already exist, create page
            release_str = sherlock['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=sherlock['name'],
                content=" ",
            )
            ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="relation", name="template", value=trilium_sherlocks_template_id)
            ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Difficulty", value=sherlock['difficulty'])
            ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Released", value=formatted_release_date)
            ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="Solved", value=str(sherlock['progress']))
            if sherlock['progress'] == 0:
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="cssClass", value="todo")
            elif sherlock['progress'] == 100:
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="cssClass", value="done")
            elif 0 < sherlock['progress'] < 100:
                ea.create_attribute(attributeId=None, isInheritable=False, noteId=new_note['note']['noteId'], type="label", name="cssClass", value="inprogress")

            print(f"[+] created ID: {sherlock['id']}, Name: {sherlock['name']}, Difficulty: {sherlock['difficulty']}")

print("[+] updating folder name                                          ")
ea.patch_note(noteId=trilium_sherlocks_folder,title="Sherlocks - "+str(total_completed)+" / "+str(len(sherlocks)))