"""
QuantumCircuit class mimicking Qiskit's API, delegating to qsys.
"""
from typing import Union, List, Optional
from .register import QuantumRegister, ClassicalRegister

# Import the core composer from existing codebase
# We treat qsys as the internal engine
try:
    from qsys.circuit.quantum_circuit import QuantumCircuit as QSysCircuit
except ImportError as e:
    # Fallback/Mock for development if qsys is not in path
    print(f"DEBUG: Failed to import qsys: {e}")
    QSysCircuit = None

class QuantumCircuit:
    """
    QuantumCircuit adaptable to Zena SDK.
    """
    def __init__(self, *regs, name: Optional[str] = None):
        self.name = name
        self.qregs = []
        self.cregs = []
        self._qubits = []
        self._clbits = []
        
        # Parse arguments similar to Qiskit (int or Register)
        if len(regs) == 0:
            # Empty circuit
            pass
        elif all(isinstance(r, int) for r in regs):
            # Case: QuantumCircuit(n_qubits, n_clbits)
            n_qubits = regs[0]
            n_clbits = regs[1] if len(regs) > 1 else 0
            self.add_register(QuantumRegister(n_qubits, 'q'))
            if n_clbits > 0:
                self.add_register(ClassicalRegister(n_clbits, 'c'))
        else:
            # Case: QuantumCircuit(qreg, creg)
            for r in regs:
                if isinstance(r, (QuantumRegister, ClassicalRegister)):
                    self.add_register(r)
                else:
                    raise TypeError("Arguments must be integers or Registers")

        # Initialize the underlying qsys Core Circuit
        if QSysCircuit:
            self._core_circuit = QSysCircuit(len(self._qubits), len(self._clbits))
        else:
            self._core_circuit = None
            print("Warning: qsys.circuit not found. Core delegation disabled.")

    def add_register(self, register: Union[QuantumRegister, ClassicalRegister]):
        if isinstance(register, QuantumRegister):
            self.qregs.append(register)
            # Flatten qubits mapping (simplistic for now)
            start_index = len(self._qubits)
            self._qubits.extend([i for i in range(start_index, start_index + register.size)])
        elif isinstance(register, ClassicalRegister):
            self.cregs.append(register)
            start_index = len(self._clbits)
            self._clbits.extend([i for i in range(start_index, start_index + register.size)])
            
    def h(self, qubit):
        """Apply Hadamard gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            self._core_circuit.h(q_idx)
        return self

    def x(self, qubit):
        """Apply Pauli-X gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            self._core_circuit.x(q_idx)
        return self

    def y(self, qubit):
        """Apply Pauli-Y gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            # Note: qsys uses string name 'y' for Instruction, which execute.py handles
            from qsys.ir.types import Instruction
            self._core_circuit._instructions.append(Instruction(name="y", qubits=(q_idx,)))
        return self

    def z(self, qubit):
        """Apply Pauli-Z gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            from qsys.ir.types import Instruction
            self._core_circuit._instructions.append(Instruction(name="z", qubits=(q_idx,)))
        return self

    def s(self, qubit):
        """Apply S gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            from qsys.ir.types import Instruction
            self._core_circuit._instructions.append(Instruction(name="s", qubits=(q_idx,)))
        return self

    def sdg(self, qubit):
        """Apply S-dagger gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            from qsys.ir.types import Instruction
            self._core_circuit._instructions.append(Instruction(name="sdg", qubits=(q_idx,)))
        return self

    def t(self, qubit):
        """Apply T gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            from qsys.ir.types import Instruction
            self._core_circuit._instructions.append(Instruction(name="t", qubits=(q_idx,)))
        return self

    def tdg(self, qubit):
        """Apply T-dagger gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            from qsys.ir.types import Instruction
            self._core_circuit._instructions.append(Instruction(name="tdg", qubits=(q_idx,)))
        return self

    def sx(self, qubit):
        """Apply sqrt-X gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            self._core_circuit.sx(q_idx)
        return self

    def rx(self, theta, qubit):
        """Apply RX gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            from qsys.ir.types import Instruction
            self._core_circuit._instructions.append(
                Instruction(name="rx", qubits=(q_idx,), params=(float(theta),))
            )
        return self

    def ry(self, theta, qubit):
        """Apply RY gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            from qsys.ir.types import Instruction
            self._core_circuit._instructions.append(
                Instruction(name="ry", qubits=(q_idx,), params=(float(theta),))
            )
        return self

    def rz(self, theta, qubit):
        """Apply RZ gate."""
        q_idx = self._resolve_index(qubit)
        if self._core_circuit:
            self._core_circuit.rz(q_idx, theta)
        return self

    def cx(self, control_qubit, target_qubit):
        """Apply CNOT gate."""
        c_idx = self._resolve_index(control_qubit)
        t_idx = self._resolve_index(target_qubit)
        if self._core_circuit:
            self._core_circuit.cx(c_idx, t_idx)
        return self

    def swap(self, qubit1, qubit2):
        """Apply SWAP gate."""
        q1_idx = self._resolve_index(qubit1)
        q2_idx = self._resolve_index(qubit2)
        if self._core_circuit:
            from qsys.ir.types import Instruction
            self._core_circuit._instructions.append(Instruction(name="swap", qubits=(q1_idx, q2_idx)))
        return self

    def measure(self, qubits, clbits):
        """Measure quantum bits into classical bits."""
        if isinstance(qubits, int):
            qubits = [qubits]
        if isinstance(clbits, int):
            clbits = [clbits]

        if len(qubits) != len(clbits):
            raise ValueError("Number of qubits and classical bits must match.")

        for q, c in zip(qubits, clbits):
            q_idx = self._resolve_index(q)
            c_idx = self._resolve_cbit_index(c)
            if self._core_circuit:
                self._core_circuit.measure(q_idx, c_idx)
        return self

    def draw(self, output='text'):
        """Draw the circuit."""
        if self._core_circuit:
            return self._core_circuit.draw(mode=output)
        return "Core circuit not initialized."

    def _resolve_index(self, qubit) -> int:
        """
        Resolve user input (int or bit object) to flat integer index for qsys.
        For now, we assume user passes integers as indices.
        """
        # TODO: Support Qubit object (e.g. qreg[0])
        if isinstance(qubit, int):
            if qubit < 0 or qubit >= len(self._qubits):
                raise IndexError(f"Qubit index {qubit} out of range.")
            return qubit
        raise TypeError("Qubit must be an integer index for now.")

    def _resolve_cbit_index(self, cbit) -> int:
        if isinstance(cbit, int):
             if cbit < 0 or cbit >= len(self._clbits):
                raise IndexError(f"Clbit index {cbit} out of range.")
             return cbit
        raise TypeError("Clbit must be an integer index for now.")
