#!/usr/bin/env python3
"""Quick contrast audit across all onepager sites.

Extracts CSS custom properties for backgrounds and text colors,
computes WCAG contrast ratio for likely text-on-bg pairs, flags
violations.

Modes:
  --dark   (default) Audit :root / [data-theme="dark"] palette
  --light            Audit [data-theme="light"] palette
  --both             Audit both, separately

Light-mode audit picks up site-specific [data-theme="light"] blocks.
Sites without their own light block fall through to shared/css/theme.css
(which is checked separately and used as a fallback).
"""
import argparse
import re
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
SITES_DIR = ROOT / "sites"
SHARED_THEME_CSS = ROOT / "shared" / "css" / "theme.css"

HEX_RE = re.compile(r"#([0-9a-fA-F]{3,8})\b")
VAR_DECL_RE = re.compile(r"--([\w-]+)\s*:\s*([^;]+);")

# Match a CSS rule block by selector. Returns body of the block.
def block_body(css, selector_pattern):
    """Return list of bodies from CSS rules whose selector matches the regex."""
    bodies = []
    # Walk char-by-char to handle nested braces correctly
    for match in re.finditer(selector_pattern, css):
        start = match.end()
        # Find opening brace right after selector
        idx = css.find("{", start)
        if idx == -1:
            continue
        depth = 1
        i = idx + 1
        while i < len(css) and depth > 0:
            if css[i] == "{":
                depth += 1
            elif css[i] == "}":
                depth -= 1
            i += 1
        if depth == 0:
            bodies.append(css[idx + 1 : i - 1])
    return bodies


def hex_to_rgb(h):
    h = h.lstrip("#")
    if len(h) == 3:
        h = "".join(c * 2 for c in h)
    if len(h) == 8:
        h = h[:6]
    if len(h) != 6:
        return None
    try:
        return tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
    except ValueError:
        return None


def relative_luminance(rgb):
    def channel(c):
        c /= 255
        return c / 12.92 if c <= 0.03928 else ((c + 0.055) / 1.055) ** 2.4
    r, g, b = (channel(c) for c in rgb)
    return 0.2126 * r + 0.7152 * g + 0.0722 * b


def contrast_ratio(rgb1, rgb2):
    l1 = relative_luminance(rgb1)
    l2 = relative_luminance(rgb2)
    lighter, darker = max(l1, l2), min(l1, l2)
    return (lighter + 0.05) / (darker + 0.05)


BG_KEYS = ("bg", "background", "surface")
TEXT_KEYS = ("text", "fg", "foreground", "color", "muted", "dim", "subtle", "secondary")


def is_bg_var(name):
    n = name.lower()
    return any(k in n for k in BG_KEYS) and "border" not in n


def is_text_var(name):
    n = name.lower()
    return any(k in n for k in TEXT_KEYS) and "border" not in n and "bg" not in n.split("-")[0]


def parse_vars(css_text):
    """Extract CSS custom property name -> RGB triplet, ignoring non-hex values."""
    vars_ = {}
    for m in VAR_DECL_RE.finditer(css_text):
        name, val = m.group(1), m.group(2).strip()
        hm = HEX_RE.search(val)
        if hm:
            rgb = hex_to_rgb(hm.group(0))
            if rgb:
                vars_[name] = rgb
    return vars_


def get_style_blocks(html):
    """Return concatenated CSS from all <style> blocks in HTML."""
    blocks = re.findall(r"<style.*?>(.*?)</style>", html, re.DOTALL | re.IGNORECASE)
    return "\n".join(blocks)


def shared_light_defaults():
    """Read the shared light defaults so sites without overrides get audited
    against the actual palette they'll receive at runtime."""
    if not SHARED_THEME_CSS.exists():
        return {}
    css = SHARED_THEME_CSS.read_text(errors="ignore")
    bodies = block_body(css, r'\[data-theme="light"\]')
    merged = "\n".join(bodies)
    return parse_vars(merged)


