# qsys/io/json_io.py
from __future__ import annotations
from typing import Any, Dict, Iterable, Optional
import json
from pathlib import Path

# Schema (lightweight) for the JSON format:
# {
#   "qsys_ir_version": 1,
#   "n_qubits": N,
#   "instructions": [
#     {"op": "RZ", "targets": [0], "params": {"theta": 1.23}},
#     {"op": "CX", "targets": [0,1], "params": null}
#   ],
#   "metadata": { ... }  # optional
# }


def _normalize_instructions(circuit_or_instrs: Any) -> Iterable[Dict]:
    """
    Similar to text_io._iter_instructions_from_circuit
    """
    if hasattr(circuit_or_instrs, "instructions"):
        instrs = circuit_or_instrs.instructions
    else:
        instrs = circuit_or_instrs

    for ins in instrs:
        if isinstance(ins, dict):
            op = ins.get("op")
            targets = list(ins.get("targets", []))
            params = ins.get("params", None)
        else:
            op = getattr(ins, "op", getattr(ins, "name", None))
            targets = list(getattr(ins, "targets", getattr(ins, "qubits", [])))
            params = getattr(ins, "params", None)
        yield {"op": str(op), "targets": [int(t) for t in targets], "params": params}


def export_json(
    circuit: Any,
    path: str | Path,
    version: int = 1,
    metadata: Optional[Dict] = None,
    n_qubits: Optional[int] = None,
) -> None:
    p = Path(path)
    if n_qubits is None:
        n_qubits = getattr(circuit, "n_qubits", None)
    if n_qubits is None:
        # determine
        targets = []
        for ins in _normalize_instructions(circuit):
            targets.extend(ins["targets"])
        n_qubits = (max(targets) + 1) if targets else 0

    payload = {
        "qsys_ir_version": version,
        "n_qubits": n_qubits,
        "instructions": list(_normalize_instructions(circuit)),
    }
    if metadata:
        payload["metadata"] = metadata

    p.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")


def _validate_json_payload(payload: Dict) -> None:
    if "qsys_ir_version" not in payload:
        raise ValueError("Missing qsys_ir_version")
    if "n_qubits" not in payload or not isinstance(payload["n_qubits"], int):
        raise ValueError("Missing/invalid n_qubits")
    if "instructions" not in payload or not isinstance(payload["instructions"], list):
        raise ValueError("Missing/invalid instructions")
    for i, ins in enumerate(payload["instructions"], start=1):
        if not isinstance(ins, dict) or "op" not in ins or "targets" not in ins:
            raise ValueError(f"Invalid instruction at index {i}: {ins!r}")


def import_json(path: str | Path) -> Dict[str, Any]:
    p = Path(path)
    payload = json.loads(p.read_text(encoding="utf-8"))
    _validate_json_payload(payload)
    return payload
