#!/usr/bin/env python3
import os
import re
import shutil

# Paths
ROOT_DIR = os.path.abspath(os.path.dirname(__file__))
BUILD_DIR = os.path.join(ROOT_DIR, "build")

# Helper to normalize URL paths by removing './'
def normalize_url(url):
    url_parts = url.split('#', 1)
    path = url_parts[0]
    anchor = '#' + url_parts[1] if len(url_parts) > 1 else ''
    
    # Remove leading './'
    if path.startswith('./'):
        path = path[2:]
    return path, anchor

def determine_target_path(filename):
    """
    Determines where an HTML file should go, relative to the ROOT_DIR.
    """
    if filename.startswith("ch") and "-" in filename:
        # e.g., ch01-00-getting-started.html -> ch01/ch01-00-getting-started.html
        prefix = filename.split("-", 1)[0] # e.g. "ch01"
        return os.path.join(prefix, filename)
    elif filename.startswith("appendix-"):
        # e.g. appendix-01-keywords.html -> appendix/appendix-01-keywords.html
        return os.path.join("appendix", filename)
    else:
        # Standard pages (toc.html, print.html, index.html, 404.html, foreword.html, etc.) -> root
        return filename

def build_mappings():
    """
    Scans the build directory and maps old relative file paths to new relative target paths.
    """
    mapping = {}
    
    # Check all files in the build root
    for file in os.listdir(BUILD_DIR):
        if file.endswith(".html"):
            target_rel = determine_target_path(file)
            mapping[file] = target_rel
            
    return mapping

def clean_html_content(content):
    """
    Cleans the HTML document by extracting title and main content, 
    and stripping scripts, styles, and decorative Ferris images.
    """
    # 1. Extract title
    title = "Rust Book"
    title_match = re.search(r"<title>(.*?)</title>", content, re.IGNORECASE | re.DOTALL)
    if title_match:
        raw_title = title_match.group(1).strip()
        # Remove common suffix
        suffix = " - The Rust Programming Language"
        if raw_title.endswith(suffix):
            title = raw_title[:-len(suffix)].strip()
        else:
            title = raw_title

    # 2. Extract content from <main> or fall back to <body>
    main_match = re.search(r"<main(?:\s+[^>]*)?>(.*?)</main>", content, re.IGNORECASE | re.DOTALL)
    if main_match:
        body_content = main_match.group(1).strip()
    else:
        body_match = re.search(r"<body(?:\s+[^>]*)?>(.*?)</body>", content, re.IGNORECASE | re.DOTALL)
        if body_match:
            body_content = body_match.group(1).strip()
        else:
            body_content = content.strip()

    # 3. Strip nested scripts or styles
    body_content = re.sub(r"<script\b[^<]*(?:(?!</script>)<[^<]*)*</script>", "", body_content, flags=re.IGNORECASE | re.DOTALL)
    body_content = re.sub(r"<style\b[^<]*(?:(?!</style>)<[^<]*)*</style>", "", body_content, flags=re.IGNORECASE | re.DOTALL)

    # 4. Strip decorative Ferris image tags to prevent broken image links
    body_content = re.sub(r"<img\b[^>]*src=\"[^\"]*img/ferris/[^\"]*\"[^>]*>", "", body_content, flags=re.IGNORECASE)

    # 5. Construct clean minimal HTML structure
    return title, body_content

def update_links(html_content, current_rel_path, mapping):
    """
    Updates all relative href and src attributes to point to the new organized folder structure.
    """
    current_new_dir = os.path.dirname(mapping.get(current_rel_path, current_rel_path))
    
    def replace_link(match):
        attr = match.group(1) # href or src
        url = match.group(2)  # the raw url
        
        path, anchor = normalize_url(url)
        
        # Check mapping for HTML files
        if path in mapping:
            target_new_path = mapping[path]
            # Calculate the relative path from current_new_dir to target_new_path
            new_rel_path = os.path.relpath(target_new_path, current_new_dir)
            return f'{attr}="{new_rel_path}{anchor}"'
            
        # Check assets (like images)
        if path.startswith("img/"):
            # Image paths are originally relative to root. Move relative path.
            # e.g., if current is in ch01/, relpath to img/trpl01-01.png is ../img/trpl01-01.png
            new_rel_path = os.path.relpath(path, current_new_dir)
            return f'{attr}="{new_rel_path}{anchor}"'
            
        return match.group(0)

    # Regex matching: href="..." or src="..."
    html_content = re.sub(r'(href|src)="([^"]*)"', replace_link, html_content)
    return html_content

def clean_and_organize():
    print(f"Reading from raw build: {BUILD_DIR}")
    
    # 1. Map all HTML files to their new organized destinations
    mapping = build_mappings()
    if not mapping:
        print("No HTML files found to process.")
        return
        
    # 2. Identify directories to clean in the workspace root
    # We will delete any existing organized folders so we don't have stale files
    dirs_to_clean = set(os.path.dirname(p) for p in mapping.values() if os.path.dirname(p))
    dirs_to_clean.add("img") # we will recreate root img directory
    
    for d in dirs_to_clean:
        d_path = os.path.join(ROOT_DIR, d)
        if os.path.exists(d_path):
            print(f"Cleaning existing organized folder: {d}/")
            shutil.rmtree(d_path)
            
    # Also clean existing loose HTML files in the root that are in our mapping
    for file_rel in mapping.values():
        if os.path.dirname(file_rel) == "":
            root_file_path = os.path.join(ROOT_DIR, file_rel)
            if os.path.exists(root_file_path):
                print(f"Removing old root file: {file_rel}")
                os.remove(root_file_path)

    # 3. Process and write cleaned HTML files to organized paths
    for old_rel, new_rel in mapping.items():
        src_path = os.path.join(BUILD_DIR, old_rel)
        dest_path = os.path.join(ROOT_DIR, new_rel)
        
        # Ensure destination folder exists
        os.makedirs(os.path.dirname(dest_path), exist_ok=True)
        
        with open(src_path, "r", encoding="utf-8") as f:
            raw_content = f.read()
            
        # Clean the HTML content
        title, clean_body = clean_html_content(raw_content)
        
        # Update hrefs and srcs
        updated_content = update_links(clean_body, old_rel, mapping)
        
        clean_html = f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>{title}</title>
</head>
<body>
{updated_content}
</body>
</html>
"""
        
        with open(dest_path, "w", encoding="utf-8") as f:
            f.write(clean_html)
            
        print(f"Saved: {new_rel}")

    # 4. Copy conceptual images inside img/ (excluding img/ferris/)
    src_img_dir = os.path.join(BUILD_DIR, "img")
    dest_img_dir = os.path.join(ROOT_DIR, "img")
    
    if os.path.exists(src_img_dir):
        os.makedirs(dest_img_dir, exist_ok=True)
        for item in os.listdir(src_img_dir):
            src_item = os.path.join(src_img_dir, item)
            dest_item = os.path.join(dest_img_dir, item)
            
            if os.path.isdir(src_item):
                # Skip ferris directory
                if item == "ferris":
                    continue
                shutil.copytree(src_item, dest_item)
            else:
                shutil.copy2(src_item, dest_item)
        print("Copied content diagrams to root img/ directory.")

    # 5. Delete raw build folder completely
    print(f"Removing build/ folder...")
    shutil.rmtree(BUILD_DIR)
    print("Reorganization and cleanup completed successfully.")

if __name__ == "__main__":
    if not os.path.exists(BUILD_DIR):
        print(f"Error: Build directory {BUILD_DIR} does not exist.")
        print("If you have already organized the files, please rebuild first if you want to clean again.")
    else:
        clean_and_organize()
