# ──────────────────────────────────────────────────────────────────────
# SCPN Control — Neuro-Symbolic Logic Compiler
# © 1998–2026 Miroslav Šotek. All rights reserved.
# Contact: www.anulum.li | protoscience@anulum.li
# ORCID: https://orcid.org/0009-0009-3560-0851
# License: MIT OR Apache-2.0
# ──────────────────────────────────────────────────────────────────────
"""
SCPN Controller Artifact (``.scpnctl.json``) loader / saver.

Defines the ``Artifact`` dataclass that mirrors the JSON schema sections
(meta, topology, weights, readout, initial_state) and provides lightweight
validation on load.
"""

from __future__ import annotations

import base64
import binascii
import json
import math
import zlib
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Union

ARTIFACT_SCHEMA_VERSION = "1.0.0"
MAX_PACKED_WORDS = 10_000_000
MAX_DECOMPRESSED_BYTES = MAX_PACKED_WORDS * 8
MAX_COMPRESSED_BYTES = 50_000_000


# ── Sub-structures ──────────────────────────────────────────────────────────


@dataclass
class FixedPoint:
    data_width: int
    fraction_bits: int
    signed: bool


@dataclass
class SeedPolicy:
    id: str
    hash_fn: str
    rng_family: str


@dataclass
class CompilerInfo:
    name: str
    version: str
    git_sha: str


@dataclass
class ArtifactMeta:
    artifact_version: str
    name: str
    dt_control_s: float
    stream_length: int
    fixed_point: FixedPoint
    firing_mode: str
    seed_policy: SeedPolicy
    created_utc: str
    compiler: CompilerInfo
    notes: str | None = None


@dataclass
class PlaceSpec:
    id: int
    name: str


@dataclass
class TransitionSpec:
    id: int
    name: str
    threshold: float
    margin: float | None = None
    delay_ticks: int = 0


@dataclass
class Topology:
    places: List[PlaceSpec]
    transitions: List[TransitionSpec]


@dataclass
class WeightMatrix:
    shape: List[int]  # [rows, cols]
    data: List[float]  # row-major


@dataclass
class PackedWeights:
    shape: List[int]  # [rows, cols, words]
    data_u64: List[int]


@dataclass
class PackedWeightsGroup:
    words_per_stream: int
    w_in_packed: PackedWeights
    w_out_packed: PackedWeights | None = None


@dataclass
class Weights:
    w_in: WeightMatrix
    w_out: WeightMatrix
    packed: PackedWeightsGroup | None = None


@dataclass
class ActionReadout:
    id: int
    name: str
    pos_place: int
    neg_place: int


@dataclass
class Readout:
    actions: List[ActionReadout]
    gains: List[float]
    abs_max: List[float]
    slew_per_s: List[float]


@dataclass
class PlaceInjection:
    place_id: int
    source: str
    scale: float
    offset: float
    clamp_0_1: bool


@dataclass
class InitialState:
    marking: List[float]
    place_injections: List[PlaceInjection]


# ── Artifact ────────────────────────────────────────────────────────────────


@dataclass
class Artifact:
    """Full SCPN controller artifact (``.scpnctl.json``)."""

    meta: ArtifactMeta
    topology: Topology
    weights: Weights
    readout: Readout
    initial_state: InitialState

    @property
    def nP(self) -> int:
        return len(self.topology.places)

    @property
    def nT(self) -> int:
        return len(self.topology.transitions)


# ── Validation ──────────────────────────────────────────────────────────────


class ArtifactValidationError(ValueError):
    """Raised when an artifact fails lightweight validation."""


def _encode_u64_compact(data_u64: List[int]) -> Dict[str, Any]:
    """Encode uint64 list as zlib-compressed base64 little-endian payload."""
    raw = bytearray()
    for value in data_u64:
        raw.extend(int(value & 0xFFFFFFFFFFFFFFFF).to_bytes(8, "little", signed=False))
    compressed = zlib.compress(bytes(raw), level=9)
    payload = base64.b64encode(compressed).decode("ascii")
    return {
        "encoding": "u64-le-zlib-base64",
        "count": len(data_u64),
        "data_u64_b64_zlib": payload,
    }


