#!/usr/bin/env python3
"""Patch Inovelli ZHA quirks to add KeepAlive client cluster on endpoint 1."""

import argparse
import importlib.util
import re
import shutil
from datetime import datetime, timezone
from pathlib import Path

MODELS = (
    "VZM31-SN",
    "VZM32-SN",
    "VZM35-SN",
    "VZM30-SN",
    "VZM36",
)

PATCH = ".adds(KeepAlive, cluster_type=ClusterType.Client, endpoint_id=1)"
IMPORTS = (
    "from zigpy.zcl import ClusterType\n"
    "from zigpy.zcl.clusters.general import KeepAlive\n"
)

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
    "--file", type=Path,
    help="Patch one explicitly specified quirk file"
)
parser.add_argument(
    "--check", action="store_true",
    help="Show proposed changes without writing files"
)
args = parser.parse_args()

if args.file:
    targets = [args.file]
else:
    spec = importlib.util.find_spec("zhaquirks.inovelli")
    if spec is None or not spec.submodule_search_locations:
        parser.error(
            "Cannot locate zhaquirks.inovelli. "
            "Run in the Home Assistant Core Python environment."
        )
    directory = Path(next(iter(spec.submodule_search_locations)))
    targets = sorted(directory.glob("*.py"))

builder_pattern = re.compile(
    r'InovelliQuirkBuilder\(\s*["\']Inovelli["\']\s*,\s*'
    r'["\']([^"\']+)["\']\s*\)'
)

# Match a builder chain that ends with its .add_to_registry() call.
registry_pattern = re.compile(
    r"(?m)^[ \t]*\.add_to_registry\(\)"
)


def supported(model):
    return model in MODELS or model.startswith("VZM36")


def patch_source(source, filename):
    matches = list(builder_pattern.finditer(source))
    edits = []
    results = []

    for match in matches:
        model = match.group(1)
        if not supported(model):
            continue

        # Locate the end of this particular builder chain.
        end_match = registry_pattern.search(source, match.end())
        if end_match is None:
            raise ValueError(
                f"{filename}: no .add_to_registry() for {model}"
            )

        # Refuse ambiguous layouts where another builder starts first.
        next_builder = builder_pattern.search(source, match.end())
        if next_builder and next_builder.start() < end_match.start():
            raise ValueError(
                f"{filename}: ambiguous builder layout for {model}"
            )

        chain = source[match.start():end_match.start()]
        if PATCH in chain:
            results.append((model, "already patched"))
            continue

        # Keep indentation consistent with existing chain calls.
        indent_match = re.search(
            r"(?m)^([ \t]*)\.[A-Za-z_]+\(",
            chain
        )
        if indent_match is None:
            raise ValueError(
                f"{filename}: cannot determine chain indentation for {model}"
            )

        indent = indent_match.group(1)
        edits.append((
            end_match.start(),
            f"{indent}{PATCH}\n"
        ))
        results.append((model, "patch needed"))

    if not results:
        return source, []

    # Apply edits backwards so earlier positions remain valid.
    patched = source
    for position, insertion in sorted(edits, reverse=True):
        patched = patched[:position] + insertion + patched[position:]

    if edits:
        if "from zigpy.zcl import ClusterType" not in patched:
            patched = "from zigpy.zcl import ClusterType\n" + patched
        if "from zigpy.zcl.clusters.general import KeepAlive" not in patched:
            patched = (
                "from zigpy.zcl.clusters.general import KeepAlive\n"
                + patched
            )

    compile(patched, str(filename), "exec")
    return patched, results


changed = 0
found = set()
timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")

for target in targets:
    if not target.is_file():
        parser.error(f"Not found: {target}")

    source = target.read_text()

    try:
        patched, results = patch_source(source, target)
    except ValueError as exc:
        parser.error(str(exc))

    if not results:
        continue

    print(f"\nFile: {target}")
    for model, status in results:
        print(f"  {model}: {status}")
        found.add(model)

    if patched == source:
        continue

    changed += 1
    if args.check:
        continue

    backup = target.with_name(
        target.name + ".backup-" + timestamp
    )
    shutil.copy2(target, backup)
    target.write_text(patched)
    print(f"  Backup: {backup}")
    print("  Patched successfully")

print("\nSummary")
print(f"Models found: {', '.join(sorted(found)) or 'none'}")
print(f"Files {'requiring changes' if args.check else 'changed'}: {changed}")

missing = [
    model for model in MODELS
    if not any(
        item == model or
        (model == "VZM36" and item.startswith("VZM36"))
        for item in found
    )
]
if missing:
    print(f"Models not found: {', '.join(missing)}")

if args.check:
    print("Check complete. No files changed.")
elif changed:
    print("Restart Home Assistant Core to load the patched quirks.")
    print("Home Assistant updates may overwrite these changes.")