#!/usr/bin/env python3
"""Replace values of existing SET variables in a Windows batch file.

Examples
--------
Overwrite the input file::

    python replace_bat_set.py make.bat tts=pyttsx3 voices=Haruka resolution=1080

Write to another file::

    python replace_bat_set.py make.bat tts=voicevox -o make_voicevox.bat

Only variables already defined by ordinary ``set`` statements are replaced.
Statements such as ``set /a`` and ``set /p`` are ignored.
"""

from __future__ import annotations

import argparse
import codecs
import re
import shutil
import sys
from dataclasses import dataclass
from pathlib import Path


# Matches both:
#   set "key=value"
#   set key=value
# while excluding SET /A and SET /P.
SET_RE = re.compile(
    r"^(?P<prefix>[ \t]*set[ \t]+)(?!/[aApP](?:[ \t]|$))"
    r'(?:(?P<quote>\")(?P<qkey>[^=\r\n\"]+?)=(?P<qvalue>[^\r\n]*)\"'
    r"|(?P<ukey>[^=\r\n]+?)=(?P<uvalue>[^\r\n]*))"
    r"(?P<ending>\r?\n|$)$",
    re.IGNORECASE,
)


@dataclass(frozen=True)
class TextFormat:
    encoding: str
    bom: bytes


def detect_text_format(data: bytes) -> TextFormat:
    """Detect common batch-file encodings without altering the original BOM."""
    if data.startswith(codecs.BOM_UTF8):
        return TextFormat("utf-8", codecs.BOM_UTF8)
    if data.startswith(codecs.BOM_UTF16_LE):
        return TextFormat("utf-16-le", codecs.BOM_UTF16_LE)
    if data.startswith(codecs.BOM_UTF16_BE):
        return TextFormat("utf-16-be", codecs.BOM_UTF16_BE)

    for encoding in ("utf-8", "cp932"):
        try:
            data.decode(encoding)
            return TextFormat(encoding, b"")
        except UnicodeDecodeError:
            pass

    raise UnicodeError("文字コードを判定できません。UTF-8 または CP932 の .bat を使用してください。")


def parse_assignments(items: list[str]) -> dict[str, tuple[str, str]]:
    """Return case-insensitive mapping: normalized key -> (original key, value)."""
    assignments: dict[str, tuple[str, str]] = {}
    for item in items:
        if "=" not in item:
            raise ValueError(f"引数は key=value 形式で指定してください: {item!r}")
        key, value = item.split("=", 1)
        key = key.strip()
        if not key:
            raise ValueError(f"変数名が空です: {item!r}")
        assignments[key.casefold()] = (key, value)
    return assignments


def replace_set_values(text: str, assignments: dict[str, tuple[str, str]]) -> tuple[str, dict[str, int]]:
    counts = {key: 0 for key in assignments}
    output: list[str] = []

    for line in text.splitlines(keepends=True):
        match = SET_RE.match(line)
        if not match:
            output.append(line)
            continue

        quoted = match.group("quote") is not None
        key_text = match.group("qkey") if quoted else match.group("ukey")
        assert key_text is not None

        # Whitespace immediately before '=' belongs to the key syntactically,
        # but ordinary batch variables normally should not include it.
        lookup_key = key_text.strip().casefold()
        replacement = assignments.get(lookup_key)
        if replacement is None:
            output.append(line)
            continue

        new_value = replacement[1]
        counts[lookup_key] += 1
        prefix = match.group("prefix")
        ending = match.group("ending")

        if quoted:
            output.append(f'{prefix}"{key_text}={new_value}"{ending}')
        else:
            output.append(f"{prefix}{key_text}={new_value}{ending}")

    return "".join(output), counts


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="バッチファイル内の既存の set 変数を key=value 引数で置換します。"
    )
    parser.add_argument("batch_file", type=Path, help="入力する .bat ファイル")
    parser.add_argument("assignments", nargs="+", metavar="key=value", help="置換する変数と値")
    parser.add_argument("-o", "--output", type=Path, help="出力先。省略時は入力ファイルを上書き")
    parser.add_argument(
        "--backup",
        action="store_true",
        help="上書き時に元ファイルを .bak として保存",
    )
    parser.add_argument(
        "--ignore-missing",
        action="store_true",
        help="指定した変数が set で定義されていなくてもエラーにしない",
    )
    return parser


def main() -> int:
    args = build_parser().parse_args()

    try:
        assignments = parse_assignments(args.assignments)
        source = args.batch_file
        destination = args.output or source

        data = source.read_bytes()
        fmt = detect_text_format(data)
        payload = data[len(fmt.bom) :]
        text = payload.decode(fmt.encoding)

        replaced_text, counts = replace_set_values(text, assignments)
        missing = [assignments[key][0] for key, count in counts.items() if count == 0]
        if missing and not args.ignore_missing:
            print(
                "[ERROR] set で定義されていない変数があります: " + ", ".join(missing),
                file=sys.stderr,
            )
            return 2

        if destination == source and args.backup:
            backup = source.with_suffix(source.suffix + ".bak")
            shutil.copy2(source, backup)
            print(f"[BACKUP] {backup}")

        encoded = fmt.bom + replaced_text.encode(fmt.encoding)
        destination.parent.mkdir(parents=True, exist_ok=True)
        destination.write_bytes(encoded)

        for normalized_key, count in counts.items():
            original_key, value = assignments[normalized_key]
            if count:
                print(f"[OK] {original_key}={value} ({count} 箇所)")
            else:
                print(f"[SKIP] {original_key}: 定義なし")
        print(f"[OUTPUT] {destination}")
        return 0

    except (OSError, UnicodeError, ValueError) as exc:
        print(f"[ERROR] {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())