def _decode_u64_compact(encoded: Dict[str, Any]) -> List[int]:
    """Decode compact uint64 payload generated by ``_encode_u64_compact``."""
    if encoded.get("encoding") != "u64-le-zlib-base64":
        raise ArtifactValidationError(f"Unsupported packed encoding: {encoded.get('encoding')}")

    payload = encoded.get("data_u64_b64_zlib")
    if not isinstance(payload, str):
        raise ArtifactValidationError("Missing compact packed payload string.")

    count_val = encoded.get("count")
    if isinstance(count_val, int) and (count_val < 0 or count_val > MAX_PACKED_WORDS):
        raise ArtifactValidationError(f"Packed count {count_val} exceeds limit {MAX_PACKED_WORDS}.")

    try:
        comp = base64.b64decode(payload.encode("ascii"), validate=True)
    except (ValueError, binascii.Error) as exc:
        raise ArtifactValidationError(f"Invalid base64 payload: {exc}") from exc

    if len(comp) > MAX_COMPRESSED_BYTES:
        raise ArtifactValidationError(f"Compressed payload too large: {len(comp)} bytes.")

    try:
        decomp = zlib.decompressobj()
        raw = decomp.decompress(comp, MAX_DECOMPRESSED_BYTES + 1)
        if decomp.unconsumed_tail:
            raise ArtifactValidationError("Decompressed packed payload exceeds configured limit.")
        raw += decomp.flush()
    except ArtifactValidationError:
        raise
    except (zlib.error, ValueError, OSError) as exc:
        raise ArtifactValidationError(f"Invalid compact packed payload: {exc}") from exc

    if len(raw) > MAX_DECOMPRESSED_BYTES:
        raise ArtifactValidationError("Decompressed packed payload exceeds configured limit.")

    if len(raw) % 8 != 0:
        raise ArtifactValidationError(f"Compact packed payload byte-length {len(raw)} is not divisible by 8.")

    available = len(raw) // 8
    if isinstance(count_val, int):
        count = count_val
    elif count_val is None:
        count = available
    else:
        raise ArtifactValidationError(f"Invalid compact packed count type: {type(count_val).__name__}.")
    if count < 0 or count > available:
        raise ArtifactValidationError(f"Invalid compact packed count {count}; available words={available}.")

    return [int.from_bytes(raw[i * 8 : (i + 1) * 8], "little", signed=False) for i in range(count)]


def encode_u64_compact(data_u64: List[int]) -> Dict[str, Any]:
    """Public compact codec helper for deterministic uint64 payload encoding."""
    return _encode_u64_compact(list(map(int, data_u64)))


def decode_u64_compact(encoded: Dict[str, Any]) -> List[int]:
    """Public compact codec helper for deterministic uint64 payload decoding."""
    return _decode_u64_compact(encoded)


