# Copyright 2024-2025 XMOS LIMITED.
# This Software is subject to the terms of the XMOS Public Licence: Version 1.

from functools import partial
import os
import queue
import threading
import sys
from PySide6.QtWidgets import (
    QApplication,
    QTabWidget,
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QSlider,
    QLabel,
    QGridLayout,
    QScrollArea,
)
from PySide6.QtCore import Qt, QTranslator, QLocale
import numpy as np
from PySide6.QtGui import QIcon
from pathlib import Path
from argparse import ArgumentParser
from tuning_utility import device
from audio_dsp.design.parse_json import make_pipeline, DspJson

from audio_dsp.design.pipeline import Pipeline, generate_dsp_main

from tuning_utility.widgets import XSlider, XDial, XButton, StageParameterGroupbox, PeqTab, BiquadWidget, Geq10bTab, tr_str, XMOSLogo
from tuning_utility.core import DspWindow
from tuning_utility.translations import register_translations




class An02031Window(DspWindow):
    """Main window for the AN02031 DSP controller application.
    This window contains tabs for controlling the DSP pipeline,
    displaying parameters, and showing the DSP graph as an image.

    The DSP control is inherited from the :class:`DspWindow` class.

    Parameters
    ----------
    params : DspJson
        The DSP parameters loaded from a JSON file.
    code_gen_dir : Path
        The directory where the generated DSP code will be saved.
    """

    def __init__(self, params: DspJson, code_gen_dir: Path):
        super().__init__(params, code_gen_dir, title="DSP Controller")

        # Main widget
        self.main_widget = QWidget()
        self.main_layout = QVBoxLayout()
        self.main_widget.setLayout(self.main_layout)
        self.setCentralWidget(self.main_widget)

        # tab widget
        self.tabs = QTabWidget()
        self.main_layout.addWidget(self.tabs)
        self.make_tabs()

        # add save and load buttons at the bottom
        self.main_layout.addLayout(self.make_save_load_gen_buttons())


    def make_tabs(self):
        """Create the tabs for the main window."""
        self.tabs.addTab(self.create_app_controls_tab(), "Control")
        self.tabs.addTab(self.create_all_parameters_tab(), "Configuration")

        # for all stages, if they are a Parametric EQ or Graphic EQ, add a tab
        for node in (self.state.state.graph.nodes):
            if node.op_type in ["ParametricEq8b", "ParametricEq16b","CascadedBiquads", "CascadedBiquads16b"]:
                self.tabs.addTab(PeqTab(self.state, node.placement.name), "Equaliser")
            elif node.op_type in ["GraphicEq10b"]:
                self.tabs.addTab(Geq10bTab(self.state, node.placement.name), node.placement.name)

        # add tab for the DSP graph image
        self.tabs.addTab(self.create_graph_image_tab(), "AN02031 Pipeline")

    def create_app_controls_tab(self):
        """
        Tab that contains the main controls for the application. This
        tab uses custom widgets to expose certain DSP controls to the
        end user.

        In AN02031, this tab contains volume sliders for the
        microphone, music, headphone, and output stages, as well as mute
        buttons and controls for the reverb, ducking, and monitoring
        switches.
        """
        tab = QWidget()
        layout = QHBoxLayout()

        # Volume sliders and mute buttons
        grid = QGridLayout()

        # We use an if so the GUI will work if the stage_name is not present in the pipeline
        # but this is not always required.
        if "mic_volume" in self.state.node_dict.keys():
            # add the microphone volume slider and mute button to the 1st column of the grid
            slider = XSlider(self.state, "mic_volume", "gain_db", label="Microphone\nVolume", widget_range=(-45, 0))
            grid.addWidget(slider, 0, 0, Qt.AlignHCenter)

            button = XButton(self.state, "mic_volume", "mute_state", label="Mute")
            grid.addWidget(button, 1, 0, Qt.AlignHCenter)

        if "music_volume" in self.state.node_dict.keys():
            # add the music volume slider and mute button to the 2nd column of the grid
            slider = XSlider(self.state, "music_volume", "gain_db", label="Music\nVolume", widget_range=(-45, 0))
            grid.addWidget(slider, 0, 1, Qt.AlignHCenter)

            button = XButton(self.state, "music_volume", "mute_state", label="Mute")
            grid.addWidget(button, 1, 1, Qt.AlignHCenter)

        if "headphone_volume" in self.state.node_dict.keys():
            # add the music volume slider and mute button to the 3rd column of the grid
            slider = XSlider(self.state, "headphone_volume", "gain_db", label="Headphone\nVolume", widget_range=(-45, 0))
            grid.addWidget(slider, 0, 2, Qt.AlignHCenter)

            button = XButton(self.state, "headphone_volume", "mute_state", label="Mute")
            grid.addWidget(button, 1, 2, Qt.AlignHCenter)

        if "output_volume" in self.state.node_dict.keys():
            # add the output volume slider and mute button to the 4th column of the grid
            slider = XSlider(self.state, "output_volume", "gain_db", label="Output\nVolume", widget_range=(-45, 0))
            grid.addWidget(slider, 0, 3, Qt.AlignHCenter)

            button = XButton(self.state, "output_volume", "mute_state", label="Mute")
            grid.addWidget(button, 1, 3, Qt.AlignHCenter)

        if "reverb" in self.state.node_dict.keys():
            # additional slider for reverb in the 5th column of the grid
            slider = XSlider(self.state, "reverb", "wet_dry_mix", label="Reverb\nWet/Dry")
            grid.addWidget(slider, 0, 4, Qt.AlignHCenter)

        if "reverb_switch" in self.state.node_dict.keys():
            # add the reverb switch button to the 5th column of the grid
            button = XButton(self.state, "reverb_switch", "position", label="Reverb")
            grid.addWidget(button, 1, 4, Qt.AlignHCenter)

        if "duck_switch" in self.state.node_dict.keys():
            # add the duck switch button to the 3rd row, 2nd column of the grid
            button = XButton(self.state, "duck_switch", "position", label="Ducking")
            grid.addWidget(button, 2, 1, Qt.AlignHCenter)

        if "monitor_switch" in self.state.node_dict.keys():
            # add the duck switch button to the 3rd row, 3rd column of the grid
            button = XButton(self.state, "monitor_switch", "position", label="Monitor")
            grid.addWidget(button, 2, 2, Qt.AlignHCenter)

        if "loopback_switch" in self.state.node_dict.keys():
            # add the duck switch button to the 3rs row, 4th column of the grid
            button = XButton(self.state, "loopback_switch", "position", label="Loopback")
            grid.addWidget(button, 2, 3, Qt.AlignHCenter)

        # add xmos logo
        logo = QLabel()
        icon = XMOSLogo()
        logo.setPixmap(icon.pixmap(68, 30*103/68))
        grid.addWidget(logo, 2, 4, Qt.AlignHCenter)

        # put in middle
        layout.addStretch(1)
        layout.addLayout(grid)
        layout.addStretch(1)
        tab.setLayout(layout)
        return tab

    def create_all_parameters_tab(self):
        """Tab that displays controls for all the available parameters.
        in the pipeline. It is intended for the DSP engineer to tune the
        product.

        For each stage, a groupbox is created containing the runtime
        controllable parameters. Note that Parametric EQs and Graphic
        EQs are handled separately.
        """
        tab = QWidget()

        widget = QScrollArea()
        widget.setWidget(tab)
        widget.setWidgetResizable(True)

        layout = QVBoxLayout(tab)

        # for each stage in the graph, create a StageParameterGroupbox
        for node in self.state.state.graph.nodes:
            if (hasattr(node, "parameters")
            and len(type(node.parameters).model_fields) > 0
            and node.op_type not in ["ParametricEq8b", "ParametricEq16b", "GraphicEq10b",
            "CascadedBiquads", "CascadedBiquads16b"]):
                layout.addWidget(StageParameterGroupbox(self.state, node.placement.name))

        tab.setLayout(layout)
        return widget

    def create_graph_image_tab(self):
        """Tab that displays the DSP Graph as an image."""
        tab = QWidget()
        layout = QVBoxLayout()

        # Load and display DSP Graph image
        label = QLabel()
        pipeline = Path(__file__).parent / "pipeline.png"
        icon = QIcon(str(pipeline))
        label.setPixmap(
            icon.pixmap(600, 400)
        )
        layout.addWidget(label, alignment=Qt.AlignCenter)

        tab.setLayout(layout)
        return tab