def audit_palette(vars_, mode_label):
    """Given a {var_name: rgb} palette, return findings dict or None."""
    bg_vars = {n: c for n, c in vars_.items() if is_bg_var(n)}
    text_vars = {n: c for n, c in vars_.items() if is_text_var(n)}

    if not bg_vars or not text_vars:
        return None

    if mode_label == "dark":
        primary_bg_name = min(bg_vars, key=lambda n: relative_luminance(bg_vars[n]))
        bg_rgb = bg_vars[primary_bg_name]
        bg_lum = relative_luminance(bg_rgb)
        if bg_lum > 0.1:
            return None  # not dark enough to be the dark mode
    else:  # light
        primary_bg_name = max(bg_vars, key=lambda n: relative_luminance(bg_vars[n]))
        bg_rgb = bg_vars[primary_bg_name]
        bg_lum = relative_luminance(bg_rgb)
        if bg_lum < 0.5:
            return None  # not light enough

    findings = []
    for tname, trgb in text_vars.items():
        ratio = contrast_ratio(trgb, bg_rgb)
        if ratio < 4.5:
            findings.append((tname, trgb, ratio))

    if not findings:
        return None
    return {
        "bg_name": primary_bg_name,
        "bg_rgb": bg_rgb,
        "findings": findings,
    }


def audit(site, mode, shared_light):
    html_path = SITES_DIR / site / "index.html"
    if not html_path.exists():
        return None
    html = html_path.read_text(errors="ignore")
    css = get_style_blocks(html)
    if not css:
        return None

    if mode == "dark":
        # :root and [data-theme="dark"] both contribute to dark palette
        bodies = block_body(css, r":root") + block_body(css, r'\[data-theme="dark"\]')
        merged = "\n".join(bodies) if bodies else css
        vars_ = parse_vars(merged)
        return audit_palette(vars_, "dark")

    if mode == "light":
        bodies = block_body(css, r'\[data-theme="light"\]')
        if bodies:
            site_vars = parse_vars("\n".join(bodies))
        else:
            site_vars = {}
        # Site overrides win. Fall back to shared defaults for missing vars.
        # AND fall back to :root for any vars the site only defines once.
        root_bodies = block_body(css, r":root")
        root_vars = parse_vars("\n".join(root_bodies))
        merged = dict(root_vars)
        merged.update(shared_light)
        merged.update(site_vars)
        return audit_palette(merged, "light")

    return None


def print_results(mode, results, total):
    label = "Light" if mode == "light" else "Dark"
    print(f"\n=== {label} mode audit ===")
    print(f"Sites with sub-AA text on {label.lower()} bg: {len(results)}/{total}\n")
    for r in results:
        bg = r["bg_rgb"]
        print(f"{r['site']}  (bg --{r['bg_name']} = #{bg[0]:02x}{bg[1]:02x}{bg[2]:02x})")
        for name, rgb, ratio in sorted(r["findings"], key=lambda x: x[2]):
            flag = "FAIL" if ratio < 3.0 else ("WEAK" if ratio < 4.5 else "OK")
            print(f"  {flag:4} ratio {ratio:5.2f}  --{name:24}  #{rgb[0]:02x}{rgb[1]:02x}{rgb[2]:02x}")
        print()
    fails = [r for r in results if any(ratio < 3.0 for _, _, ratio in r["findings"])]
    print(f"  {label}-bg sites with at least one FAIL (<3:1): {len(fails)}")
    print(f"  {label}-bg sites with WEAK (<4.5 but >=3): {len(results) - len(fails)}")


def run(mode):
    sites = sorted(p.name for p in SITES_DIR.iterdir() if p.is_dir())
    shared_light = shared_light_defaults() if mode in ("light", "both") else {}

    if mode in ("dark", "both"):
        results = []
        for site in sites:
            r = audit(site, "dark", shared_light)
            if r:
                r["site"] = site
                results.append(r)
        print_results("dark", results, len(sites))

    if mode in ("light", "both"):
        results = []
        for site in sites:
            r = audit(site, "light", shared_light)
            if r:
                r["site"] = site
                results.append(r)
        print_results("light", results, len(sites))


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    grp = ap.add_mutually_exclusive_group()
    grp.add_argument("--dark", action="store_const", const="dark", dest="mode", help="Audit dark palette only (default)")
    grp.add_argument("--light", action="store_const", const="light", dest="mode", help="Audit light palette only")
    grp.add_argument("--both", action="store_const", const="both", dest="mode", help="Audit both palettes")
    args = ap.parse_args()
    mode = args.mode or "dark"
    run(mode)


if __name__ == "__main__":
    main()
