import json
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument(
    "-s",
    "--storage",
    dest="storage_path",
    default="/srv/homeassistant/config/.storage",
    type=str,
    help="Path to storage",
)
parser.add_argument(
    "-y",
    "--skip_confirm",
    default=False,
    dest="skip_confirm",
    action="store_true",
    help="Skip confirm",
)
try:
    args = parser.parse_args()
except:
    exit(1)


def collect_devices() -> dict[str, str]:
    devices: dict[str, str] = {}
    with open(f"{args.storage_path}/core.device_registry") as devices_file:
        raw_devices = json.load(devices_file)

        for d in filter(
            lambda x: x["name_by_user"] != None, raw_devices["data"]["devices"]
        ):
            devices[d["id"]] = d["name_by_user"]

    return devices


def collect_entities() -> list[dict]:
    with open(f"{args.storage_path}/core.entity_registry") as entities_file:
        entities = json.load(entities_file)["data"]["entities"]

    return entities


def valid_entity_id(inp: str) -> str:
    return (
        inp.lower()
        .replace(" ", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("\\", "_")
    )


def confirm(msg):
    yon = input(f"{msg} (y/n): ")
    return True if yon == "y" else False


def main():
    devices = collect_devices()
    entities = collect_entities()

    for e in entities:
        if e["device_id"] in devices:
            old = e["entity_id"]
            suffix = valid_entity_id(e["original_name"])
            prefix = e["entity_id"].split(".")[0]
            name = valid_entity_id(devices[e["device_id"]])
            new = f"{prefix}.{name}_{suffix}"

            if old != new:
                print(f"{old} -> {new}")
                if not args.skip_confirm:
                    if confirm("Do you want to replace the name for this entity?"):
                        e["entity_id"] = new
                else:
                    e["entity_id"] = new

    with open(f"{args.storage_path}/core.entity_registry") as entities_file, open(
        "debug.json", "w"
    ) as of:
        entities_object = json.load(entities_file)
        entities_object["data"]["entities"] = entities

        json.dump(entities_object, of)

    with open(f"{args.storage_path}/core.entity_registry", "w") as entities_file:
        if confirm("Check 'debug.json' if the output looks reasonable then confirm"):
            json.dump(entities_object, entities_file)
            print("Replaced 'core.entity_registry' with the contents of debug.json")
            print("PLEASE RESTART THE SERVER NOW")
        else:
            print(
                "Nothing changed. You could try to fix debug.json and replace the original file with it."
            )


main()