def parse_args(tuning_path):
    parser = ArgumentParser()
    parser.add_argument(
        "--code-gen",
        "-c",
        type=Path,
        nargs='?',
        const=tuning_path,
        help=("If set, generate the pipeline code from the given tuning path and exit. "
              "Optionally, the path to a specific .json file can be supplied. "
              f"By default uses the file {tuning_path.name}"),
    )
    parser.add_argument(
        "--draw",
        "-d",
        type=Path,
        help="if set, draw the pipeline and exit",
    )
    return parser.parse_args()


if __name__ == "__main__":

    this_dir = Path(__file__).parent
    tuning_path = this_dir / "live_streaming_sound_card.json"

    if not tuning_path.exists():
        json_obj = None
    else:
        json_obj = DspJson.model_validate_json(tuning_path.read_text())

    code_gen_dir = this_dir / "app_an02031/src/generated_dsp"


    args = parse_args(tuning_path)

    if args.code_gen:
        p = make_pipeline(json_obj)
        generate_dsp_main(p, out_dir=code_gen_dir)
        exit(0)

    if args.draw:
        p = make_pipeline(json_obj)
        p.draw(args.draw)
        exit(0)

    # this is the main call to the GUI application
    app = QApplication(sys.argv)

    register_translations(app)

    window = An02031Window(json_obj, code_gen_dir)
    window.show()
    sys.exit(app.exec())