def _validate(artifact: Artifact) -> None:
    """Lightweight checks: required fields, ranges, shape consistency."""
    meta = artifact.meta

    if meta.firing_mode not in ("binary", "fractional"):
        raise ArtifactValidationError(f"firing_mode must be 'binary' or 'fractional', got '{meta.firing_mode}'")

    if isinstance(meta.fixed_point.data_width, bool) or not isinstance(meta.fixed_point.data_width, int):
        raise ArtifactValidationError("fixed_point.data_width must be an integer >= 1")
    if meta.fixed_point.data_width < 1:
        raise ArtifactValidationError("fixed_point.data_width must be >= 1")
    if isinstance(meta.fixed_point.fraction_bits, bool) or not isinstance(meta.fixed_point.fraction_bits, int):
        raise ArtifactValidationError("fixed_point.fraction_bits must be an integer >= 0")
    if meta.fixed_point.fraction_bits < 0:
        raise ArtifactValidationError("fixed_point.fraction_bits must be >= 0")
    if meta.fixed_point.fraction_bits >= meta.fixed_point.data_width:
        raise ArtifactValidationError("fixed_point.fraction_bits must be < fixed_point.data_width")
    if not isinstance(meta.fixed_point.signed, bool):
        raise ArtifactValidationError("fixed_point.signed must be a boolean")

    if isinstance(meta.stream_length, bool) or not isinstance(meta.stream_length, int):
        raise ArtifactValidationError("stream_length must be an integer >= 1")
    if meta.stream_length < 1:
        raise ArtifactValidationError("stream_length must be >= 1")

    if isinstance(meta.dt_control_s, bool) or not isinstance(meta.dt_control_s, (int, float)):
        raise ArtifactValidationError("dt_control_s must be finite and > 0")
    if not math.isfinite(meta.dt_control_s):
        raise ArtifactValidationError("dt_control_s must be finite and > 0")
    if meta.dt_control_s <= 0:
        raise ArtifactValidationError("dt_control_s must be > 0")

    # Weight ranges
    for val in artifact.weights.w_in.data:
        if not (-1.0 <= val <= 1.0):
            raise ArtifactValidationError(f"w_in weight {val} outside [-1, 1]")
    for val in artifact.weights.w_out.data:
        if not (0.0 <= val <= 1.0):
            raise ArtifactValidationError(f"w_out weight {val} outside [0, 1]")

    # Threshold ranges
    for t in artifact.topology.transitions:
        if isinstance(t.threshold, bool) or not isinstance(t.threshold, (int, float)):
            raise ArtifactValidationError(f"threshold {t.threshold} for '{t.name}' must be finite and in [0, 1]")
        if not math.isfinite(t.threshold):
            raise ArtifactValidationError(f"threshold {t.threshold} for '{t.name}' must be finite and in [0, 1]")
        if not (0.0 <= t.threshold <= 1.0):
            raise ArtifactValidationError(f"threshold {t.threshold} for '{t.name}' outside [0, 1]")
        if t.margin is not None:
            if isinstance(t.margin, bool) or not isinstance(t.margin, (int, float)):
                raise ArtifactValidationError(f"margin {t.margin} for '{t.name}' must be finite and >= 0")
            if not math.isfinite(t.margin) or t.margin < 0.0:
                raise ArtifactValidationError(f"margin {t.margin} for '{t.name}' must be finite and >= 0")
        if isinstance(t.delay_ticks, bool) or not isinstance(t.delay_ticks, int):
            raise ArtifactValidationError(f"delay_ticks {t.delay_ticks} for '{t.name}' must be an integer >= 0")
        if t.delay_ticks < 0:
            raise ArtifactValidationError(f"delay_ticks {t.delay_ticks} for '{t.name}' must be >= 0")

    # Shape consistency
    nP = artifact.nP
    nT = artifact.nT
    expected_w_in = nT * nP
    expected_w_out = nP * nT

    if len(artifact.weights.w_in.data) != expected_w_in:
        raise ArtifactValidationError(f"w_in data length {len(artifact.weights.w_in.data)} != nT*nP={expected_w_in}")
    if len(artifact.weights.w_out.data) != expected_w_out:
        raise ArtifactValidationError(f"w_out data length {len(artifact.weights.w_out.data)} != nP*nT={expected_w_out}")

    # Marking length
    if len(artifact.initial_state.marking) != nP:
        raise ArtifactValidationError(f"marking length {len(artifact.initial_state.marking)} != nP={nP}")
    for val in artifact.initial_state.marking:
        if not (0.0 <= val <= 1.0):
            raise ArtifactValidationError(f"initial marking {val} outside [0, 1]")
    for inj in artifact.initial_state.place_injections:
        if isinstance(inj.place_id, bool) or not isinstance(inj.place_id, int):
            raise ArtifactValidationError("place_injections.place_id must be an integer")
        if inj.place_id < 0 or inj.place_id >= nP:
            raise ArtifactValidationError(f"place_injections.place_id {inj.place_id} out of bounds for nP={nP}")
        if not isinstance(inj.source, str) or not inj.source:
            raise ArtifactValidationError("place_injections.source must be a non-empty string")
        if isinstance(inj.scale, bool) or not isinstance(inj.scale, (int, float)) or not math.isfinite(inj.scale):
            raise ArtifactValidationError("place_injections.scale must be finite numeric")
        if isinstance(inj.offset, bool) or not isinstance(inj.offset, (int, float)) or not math.isfinite(inj.offset):
            raise ArtifactValidationError("place_injections.offset must be finite numeric")
        if not isinstance(inj.clamp_0_1, bool):
            raise ArtifactValidationError("place_injections.clamp_0_1 must be a boolean")

    # Readout consistency
    n_actions = len(artifact.readout.actions)
    for action in artifact.readout.actions:
        if isinstance(action.id, bool) or not isinstance(action.id, int) or action.id < 0:
            raise ArtifactValidationError("readout.actions.id must be an integer >= 0")
        if not isinstance(action.name, str) or not action.name:
            raise ArtifactValidationError("readout.actions.name must be a non-empty string")
        if isinstance(action.pos_place, bool) or not isinstance(action.pos_place, int):
            raise ArtifactValidationError("readout.actions.pos_place must be an integer")
        if isinstance(action.neg_place, bool) or not isinstance(action.neg_place, int):
            raise ArtifactValidationError("readout.actions.neg_place must be an integer")
        if action.pos_place < 0 or action.pos_place >= nP:
            raise ArtifactValidationError(f"readout.actions.pos_place {action.pos_place} out of bounds for nP={nP}")
        if action.neg_place < 0 or action.neg_place >= nP:
            raise ArtifactValidationError(f"readout.actions.neg_place {action.neg_place} out of bounds for nP={nP}")
    if len(artifact.readout.gains) != n_actions:
        raise ArtifactValidationError("readout.gains length must equal number of actions")
    if len(artifact.readout.abs_max) != n_actions:
        raise ArtifactValidationError("readout.abs_max length must equal number of actions")
    if len(artifact.readout.slew_per_s) != n_actions:
        raise ArtifactValidationError("readout.slew_per_s length must equal number of actions")
    for val in artifact.readout.gains:
        if isinstance(val, bool) or not isinstance(val, (int, float)) or not math.isfinite(val):
            raise ArtifactValidationError("readout.gains must contain finite numeric values")
    for val in artifact.readout.abs_max:
        if isinstance(val, bool) or not isinstance(val, (int, float)) or not math.isfinite(val) or val < 0.0:
            raise ArtifactValidationError("readout.abs_max must contain finite numeric values >= 0")
    for val in artifact.readout.slew_per_s:
        if isinstance(val, bool) or not isinstance(val, (int, float)) or not math.isfinite(val) or val < 0.0:
            raise ArtifactValidationError("readout.slew_per_s must contain finite numeric values >= 0")


