# qsys/runtime/execute.py
from __future__ import annotations
from typing import Any, Dict

from qsys.circuit.quantum_circuit import QuantumCircuit
from qsys.errors import ValidationError
try:
    from simulator_statevector import QuantumCircuit as SimulatorCircuit
except ImportError:
    # Fallback or error - for now we assume it's there as per request
    from simulator_statevector import StateVecPy as SimulatorCircuit # This wont work, but prevents import error if structure mismatch.
    # Actually, let's just import it. If it fails, the user needs to compile.
    pass
from simulator_statevector import QuantumCircuit as SimulatorCircuit

# --- helper: try to call the transpiler / pass manager if available --- #
from importlib import import_module


def _maybe_run_transpiler(qc: QuantumCircuit, backend: Any) -> QuantumCircuit:
    try:
        transpiler = import_module("qsys.transpiler")
    except Exception:
        return qc

    candidates = ("run_pass_manager", "pass_manager", "apply_passes", "transpile")
    for name in candidates:
        fn = getattr(transpiler, name, None)
        if not callable(fn):
            continue
        try:
            return fn(qc, backend=backend)
        except TypeError:
            pass
        try:
            return fn(qc, backend)
        except TypeError:
            pass
        try:
            return fn(qc)
        except Exception:
            pass
    return qc


def _validate_circuit_indices(circuit: QuantumCircuit) -> None:
    """
    Strict runtime validation: ensure all qubit indices are within [0, n_qubits-1]
    and classical indices are within [0, n_clbits-1].
    Raises ValidationError with helpful message on the first problem found.
    """
    n_q = circuit.n_qubits
    n_c = circuit.n_clbits
    for instr in circuit.instructions:
        # check qubit indices
        for q in instr.qubits or ():
            if not isinstance(q, int):
                raise ValidationError(f"qubit index must be int, got {type(q)}")
            if q < 0 or q >= n_q:
                raise ValidationError(f"invalid qubit index {q}")
        # check classical bits if present (some instr may have clbits)
        for c in getattr(instr, "clbits", ()) or ():
            if not isinstance(c, int):
                raise ValidationError(f"classical bit index must be int, got {type(c)}")
            if c < 0 or c >= n_c:
                raise ValidationError(f"invalid classical bit index {c}")


def execute(
    circuit: QuantumCircuit,
    backend: Any = None,
    shots: int = 1024,
    seed: int | None = None,
    use_transpiler: bool = True,
) -> Dict[str, Any]:
    """
    Execute `circuit` on the local statevector simulator.

    - If use_transpiler is True, attempt to run the transpiler/pass manager.
    - Validate indices strictly before running (raises ValidationError).
    - If shots > 0, return sampled counts; if shots == 0 return the statevector.
    """
    if backend is not None:
        print("Warning: backend ignored – using simulator")

    # Optionally run the transpiler first (defensive)
    if use_transpiler:
        circuit = _maybe_run_transpiler(circuit, backend)

    # Runtime validation (strict)
    _validate_circuit_indices(circuit)

    # instantiate simulator circuit
    sim = SimulatorCircuit(circuit.n_qubits)

    # apply instructions
    for instr in circuit.instructions:
        op = instr.name
        qubits = instr.qubits
        params = instr.params or []

        if op == "rz":
            sim.rz(qubits[0], params[0])
        elif op == "sx":
            sim.sx(qubits[0])
        elif op == "x":
            sim.x(qubits[0])
        elif op == "y":
            sim.y(qubits[0])
        elif op == "z":
            sim.z(qubits[0])
        elif op == "h":
            sim.h(qubits[0])
        elif op == "s":
            sim.s(qubits[0])
        elif op == "sdg":
            sim.sdg(qubits[0])
        elif op == "t":
            sim.t(qubits[0])
        elif op == "tdg":
            sim.tdg(qubits[0])
        elif op == "rx":
            sim.rx(qubits[0], params[0])
        elif op == "ry":
            sim.ry(qubits[0], params[0])
        elif op == "cx":
            sim.cx(qubits[0], qubits[1])
        elif op == "swap":
            sim.swap(qubits[0], qubits[1])
        # measures are handled by the run() method implicitly or we ignore them here 
        # because the rust `run` method just measures all at the end?
        # WAIT. The Rust `run` method measures ALL qubits.
        # The Python `QuantumCircuit` has explicit `measure`.
        # The current `StateVecPy` implementation of `measure_counts` measures ALL qubits into a bitstring.
        # So it matches the implicit behavior if we ignore specific measure instructions 
        # (assuming we want full register measurement).
        # If the user wants partial measurement, the current Rust implementation doesn't support it well yet 
        # (it returns full n_qubits string).
        # So we just ignore 'measure' op here and rely on `run` doing full measurement.

    result: Dict[str, Any] = {"shots": shots, "seed": seed}
    if backend is not None:
        result["backend_id"] = getattr(backend, "id", "local_sim")

    if shots is not None and shots > 0:
        counts = sim.run(shots, seed)
        # ensure a plain dict is returned to tests
        result["counts"] = dict(counts)
    else:
        # return final statevector
        # sim is SimulatorCircuit (QuantumCircuit). We need to call statevector().
        sv = sim.statevector()
        result["statevector"] = sv

    return result
