import os
import sys
import time
import random
import re
import requests
from urllib.parse import urljoin, urlparse
from dotenv import load_dotenv

def rewrite_image_paths(content, session, output_dir, base_url):
    def replace_html_img(match):
        return match.group(1) + download_image(match.group(2), session, output_dir, base_url) + match.group(3)

    def replace_markdown_img(match):
        return match.group(1) + download_image(match.group(2), session, output_dir, base_url) + match.group(3)

    content = re.sub(r'(<img\b[^>]*\bsrc=["\'])([^"\']+)(["\'][^>]*>)', replace_html_img, content, flags=re.I)
    return re.sub(r'(!\[[^\]]*\]\()([^)\s]+)(\))', replace_markdown_img, content)

def print_status(current, total, section, sleep_duration=0):
    percent = int(current / total * 100) if total else 0
    filled = int(20 * current / total) if total else 0
    bar = "#" * filled + "-" * (20 - filled)
    print(f"\r\033[K[{bar}] {current}/{total} - {percent}% | {section} | sleep {sleep_duration}s", end="", flush=True)


def download_image(url, session, output_dir, base_url):
    try:
        if not url.startswith('http'):
            if url.startswith('/'):
                url = urljoin(base_url, url)
            else:
                return url
                
        parsed_url = urlparse(url)
        filename = os.path.basename(parsed_url.path)
        if not filename:
            filename = f"image_{int(time.time())}.png"
            
        local_path = os.path.join(output_dir, filename)
        if os.path.exists(local_path):
            return local_path
            
        # Keep terminal output to the single progress line; failures print below.
        response = session.get(url, stream=True)
        response.raise_for_status()
        with open(local_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        return local_path
    except Exception as e:
        print(f"\n[-] Failed to download {url}: {e}")
        return url

def main():
    load_dotenv()
    cookie_string = os.getenv('HTB_COOKIE')
    user_agent = os.getenv('HTB_USER_AGENT', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36')
    
    if not cookie_string:
        print("Error: HTB_COOKIE not found in .env file.")
        sys.exit(1)
        
    if len(sys.argv) < 2:
        print("Usage: python htb_scraper.py https://academy.hackthebox.com/app/module/144")
        sys.exit(1)
        
    module_url = sys.argv[1]
    
    # Extract module ID
    # URLs typically: /app/module/144 or /module/144
    parts = module_url.strip('/').split('/')
    try:
        module_idx = parts.index('module')
        module_id = parts[module_idx + 1]
    except (ValueError, IndexError):
        print("Error: Could not extract module ID from URL. Make sure it contains '/module/<id>'")
        sys.exit(1)
    
    session = requests.Session()
    headers = {
        'accept': 'application/json',
        'referer': module_url,
        'user-agent': user_agent,
        'sec-ch-ua': '"Brave";v="149", "Chromium";v="149", "Not)A;Brand";v="24"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"Linux"',
        'sec-fetch-dest': 'empty',
        'sec-fetch-mode': 'cors',
        'sec-fetch-site': 'same-origin',
        'sec-gpc': '1',
        'Cookie': cookie_string
    }
    
    session.headers.update(headers)
    
    print_status(0, 0, f"module {module_id}", 0)
    api_base = "https://academy.hackthebox.com"
    sections_url = f"{api_base}/api/v3/modules/{module_id}/sections"
    
    res = session.get(sections_url)
    if res.status_code != 200:
        print(f"\n[-] API returned status {res.status_code}. Cookie invalid or expired.")
        sys.exit(1)
        
    data = res.json()
    # API v3 returns list of groups containing sections
    sections = []
    if isinstance(data, dict) and "data" in data:
        groups = data["data"]
    elif isinstance(data, list):
        groups = data
    else:
        groups = []
        
    for group in groups:
        if isinstance(group, dict) and "sections" in group:
            for s in group["sections"]:
                if isinstance(s, dict) and "id" in s:
                    sections.append(s)
                    
    if not sections:
        print("\n[-] No sections found in the API response.")
        sys.exit(1)
        
    print_status(0, len(sections), f"module {module_id}", 0)
    
    output_file = f"module_{module_id}.md"
    images_dir = f"images_module_{module_id}"
    os.makedirs(images_dir, exist_ok=True)
    
    with open(output_file, "w", encoding="utf-8") as f:
        f.write(f"# HTB Academy Module {module_id}\n\n")
        f.write(f"Source: {module_url}\n\n")
        
    for i, sec in enumerate(sections, 1):
        sec_id = sec['id']
        sec_title = sec.get('title', f"Section {sec_id}")
        
        delay = random.randint(3, 8)
        for remaining in range(delay, 0, -1):
            print_status(i, len(sections), sec_title, remaining)
            time.sleep(1)

        print_status(i, len(sections), sec_title, 0)
        content_url = f"{api_base}/api/v2/modules/{module_id}/sections/{sec_id}"
        
        sec_res = session.get(content_url)
        if sec_res.status_code != 200:
            print(f"\n[-] Failed to fetch {sec_title}. Status: {sec_res.status_code}")
            continue
            
        sec_data = sec_res.json()
        if isinstance(sec_data, dict) and "data" in sec_data:
            sec_data = sec_data["data"]
            
        html_content = sec_data.get("content", "")
        if not html_content:
            print(f"\n[-] No content found for {sec_title}.")
            continue
            
        # HTB's API content is already Markdown. Do not run it through an
        # HTML parser; XML examples inside fenced code blocks are real content.
        markdown_text = rewrite_image_paths(html_content, session, images_dir, api_base)
        
        with open(output_file, "a", encoding="utf-8") as f:
            f.write(f"\n\n---\n\n## {sec_title}\n\n")
            f.write(markdown_text)
            
    print(f"\r\033[K[+] Done! Saved to {output_file}")

if __name__ == "__main__":
    main()