#!/usr/bin/env python3
"""Align BoundaryFlag with sbc/ef/mag/scd/pbc index blocks in OpticsFEM JSON files."""
from __future__ import annotations

import json
import sys
from pathlib import Path

PMC, PEC, SBC, ELE, PBC, MAG, SCD = 0, 1, 2, 3, 4, 8, 9

ROOT = Path(__file__).resolve().parents[2]

OPTICS_JSON_GLOBS = [
    ROOT / "3D opticsfem-master" / "*.json",
    ROOT / "3D opticsfem-master" / "scat3D" / "Release" / "*.json",
    ROOT / "3D opticsfem-master" / "scat3D_SBC" / "Release" / "*.json",
    ROOT / "三维matlab代码" / "matlab 3D一阶散射问题" / "*.json",
    ROOT / "三维matlab代码" / "matlab 3D一阶基+散射边界条件" / "*.json",
    ROOT / "三维matlab代码" / "matlab 3D一阶基 + bele" / "*.json",
    ROOT / "三维matlab代码" / "matlab 3D一阶基+散射边界条件+单周期边界" / "PBCmesh.json",
]


def collect_pbc_domains(doc: dict) -> set[int]:
    domains: set[int] = set()
    for key in ("pbc", "pbc2"):
        if key not in doc:
            continue
        block = doc[key]
        for field in ("srcIndex", "dstIndex"):
            if field in block:
                domains.update(int(x) for x in block[field])
    return domains


def build_boundary_flag(doc: dict) -> list[int] | None:
    if "NbrBoundary" not in doc or "BoundaryFlag" not in doc:
        return None
    n = int(doc["NbrBoundary"])
    flags = [PMC] * n

    for d in collect_pbc_domains(doc):
        if 1 <= d <= n:
            flags[d - 1] = PBC

    if "sbc" in doc:
        for d in doc["sbc"]["Index"]:
            d = int(d)
            if 1 <= d <= n:
                flags[d - 1] = SBC

    if "ef" in doc:
        for d in doc["ef"]["index"]:
            d = int(d)
            if 1 <= d <= n:
                flags[d - 1] = ELE

    if "mag" in doc:
        for d in doc["mag"]["index"]:
            d = int(d)
            if 1 <= d <= n:
                flags[d - 1] = MAG

    if "scd" in doc:
        for d in doc["scd"]["index"]:
            d = int(d)
            if 1 <= d <= n:
                flags[d - 1] = SCD

    # Legacy 14-face cube: domain 6 is material interface, not an outer SBC face.
    if n == 14 and "sbc" in doc:
        sbc_set = {int(x) for x in doc["sbc"]["Index"]}
        if 6 not in sbc_set and flags[5] == SBC:
            flags[5] = PMC

    return flags


def validate(doc: dict, path: Path) -> list[str]:
    issues: list[str] = []
    if "BoundaryFlag" not in doc:
        return issues
    flags = doc["BoundaryFlag"]
    n = len(flags)

    if "sbc" in doc:
        for d in doc["sbc"]["Index"]:
            d = int(d)
            if not (1 <= d <= n):
                issues.append(f"sbc.Index {d} out of range")
            elif flags[d - 1] != SBC:
                issues.append(f"sbc.Index {d} but BoundaryFlag[{d}]={flags[d - 1]} (expected {SBC})")

    if "ef" in doc:
        for d in doc["ef"]["index"]:
            d = int(d)
            if flags[d - 1] != ELE:
                issues.append(f"ef.index {d} but BoundaryFlag[{d}]={flags[d - 1]} (expected {ELE})")

    if "mag" in doc:
        for d in doc["mag"]["index"]:
            d = int(d)
            if flags[d - 1] != MAG:
                issues.append(f"mag.index {d} flag={flags[d - 1]} (expected {MAG})")

    for d in collect_pbc_domains(doc):
        if 1 <= d <= n and flags[d - 1] != PBC:
            issues.append(f"pbc domain {d} but BoundaryFlag[{d}]={flags[d - 1]} (expected {PBC})")

    return issues


def main() -> int:
    write = "--write" in sys.argv
    paths: list[Path] = []
    for pattern in OPTICS_JSON_GLOBS:
        if pattern.parent.exists():
            paths.extend(sorted(pattern.parent.glob(pattern.name)))

    changed = 0
    for path in paths:
        if path.name in ("CMakeSettings.json",):
            continue
        try:
            doc = json.loads(path.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        if "FemType" not in doc:
            continue

        new_flags = build_boundary_flag(doc)
        if new_flags is None:
            continue

        old_flags = doc.get("BoundaryFlag")
        issues_before = validate(doc, path) if old_flags else ["missing BoundaryFlag"]

        if old_flags == new_flags and not issues_before:
            print(f"OK  {path.relative_to(ROOT)}")
            continue

        print(f"FIX {path.relative_to(ROOT)}")
        if old_flags != new_flags:
            print(f"    old: {old_flags}")
            print(f"    new: {new_flags}")
        for msg in issues_before:
            print(f"    was: {msg}")

        if write:
            doc["BoundaryFlag"] = new_flags
            path.write_text(json.dumps(doc, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
            changed += 1

    if write:
        print(f"\nUpdated {changed} file(s).")
    else:
        print("\nDry run. Re-run with --write to apply.")
    return 0


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