# ── Load / Save ─────────────────────────────────────────────────────────────


def load_artifact(path: Union[str, Path]) -> Artifact:
    """Parse a ``.scpnctl.json`` file into an ``Artifact`` dataclass."""
    with open(path, "r", encoding="utf-8") as f:
        obj = json.load(f)

    # Meta
    m = obj["meta"]
    meta = ArtifactMeta(
        artifact_version=m["artifact_version"],
        name=m["name"],
        dt_control_s=m["dt_control_s"],
        stream_length=m["stream_length"],
        fixed_point=FixedPoint(
            data_width=m["fixed_point"]["data_width"],
            fraction_bits=m["fixed_point"]["fraction_bits"],
            signed=m["fixed_point"]["signed"],
        ),
        firing_mode=m["firing_mode"],
        seed_policy=SeedPolicy(
            id=m["seed_policy"]["id"],
            hash_fn=m["seed_policy"]["hash_fn"],
            rng_family=m["seed_policy"]["rng_family"],
        ),
        created_utc=m["created_utc"],
        compiler=CompilerInfo(
            name=m["compiler"]["name"],
            version=m["compiler"]["version"],
            git_sha=m["compiler"]["git_sha"],
        ),
        notes=m.get("notes"),
    )

    # Topology
    places = [PlaceSpec(id=p["id"], name=p["name"]) for p in obj["topology"]["places"]]
    transitions = [
        TransitionSpec(
            id=t["id"],
            name=t["name"],
            threshold=t["threshold"],
            margin=t.get("margin"),
            delay_ticks=t.get("delay_ticks", 0),
        )
        for t in obj["topology"]["transitions"]
    ]
    topology = Topology(places=places, transitions=transitions)

    # Weights
    w_in = WeightMatrix(
        shape=obj["weights"]["w_in"]["shape"],
        data=list(map(float, obj["weights"]["w_in"]["data"])),
    )
    w_out = WeightMatrix(
        shape=obj["weights"]["w_out"]["shape"],
        data=list(map(float, obj["weights"]["w_out"]["data"])),
    )
    packed = None
    if "packed" in obj["weights"]:
        pw = obj["weights"]["packed"]
        w_in_obj = pw["w_in_packed"]
        if "data_u64" in w_in_obj:
            w_in_data = list(map(int, w_in_obj["data_u64"]))
        else:
            w_in_data = _decode_u64_compact(w_in_obj)
        pw_in = PackedWeights(
            shape=w_in_obj["shape"],
            data_u64=w_in_data,
        )
        pw_out = None
        if "w_out_packed" in pw:
            w_out_obj = pw["w_out_packed"]
            if "data_u64" in w_out_obj:
                w_out_data = list(map(int, w_out_obj["data_u64"]))
            else:
                w_out_data = _decode_u64_compact(w_out_obj)
            pw_out = PackedWeights(
                shape=w_out_obj["shape"],
                data_u64=w_out_data,
            )
        packed = PackedWeightsGroup(
            words_per_stream=int(pw["words_per_stream"]),
            w_in_packed=pw_in,
            w_out_packed=pw_out,
        )
    weights = Weights(w_in=w_in, w_out=w_out, packed=packed)

    # Readout
    actions = [
        ActionReadout(
            id=a["id"],
            name=a["name"],
            pos_place=a["pos_place"],
            neg_place=a["neg_place"],
        )
        for a in obj["readout"]["actions"]
    ]
    readout = Readout(
        actions=actions,
        gains=obj["readout"]["gains"]["per_action"],
        abs_max=obj["readout"]["limits"]["per_action_abs_max"],
        slew_per_s=obj["readout"]["limits"]["slew_per_s"],
    )

    # Initial state
    injections = [
        PlaceInjection(
            place_id=inj["place_id"],
            source=inj["source"],
            scale=inj["scale"],
            offset=inj["offset"],
            clamp_0_1=inj["clamp_0_1"],
        )
        for inj in obj["initial_state"]["place_injections"]
    ]
    initial_state = InitialState(
        marking=list(map(float, obj["initial_state"]["marking"])),
        place_injections=injections,
    )

    artifact = Artifact(
        meta=meta,
        topology=topology,
        weights=weights,
        readout=readout,
        initial_state=initial_state,
    )
    _validate(artifact)
    return artifact


