#!/usr/bin/env python3
import os
import re
import sys
import glob

OLD_DIR = os.path.dirname(os.path.abspath(__file__))
CSS_FILENAME = "archive-banner.css"
MARKER_COMMENT = "<!-- ARCHIVE-BANNER -->"
BANNER_HTML = """{marker}
<link rel="stylesheet" href="{css_path}">
<div class="archive-banner">
  This is an archived website. <a href="http://erictaylor.me">Return to erictaylor.me</a>
  <button class="archive-banner-close" onclick="this.parentElement.style.display='none'">&times;</button>
</div>
{marker_end}"""


def get_css_relative_path(html_file_path):
    html_dir = os.path.dirname(html_file_path)
    rel_dir = os.path.relpath(OLD_DIR, html_dir)
    if rel_dir == ".":
        return CSS_FILENAME
    return rel_dir + "/" + CSS_FILENAME


def find_html_files(base_dir):
    html_files = []
    for root, dirs, files in os.walk(base_dir):
        for f in files:
            if f.lower().endswith((".html", ".htm")):
                html_files.append(os.path.join(root, f))
    return html_files


def inject_banner(filepath):
    with open(filepath, "r", encoding="utf-8", errors="replace") as f:
        content = f.read()

    if MARKER_COMMENT in content:
        return False

    css_path = get_css_relative_path(filepath)
    banner = BANNER_HTML.format(
        marker=MARKER_COMMENT, css_path=css_path, marker_end="<!-- /ARCHIVE-BANNER -->"
    )

    body_match = re.search(r"(<body[^>]*>)", content, re.IGNORECASE)
    if body_match:
        insert_pos = body_match.end()
        new_content = content[:insert_pos] + "\n" + banner + "\n" + content[insert_pos:]
    else:
        new_content = banner + "\n" + content

    with open(filepath, "w", encoding="utf-8") as f:
        f.write(new_content)
    return True


def remove_banner(filepath):
    with open(filepath, "r", encoding="utf-8", errors="replace") as f:
        content = f.read()

    if MARKER_COMMENT not in content:
        return False

    pattern = re.compile(
        r"\n?\s*"
        + re.escape(MARKER_COMMENT)
        + r".*?"
        + re.escape("<!-- /ARCHIVE-BANNER -->")
        + r"\n?\s*",
        re.DOTALL,
    )
    new_content = pattern.sub("\n", content)

    with open(filepath, "w", encoding="utf-8") as f:
        f.write(new_content)
    return True


def main():
    if len(sys.argv) > 1 and sys.argv[1] == "--undo":
        mode = "undo"
    else:
        mode = "inject"

    html_files = find_html_files(OLD_DIR)

    if mode == "undo":
        print(f"Removing archive banner from files in: {OLD_DIR}")
        removed = 0
        for fp in html_files:
            if remove_banner(fp):
                print(f"  Removed: {os.path.relpath(fp, OLD_DIR)}")
                removed += 1
        print(f"\nDone. Removed banner from {removed} files.")
    else:
        print(f"Injecting archive banner into files in: {OLD_DIR}")
        injected = 0
        skipped = 0
        for fp in html_files:
            if inject_banner(fp):
                print(f"  Injected: {os.path.relpath(fp, OLD_DIR)}")
                injected += 1
            else:
                skipped += 1
        print(
            f"\nDone. Injected {injected} files, skipped {skipped} (already have banner)."
        )


if __name__ == "__main__":
    main()