# Copyright 2026 The PECOS Developers
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.

"""Tests for AST to Stim code generator."""

import pytest

stim = pytest.importorskip("stim")

from pecos.slr import Barrier, CReg, If, Main, QReg, Repeat  # noqa: E402
from pecos.slr.ast import slr_to_ast  # noqa: E402
from pecos.slr.ast.codegen import AstToStim, ast_to_stim, ast_to_stim_str  # noqa: E402
from pecos.slr.qeclib import qubit as qb  # noqa: E402


class TestAstToStimBasic:
    """Basic code generation tests."""

    def test_empty_program(self) -> None:
        """Empty program generates empty Stim circuit."""
        prog = Main()
        ast = slr_to_ast(prog)

        circuit = ast_to_stim(ast)

        assert isinstance(circuit, stim.Circuit)
        assert len(circuit) == 0

    def test_program_with_qreg(self) -> None:
        """Program with QReg generates non-empty circuit."""
        prog = Main(
            q := QReg("q", 2),
            qb.H(q[0]),
        )
        ast = slr_to_ast(prog)

        circuit = ast_to_stim(ast)

        assert isinstance(circuit, stim.Circuit)
        assert len(circuit) > 0

    def test_string_output(self) -> None:
        """String output function returns circuit as string."""
        prog = Main(
            q := QReg("q", 1),
            qb.H(q[0]),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert isinstance(code, str)
        assert "H" in code


class TestAstToStimGates:
    """Gate code generation tests."""

    def test_hadamard_gate(self) -> None:
        """Hadamard gate generates H instruction."""
        prog = Main(
            q := QReg("q", 1),
            qb.H(q[0]),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert "H 0" in code

    def test_pauli_gates(self) -> None:
        """Pauli gates generate X, Y, Z instructions."""
        prog = Main(
            q := QReg("q", 1),
            qb.X(q[0]),
            qb.Y(q[0]),
            qb.Z(q[0]),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert "X 0" in code
        assert "Y 0" in code
        assert "Z 0" in code

    def test_phase_gates(self) -> None:
        """Phase gates generate S and S_DAG instructions."""
        prog = Main(
            q := QReg("q", 1),
            qb.SZ(q[0]),
            qb.SZdg(q[0]),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert "S 0" in code
        assert "S_DAG 0" in code

    def test_t_gates(self) -> None:
        """T gate handles gracefully (may be unsupported in Stim)."""
        # Note: T gates are non-Clifford and Stim uses them for noise modeling
        # The Stim gate is called "T" not "T_DAG" for the adjoint
        prog = Main(
            q := QReg("q", 1),
            qb.T(q[0]),
        )
        ast = slr_to_ast(prog)

        # T gate may not be directly supported - check the generator handles it
        # If T isn't supported, it should skip or the test should be adjusted
        try:
            code = ast_to_stim_str(ast)
            # If T is supported, it should appear in output
            assert "T" in code or len(code) == 0  # May be skipped if unsupported
        except (IndexError, ValueError):
            # T gate not supported in Stim - this is expected
            pass

    def test_two_qubit_cx_gate(self) -> None:
        """CX gate generates CX instruction with correct qubits."""
        prog = Main(
            q := QReg("q", 2),
            qb.CX(q[0], q[1]),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert "CX 0 1" in code

    def test_two_qubit_cz_gate(self) -> None:
        """CZ gate generates CZ instruction with correct qubits."""
        prog = Main(
            q := QReg("q", 2),
            qb.CZ(q[0], q[1]),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert "CZ 0 1" in code

    def test_multiple_gates(self) -> None:
        """Multiple gates generate in sequence."""
        prog = Main(
            q := QReg("q", 2),
            qb.H(q[0]),
            qb.X(q[1]),
            qb.CZ(q[0], q[1]),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert "H 0" in code
        assert "X 1" in code
        assert "CZ 0 1" in code


class TestAstToStimPrepMeasure:
    """Prep and measure code generation tests."""

    def test_measurement(self) -> None:
        """Measurement generates M instruction."""
        prog = Main(
            q := QReg("q", 1),
            c := CReg("c", 1),
            qb.Measure(q[0]) > c[0],
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert "M 0" in code

    def test_multiple_measurements(self) -> None:
        """Multiple measurements generate M instructions for all qubits."""
        prog = Main(
            q := QReg("q", 2),
            c := CReg("c", 2),
            qb.Measure(q[0]) > c[0],
            qb.Measure(q[1]) > c[1],
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        # Stim may combine measurements: "M 0 1" or "M 0\nM 1"
        assert "M" in code
        # Check both qubits are measured (format may vary)
        assert "0" in code
        assert "1" in code

    def test_prep_reset(self) -> None:
        """Prep generates R (reset) instruction."""
        prog = Main(
            q := QReg("q", 1),
            qb.Prep(q[0]),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert "R 0" in code


class TestAstToStimControlFlow:
    """Control flow code generation tests."""

    def test_barrier_becomes_tick(self) -> None:
        """Barrier generates TICK instruction."""
        prog = Main(
            q := QReg("q", 2),
            qb.H(q[0]),
            Barrier(q),
            qb.CX(q[0], q[1]),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert "TICK" in code

    def test_repeat_uses_repeat_block(self) -> None:
        """Repeat generates REPEAT block."""
        prog = Main(
            q := QReg("q", 1),
            Repeat(cond=3).block(
                qb.H(q[0]),
            ),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        # Stim uses REPEAT blocks
        assert "REPEAT 3" in code
        assert "H 0" in code

    def test_if_statement_adds_tick(self) -> None:
        """If statement adds TICK marker (conditionals unsupported in Stim)."""
        prog = Main(
            q := QReg("q", 1),
            c := CReg("c", 1),
            If(c[0] == 1).Then(
                qb.H(q[0]),
            ),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        # Conditionals add TICK markers since Stim doesn't support conditionals
        assert "TICK" in code
        assert "H 0" in code


class TestAstToStimQEC:
    """QEC pattern code generation tests."""

    def test_syndrome_extraction(self) -> None:
        """Syndrome extraction generates correct qubit indices."""
        prog = Main(
            data := QReg("data", 2),
            ancilla := QReg("ancilla", 1),
            c := CReg("c", 1),
            qb.CX(data[0], ancilla[0]),
            qb.CX(data[1], ancilla[0]),
            qb.Measure(ancilla[0]) > c[0],
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        # data[0] -> qubit 0, data[1] -> qubit 1, ancilla[0] -> qubit 2
        # Stim may combine CX gates into one line: "CX 0 2 1 2"
        assert "CX" in code
        assert "0 2" in code or "0, 2" in code  # First CX pair
        assert "1 2" in code or "1, 2" in code  # Second CX pair
        assert "M 2" in code

    def test_repeated_syndrome_extraction(self) -> None:
        """Repeated syndrome extraction generates REPEAT block."""
        prog = Main(
            data := QReg("data", 2),
            ancilla := QReg("ancilla", 1),
            c := CReg("c", 1),
            Repeat(cond=3).block(
                qb.CX(data[0], ancilla[0]),
                qb.CX(data[1], ancilla[0]),
                qb.Measure(ancilla[0]) > c[0],
                qb.Prep(ancilla[0]),
            ),
        )
        ast = slr_to_ast(prog)

        code = ast_to_stim_str(ast)

        assert "REPEAT 3" in code


class TestAstToStimGenerator:
    """Tests for AstToStim generator class."""

    def test_generator_reusable(self) -> None:
        """Generator can be reused for multiple programs."""
        generator = AstToStim()

        prog1 = Main(
            q := QReg("q", 1),
            qb.H(q[0]),
        )

        prog2 = Main(
            r := QReg("r", 2),
            qb.X(r[0]),
        )

        ast1 = slr_to_ast(prog1)
        ast2 = slr_to_ast(prog2)

        circuit1 = generator.generate(ast1)
        circuit2 = generator.generate(ast2)

        code1 = str(circuit1)
        code2 = str(circuit2)

        assert "H 0" in code1
        assert "X 0" in code2

    def test_measurement_count_tracked(self) -> None:
        """Generator tracks measurement count."""
        prog = Main(
            q := QReg("q", 3),
            c := CReg("c", 3),
            qb.Measure(q[0]) > c[0],
            qb.Measure(q[1]) > c[1],
            qb.Measure(q[2]) > c[2],
        )
        ast = slr_to_ast(prog)

        generator = AstToStim()
        generator.generate(ast)

        assert generator.context.measurement_count == 3


class TestAstToStimFullPipeline:
    """End-to-end tests: SLR -> AST -> Stim."""

    def test_bell_state_circuit(self) -> None:
        """Bell state generates H and CX instructions."""
        prog = Main(
            q := QReg("q", 2),
            qb.H(q[0]),
            qb.CX(q[0], q[1]),
        )

        ast = slr_to_ast(prog)
        circuit = ast_to_stim(ast)

        # Verify circuit structure
        code = str(circuit)
        assert "H 0" in code
        assert "CX 0 1" in code

    def test_circuit_is_valid_stim(self) -> None:
        """Test that generated circuit can be used by Stim."""
        prog = Main(
            q := QReg("q", 2),
            c := CReg("c", 2),
            qb.H(q[0]),
            qb.CX(q[0], q[1]),
            qb.Measure(q[0]) > c[0],
            qb.Measure(q[1]) > c[1],
        )

        ast = slr_to_ast(prog)
        circuit = ast_to_stim(ast)

        # Verify Stim can sample from the circuit
        sampler = circuit.compile_sampler()
        samples = sampler.sample(shots=10)

        assert samples.shape == (10, 2)

    def test_ghz_state_circuit(self) -> None:
        """Test a GHZ state circuit."""
        prog = Main(
            q := QReg("q", 3),
            qb.H(q[0]),
            qb.CX(q[0], q[1]),
            qb.CX(q[1], q[2]),
        )

        ast = slr_to_ast(prog)
        code = ast_to_stim_str(ast)

        assert "H 0" in code
        # Stim may combine CX gates: "CX 0 1 1 2" or separate lines
        assert "CX" in code
        assert "0 1" in code or "0, 1" in code
        assert "1 2" in code or "1, 2" in code