def save_artifact(
    artifact: Artifact,
    path: Union[str, Path],
    compact_packed: bool = False,
) -> None:
    """Serialize an ``Artifact`` to indented JSON."""

    def _weight_matrix_dict(wm: WeightMatrix) -> Dict[str, Any]:
        return {"shape": wm.shape, "data": wm.data}

    packed_dict: Dict[str, Any] | None = None
    if artifact.weights.packed is not None:
        pg = artifact.weights.packed
        if compact_packed:
            pw_in_d = {"shape": pg.w_in_packed.shape}
            pw_in_d.update(_encode_u64_compact(pg.w_in_packed.data_u64))
        else:
            pw_in_d = {"shape": pg.w_in_packed.shape, "data_u64": pg.w_in_packed.data_u64}
        pw_out_d = None
        if pg.w_out_packed is not None:
            if compact_packed:
                pw_out_d = {"shape": pg.w_out_packed.shape}
                pw_out_d.update(_encode_u64_compact(pg.w_out_packed.data_u64))
            else:
                pw_out_d = {
                    "shape": pg.w_out_packed.shape,
                    "data_u64": pg.w_out_packed.data_u64,
                }
        packed_dict = {
            "words_per_stream": pg.words_per_stream,
            "w_in_packed": pw_in_d,
        }
        if pw_out_d is not None:
            packed_dict["w_out_packed"] = pw_out_d

    obj: Dict[str, Any] = {
        "meta": {
            "artifact_version": artifact.meta.artifact_version,
            "name": artifact.meta.name,
            "dt_control_s": artifact.meta.dt_control_s,
            "stream_length": artifact.meta.stream_length,
            "fixed_point": {
                "data_width": artifact.meta.fixed_point.data_width,
                "fraction_bits": artifact.meta.fixed_point.fraction_bits,
                "signed": artifact.meta.fixed_point.signed,
            },
            "firing_mode": artifact.meta.firing_mode,
            "seed_policy": {
                "id": artifact.meta.seed_policy.id,
                "hash_fn": artifact.meta.seed_policy.hash_fn,
                "rng_family": artifact.meta.seed_policy.rng_family,
            },
            "created_utc": artifact.meta.created_utc,
            "compiler": {
                "name": artifact.meta.compiler.name,
                "version": artifact.meta.compiler.version,
                "git_sha": artifact.meta.compiler.git_sha,
            },
        },
        "topology": {
            "places": [{"id": p.id, "name": p.name} for p in artifact.topology.places],
            "transitions": [
                {
                    "id": t.id,
                    "name": t.name,
                    "threshold": t.threshold,
                    **({"margin": t.margin} if t.margin is not None else {}),
                    "delay_ticks": int(t.delay_ticks),
                }
                for t in artifact.topology.transitions
            ],
        },
        "weights": {
            "w_in": _weight_matrix_dict(artifact.weights.w_in),
            "w_out": _weight_matrix_dict(artifact.weights.w_out),
        },
        "readout": {
            "actions": [
                {
                    "id": a.id,
                    "name": a.name,
                    "pos_place": a.pos_place,
                    "neg_place": a.neg_place,
                }
                for a in artifact.readout.actions
            ],
            "gains": {"per_action": artifact.readout.gains},
            "limits": {
                "per_action_abs_max": artifact.readout.abs_max,
                "slew_per_s": artifact.readout.slew_per_s,
            },
        },
        "initial_state": {
            "marking": artifact.initial_state.marking,
            "place_injections": [
                {
                    "place_id": inj.place_id,
                    "source": inj.source,
                    "scale": inj.scale,
                    "offset": inj.offset,
                    "clamp_0_1": inj.clamp_0_1,
                }
                for inj in artifact.initial_state.place_injections
            ],
        },
    }

    if artifact.meta.notes is not None:
        obj["meta"]["notes"] = artifact.meta.notes

    if packed_dict is not None:
        obj["weights"]["packed"] = packed_dict

    with open(path, "w", encoding="utf-8") as f:
        json.dump(obj, f, indent=2)
        f.write("\n")
