#!/usr/bin/env python3
"""Safely add the Codex usage module to an existing Waybar JSONC config."""

from __future__ import annotations

import argparse
import re
import shutil
import sys
from datetime import datetime
from pathlib import Path

MODULE_NAME = "custom/codex-usage"
MODULE_BLOCK = '''  "custom/codex-usage": {
    "exec": "codex-waybar",
    "return-type": "json",
    "interval": 60,
    "tooltip": true
  }
'''
CSS_BLOCK = """

/* Codex usage module */
#custom-codex-usage {
  min-width: 12px;
  margin: 0 7.5px;
}

#custom-codex-usage.warning {
  color: #d8a657;
}

#custom-codex-usage.critical,
#custom-codex-usage.unavailable {
  color: #a55555;
}

#custom-codex-usage.stale {
  opacity: 0.65;
}
"""


class ConfigureError(RuntimeError):
    pass


def backup(path: Path) -> Path:
    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    destination = path.with_name(f"{path.name}.bak.{stamp}")
    shutil.copy2(path, destination)
    return destination


def insert_module_reference(text: str, section: str) -> str:
    pattern = re.compile(rf'("modules-{section}"\s*:\s*\[)(.*?)(\])', re.DOTALL)
    matches = list(pattern.finditer(text))
    if len(matches) != 1:
        raise ConfigureError(f'Could not uniquely locate "modules-{section}"')
    match = matches[0]
    if f'"{MODULE_NAME}"' in match.group(2):
        return text
    return text[: match.end(1)] + f'"{MODULE_NAME}", ' + text[match.end(1) :]


def insert_module_block(text: str) -> str:
    if re.search(rf'"{re.escape(MODULE_NAME)}"\s*:', text):
        return text
    close = text.rfind("}")
    if close < 0:
        raise ConfigureError("Could not locate the top-level closing brace")
    before = text[:close].rstrip()
    comma = "" if before.endswith("{") else ","
    return f"{before}{comma}\n{MODULE_BLOCK}{text[close:]}"


def configure(config_path: Path, style_path: Path, section: str) -> tuple[Path | None, Path | None]:
    if not config_path.is_file():
        raise ConfigureError(f"Waybar config not found: {config_path}")
    original = config_path.read_text()
    updated = insert_module_block(insert_module_reference(original, section))
    config_backup = None
    if updated != original:
        config_backup = backup(config_path)
        config_path.write_text(updated)

    style_backup = None
    if style_path.exists():
        style = style_path.read_text()
        if "#custom-codex-usage" not in style:
            style_backup = backup(style_path)
            style_path.write_text(style.rstrip() + CSS_BLOCK)
    else:
        style_path.parent.mkdir(parents=True, exist_ok=True)
        style_path.write_text(CSS_BLOCK.lstrip())
    return config_backup, style_backup


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    default_dir = Path.home() / ".config" / "waybar"
    parser.add_argument("--section", choices=("left", "center", "right"), required=True)
    parser.add_argument("--config", type=Path, default=default_dir / "config.jsonc")
    parser.add_argument("--style", type=Path, default=default_dir / "style.css")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        config_backup, style_backup = configure(args.config, args.style, args.section)
    except (OSError, ConfigureError) as exc:
        print(f"codex-waybar-configure: {exc}", file=sys.stderr)
        return 1
    print(f"Configured {args.config}")
    if config_backup:
        print(f"Backup: {config_backup}")
    if style_backup:
        print(f"Backup: {style_backup}")
    return 0


if __name__ == "__main__":
    sys.exit(main())