#!/usr/bin/env python3
"""
Live QPU Nudge - Single-bit measurement for decode influence
Uses IBM Torino via ~/.qiskit/qiskit-ibm.json

CDR Integration: Each QPU measurement is checkpointed for reversible
state tracking. This enables:
- Replay of quantum measurement sequences
- Rollback to previous QPU states
- Audit trail of all quantum operations
"""
import os
import hashlib
from typing import Tuple, Dict, Any, Optional

# CDR checkpointing (optional - graceful fallback if unavailable)
CDR_ENABLED = os.environ.get("QPU_CDR_CHECKPOINTS", "1") != "0"
_cdr_manager = None

def get_cdr_manager():
    """Get or create CDR state manager singleton for QPU tracking."""
    global _cdr_manager
    if _cdr_manager is None:
        try:
            from helix_cdc.cdr_state_manager import CDRStateManager
            _cdr_manager = CDRStateManager(tape_size=65536, max_checkpoints=1024)
            # Initialize CDR with deterministic seed for QPU tracking
            _cdr_manager.cdr_init(seed=b"qpu_measurement_tracker_seed_v1" + b"\x00" * 32)
        except ImportError:
            pass  # CDR not available, graceful fallback
    return _cdr_manager


def nudge() -> Tuple[int, Dict[str, Any]]:
    """
    Run a minimal quantum circuit (1 qubit Hadamard + measure).
    Returns: (measurement_bit, metadata)

    The bit (0 or 1) can influence decode (e.g., temperature adjustment).
    Fails noisy if QPU unavailable (so you notice).
    """
    try:
        from qiskit import QuantumCircuit, transpile
        from qiskit_ibm_runtime import QiskitRuntimeService, SamplerV2

        # Load service from ~/.qiskit/qiskit-ibm.json
        svc = QiskitRuntimeService()
        backend_name = "ibm_torino"
        backend = svc.backend(backend_name)

        # Minimal circuit: H gate + measure (creates 50/50 superposition)
        qc = QuantumCircuit(1, 1)
        qc.h(0)
        qc.measure(0, 0)

        # Transpile circuit for target backend (required for ISA compliance)
        qc_transpiled = transpile(qc, backend=backend, optimization_level=1)

        # Run on real hardware (SamplerV2 for free tier)
        sampler = SamplerV2(mode=backend)
        job = sampler.run([qc_transpiled], shots=256)
        res = job.result()

        # Extract measurement: majority vote over 256 shots
        # SamplerV2 returns BitArray, use get_counts()
        bit_array = res[0].data.c
        counts = bit_array.get_counts()

        total_shots = sum(counts.values())
        p1 = counts.get('1', 0) / total_shots  # Probability of |1⟩
        m = int(p1 >= 0.5)  # 1 if >50% measured |1⟩, else 0

        # Capture backend metadata for reproducibility
        meta = {
            "backend": backend_name,
            "job_id": job.job_id(),
            "p1": p1,  # Probability of |1⟩
            "shots": 256,
            "hardware": True,
            "backend_version": getattr(backend, 'version', None),
            "basis_gates": list(backend.operation_names) if hasattr(backend, 'operation_names') else None,
            "num_qubits": backend.num_qubits if hasattr(backend, 'num_qubits') else None,
        }

        # Add coupling map (serialized)
        if hasattr(backend, 'coupling_map') and backend.coupling_map:
            try:
                meta["coupling_map"] = backend.coupling_map.get_edges() if hasattr(backend.coupling_map, 'get_edges') else str(backend.coupling_map)
            except:
                pass

        # CDR checkpoint for QPU measurement (enables reversible tracking)
        if CDR_ENABLED:
            mgr = get_cdr_manager()
            if mgr is not None:
                # Create circuit hash for reproducibility
                circuit_hash = hashlib.sha256(
                    f"H|measure|{backend_name}".encode()
                ).hexdigest()[:16]

                # CDR absorb: record measurement in sponge state
                state_msg = f"qpu|{m}|{p1:.4f}|{job.job_id()[:8]}".encode()[:64]
                mgr.cdr_absorb(state_msg, store_message=True)

                # Named checkpoint for random-access rollback
                mgr.checkpoint_qpu_measurement(
                    measurement=m,
                    p1=p1,
                    job_id=job.job_id(),
                    circuit_hash=circuit_hash,
                )

                # Add CDR state to metadata
                meta["cdr"] = {
                    "enabled": True,
                    "step_count": mgr.cdr_state.step_count,
                    "state_hash": mgr.cdr_state_hash()[:16],
                    "checkpoint_id": f"qpu_{job.job_id()[:8]}",
                }

        return m, meta

    except Exception as e:
        raise RuntimeError(f"QPU nudge failed: {e}")


def nudge_mock(seed: int = None) -> Tuple[int, Dict[str, Any]]:
    """
    Mock QPU nudge for testing without real quantum hardware.
    Uses deterministic PRNG seeded by input or time.

    Args:
        seed: Optional seed for reproducibility

    Returns: (measurement_bit, metadata) - same format as real nudge()
    """
    import time
    import random

    if seed is None:
        seed = int(time.time() * 1000) % (2**32)

    rng = random.Random(seed)
    p1 = rng.random()  # Simulated probability of |1⟩
    m = int(p1 >= 0.5)

    job_id = f"mock_{seed:08x}"

    meta = {
        "backend": "mock_simulator",
        "job_id": job_id,
        "p1": p1,
        "shots": 256,
        "hardware": False,
        "seed": seed,
    }

    # CDR checkpoint for mock QPU measurement
    if CDR_ENABLED:
        mgr = get_cdr_manager()
        if mgr is not None:
            circuit_hash = hashlib.sha256(b"H|measure|mock").hexdigest()[:16]

            # CDR absorb
            state_msg = f"qpu_mock|{m}|{p1:.4f}|{job_id[:8]}".encode()[:64]
            mgr.cdr_absorb(state_msg, store_message=True)

            # Named checkpoint
            mgr.checkpoint_qpu_measurement(
                measurement=m,
                p1=p1,
                job_id=job_id,
                circuit_hash=circuit_hash,
            )

            meta["cdr"] = {
                "enabled": True,
                "step_count": mgr.cdr_state.step_count,
                "state_hash": mgr.cdr_state_hash()[:16],
                "checkpoint_id": f"qpu_{job_id[:8]}",
            }

    return m, meta


def get_qpu_cdr_status() -> Dict[str, Any]:
    """Get CDR status for QPU measurement tracking."""
    if not CDR_ENABLED:
        return {"enabled": False, "message": "CDR disabled (QPU_CDR_CHECKPOINTS=0)"}

    mgr = get_cdr_manager()
    if mgr is None:
        return {"enabled": False, "message": "CDR manager unavailable"}

    return {
        "enabled": True,
        "cdr_state": {
            "initialized": mgr.cdr_state.initialized,
            "step_count": mgr.cdr_state.step_count,
            "state_hash": mgr.cdr_state_hash()[:16],
        },
        "checkpoints": {
            "count": len(mgr.checkpoints),
            "ids": [k for k in mgr.checkpoints.keys() if k.startswith("qpu_")][-10:],
        },
        "stats": {
            "absorbs": mgr.stats.get("cdr_absorbs", 0),
            "rollbacks": mgr.stats.get("rollbacks_performed", 0),
        },
    }
