from textwrap import dedent

import black
import pytest

from bec_widgets.cli.generate_cli import ClientGenerator


# Mock classes to test the generator
class MockBECWaveform1D:
    USER_ACCESS = ["set_frequency", "set_amplitude"]

    def set_frequency(self, frequency: float) -> list:
        """Set the frequency of the waveform."""
        return [frequency]

    def set_amplitude(self, amplitude: float) -> tuple[float, float]:
        """Set the amplitude of the waveform."""
        return amplitude, amplitude


class MockBECFigure:
    USER_ACCESS = ["add_plot", "remove_plot"]

    def add_plot(self, plot_id: str):
        """Add a plot to the figure."""
        pass

    def remove_plot(self, plot_id: str):
        """Remove a plot from the figure."""
        pass


def test_client_generator_with_black_formatting():
    generator = ClientGenerator()
    generator.generate_client([MockBECWaveform1D, MockBECFigure])

    # Format the expected output with black to ensure it matches the generator output
    expected_output = dedent(
        '''\
        # This file was automatically generated by generate_cli.py

        from bec_widgets.cli.client_utils import rpc_call, RPCBase, BECFigureClientMixin
        from typing import Literal, Optional, overload

        class MockBECWaveform1D(RPCBase):
            @rpc_call
            def set_frequency(self, frequency: float) -> list:
                """
                Set the frequency of the waveform.
                """
            @rpc_call
            def set_amplitude(self, amplitude: float) -> tuple[float, float]:
                """
                Set the amplitude of the waveform.
                """ 
        
        class MockBECFigure(RPCBase):
            @rpc_call
            def add_plot(self, plot_id: str):
                """
                Add a plot to the figure.
                """

            @rpc_call
            def remove_plot(self, plot_id: str):
                """
                Remove a plot from the figure.
                """
    '''
    )

    expected_output_formatted = black.format_str(
        expected_output, mode=black.FileMode(line_length=100)
    ).lstrip()

    generated_output_formatted = black.format_str(
        generator.header + "\n" + generator.content, mode=black.FileMode(line_length=100)
    )

    assert expected_output_formatted == generated_output_formatted