"""
P3Q Classical Quantum Simulator Backend
Classical simulation of OpenQASM circuits for P4 interlock testing.

Authors: Ahmad Ali Parr, Jessica L. Williams (SNAPKITTYWEST)

NOT a quantum computer — pure classical simulation for verification.
"""

import os
import numpy as np
from dataclasses import dataclass
from enum import IntEnum
from typing import Optional


class QSimCommandType(IntEnum):
    KEYGEN_256      = 0x01
    NONCE_128       = 0x02
    GROVER_AES4     = 0x03
    AMP_EST_LEAKAGE = 0x04

class QSimStatus(IntEnum):
    SUCCESS     = 0x00
    TIMEOUT     = 0x01
    ERROR       = 0x02
    UNSUPPORTED = 0xFF


@dataclass
class QSimCommand:
    cmd_type:   QSimCommandType
    qubits:     int
    shots:      int
    params:     bytes
    event_id:   int


@dataclass
class QSimResponse:
    event_id:        int
    status:          QSimStatus
    data:            bytes
    shots_completed: int


class P3QSimulator:
    """Classical reference simulator for P3Q quantum circuits."""

    MAX_QUBITS = 32  # Classical simulation limit

    def __init__(self, seed: Optional[int] = None):
        self.rng = np.random.default_rng(seed)

    def execute(self, cmd: QSimCommand) -> QSimResponse:
        # Classical simulation — no qubit limit enforced (all commands simulated)

        dispatch = {
            QSimCommandType.KEYGEN_256:      self._sim_keygen,
            QSimCommandType.NONCE_128:       self._sim_nonce,
            QSimCommandType.GROVER_AES4:     self._sim_grover_aes4,
            QSimCommandType.AMP_EST_LEAKAGE: self._sim_amp_est,
        }
        fn = dispatch.get(cmd.cmd_type)
        if fn is None:
            return QSimResponse(cmd.event_id, QSimStatus.UNSUPPORTED, b'', 0)
        return fn(cmd)

    def _sim_keygen(self, cmd: QSimCommand) -> QSimResponse:
        # Hadamard ⊗ 256 + measure = uniform random bits
        return QSimResponse(cmd.event_id, QSimStatus.SUCCESS, self.rng.bytes(32), 1)

    def _sim_nonce(self, cmd: QSimCommand) -> QSimResponse:
        return QSimResponse(cmd.event_id, QSimStatus.SUCCESS, self.rng.bytes(16), 1)

    def _sim_grover_aes4(self, cmd: QSimCommand) -> QSimResponse:
        """
        Classical simulation of 4-round AES Grover.
        Returns random key candidates (uniform measurement distribution).
        Note: real Grover would need ~2^64 iterations for 128-bit key.
        """
        if len(cmd.params) < 32:
            return QSimResponse(cmd.event_id, QSimStatus.ERROR, b'', 0)
        n_results = min(cmd.shots, 100)
        results = self.rng.bytes(16 * n_results)
        return QSimResponse(cmd.event_id, QSimStatus.SUCCESS, results, n_results)

    def _sim_amp_est(self, cmd: QSimCommand) -> QSimResponse:
        # Simulate amplitude estimation: random 32-bit phase estimate
        phase = int(self.rng.integers(0, 2**32))
        return QSimResponse(
            cmd.event_id, QSimStatus.SUCCESS,
            phase.to_bytes(4, 'big'), cmd.shots
        )


# ── Self-test ──────────────────────────────────────────────────────────────

def test_p4_interlock():
    sim = P3QSimulator(seed=42)

    # KeyGen
    rsp = sim.execute(QSimCommand(QSimCommandType.KEYGEN_256, 256, 1, b'', 1))
    assert rsp.status == QSimStatus.SUCCESS and len(rsp.data) == 32
    print(f"KeyGen256: {rsp.data.hex()}")

    # Nonce
    rsp = sim.execute(QSimCommand(QSimCommandType.NONCE_128, 128, 1, b'', 2))
    assert rsp.status == QSimStatus.SUCCESS and len(rsp.data) == 16
    print(f"Nonce128:  {rsp.data.hex()}")

    # Grover AES-4 (simulated — uniform distribution)
    pt = bytes.fromhex("00112233445566778899aabbccddeeff")
    ct = bytes.fromhex("69c4e0d86a7b0430d8cdb78070b4c55a")
    rsp = sim.execute(QSimCommand(QSimCommandType.GROVER_AES4, 512, 10, pt+ct, 3))
    assert rsp.status == QSimStatus.SUCCESS
    print(f"Grover shots completed: {rsp.shots_completed}")

    # Amplitude estimation
    rsp = sim.execute(QSimCommand(QSimCommandType.AMP_EST_LEAKAGE, 257, 100, b'', 4))
    assert rsp.status == QSimStatus.SUCCESS
    print(f"AmpEst phase: {int.from_bytes(rsp.data,'big'):#010x}")

    print("All P4 interlock tests passed.")


if __name__ == "__main__":
    test_p4_interlock()