"""TUI for the distortion-tui package."""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, NamedTuple

import librosa
import matplotlib.pyplot as plt
import soundfile
from textual.app import App, ComposeResult
from textual.containers import Horizontal, Vertical
from textual.widgets import Button, Checkbox, Footer, Header, Input, Label, Log, Select

import distortion_tui

# Maps display label → (function name in package, output filename suffix)
_FUNCTIONS: list[tuple[str, str, str]] = [
    ("Isochronous", "isochronous", "_isochronous"),
    ("Infinite Peak Clipping", "infinite_peak_clipped", "_clipped"),
]
_SELECT_OPTIONS = [(label, key) for label, key, _ in _FUNCTIONS]
_SUFFIX_BY_KEY = {key: suffix for _, key, suffix in _FUNCTIONS}

if TYPE_CHECKING:
    import numpy as np


class PlotPayload(NamedTuple):
    """Container for waveform data used for plotting."""

    waveform: np.ndarray
    processed: np.ndarray
    label: str
    sr: float


class DistortionApp(App):
    """Textual TUI for audio distortion functions."""

    CSS = """
    /* UZH brand colors:
       Blue #0028A5 | Dark #001A6E | Light bg #F0F3FC | Text #1A1A1A | Muted #6B7280 */

    Screen {
        background: #F0F3FC;
        color: #1A1A1A;
        padding: 1 2;
    }

    Header {
        background: #0028A5;
        color: #FFFFFF;
        text-style: bold;
    }

    Footer {
        background: #001A6E;
        color: #FFFFFF;
    }

    #form {
        background: #FFFFFF;
        border: solid #0028A5;
        padding: 1 2;
        height: auto;
        margin-bottom: 1;
    }

    Label {
        color: #0028A5;
        text-style: bold;
        margin-top: 1;
    }

    Input {
        background: #F0F3FC;
        border: tall #0028A5;
        color: #1A1A1A;
        margin-bottom: 0;
    }

    Input:focus {
        border: tall #001A6E;
        background: #FFFFFF;
    }

    Select {
        background: #F0F3FC;
        border: tall #0028A5;
        color: #1A1A1A;
        margin-bottom: 0;
    }

    Select:focus {
        border: tall #001A6E;
    }

    SelectOverlay {
        background: #FFFFFF;
        border: tall #0028A5;
        color: #1A1A1A;
    }

    #plot-row {
        height: auto;
        margin-top: 1;
    }

    #plot_path {
        width: 1fr;
    }

    #save-plot {
        width: auto;
        margin-right: 1;
        color: #1A1A1A;
        background: #FFFFFF;
    }

    Checkbox:focus {
        border: tall #0028A5;
    }

    #buttons {
        height: auto;
        margin-top: 1;
        align: center middle;
    }

    Button {
        margin: 0 1;
        background: #0028A5;
        color: #FFFFFF;
        border: none;
        text-style: bold;
    }

    Button:hover {
        background: #001A6E;
    }

    Button:focus {
        background: #001A6E;
        border: none;
    }

    Button.-error {
        background: #6B0028;
    }

    Button.-error:hover {
        background: #4A001C;
    }

    #log {
        background: #FFFFFF;
        color: #1A1A1A;
        border: solid #0028A5;
        height: 1fr;
        margin-top: 1;
    }
    """

    BINDINGS = (("ctrl+q", "quit", "Quit"),)

    def compose(self) -> ComposeResult:  # noqa: PLR6301
        """
        Build the UI.

        Yields
        ------
        ComposeResult
            The UI elements to compose.

        """
        yield Header()
        with Vertical(id="form"):
            yield Label("Function")
            yield Select(_SELECT_OPTIONS, value="isochronous", id="function")
            yield Label("Input audio file")
            yield Input(placeholder="/path/to/audio.wav", id="input_path")
            yield Label("Output audio file")
            yield Input(placeholder="/path/to/output.wav", id="output_path")
            with Horizontal(id="plot-row"):
                yield Checkbox("Save waveform plot", id="save-plot")
                yield Input(
                    placeholder="/path/to/plot.png",
                    id="plot_path",
                    disabled=True,
                )
        with Horizontal(id="buttons"):
            yield Button("Process", variant="primary", id="process")
            yield Button("Quit", variant="error", id="quit")
        yield Log(id="log", highlight=True)
        yield Footer()

    def _refresh_output_path(self) -> None:
        """Recompute output and plot paths from current input path and function."""
        input_val = self.query_one("#input_path", Input).value.strip()
        fn_key = self.query_one("#function", Select).value
        if not input_val or fn_key is Select.BLANK:
            return
        p = Path(input_val)
        suffix = _SUFFIX_BY_KEY.get(str(fn_key), "")
        out = p.with_stem(p.stem + suffix)
        self.query_one("#output_path", Input).value = str(out)

    def on_select_changed(self, event: Select.Changed) -> None:
        """Recompute output path when function changes."""
        if event.select.id == "function":
            self._refresh_output_path()

    def on_input_changed(self, event: Input.Changed) -> None:
        """Auto-fill output and plot paths."""
        if event.input.id == "input_path":
            self._refresh_output_path()
        if event.input.id in {"input_path", "output_path"}:
            out_val = self.query_one("#output_path", Input).value.strip()
            if out_val:
                p = Path(out_val)
                if p.stem:
                    self.query_one("#plot_path", Input).value = str(
                        p.with_suffix(".png")
                    )

    def on_checkbox_changed(self, event: Checkbox.Changed) -> None:
        """Enable/disable the plot path input alongside the checkbox."""
        self.query_one("#plot_path", Input).disabled = not event.value

    def on_button_pressed(self, event: Button.Pressed) -> None:
        """Handle button presses."""
        if event.button.id == "quit":
            self.exit()
        elif event.button.id == "process":
            self._process()

    def _process(self) -> None:
        log = self.query_one("#log", Log)
        input_path = Path(self.query_one("#input_path", Input).value.strip())
        output_path = Path(self.query_one("#output_path", Input).value.strip())
        save_plot = self.query_one("#save-plot", Checkbox).value
        plot_path = Path(self.query_one("#plot_path", Input).value.strip())
        fn_key = self.query_one("#function", Select).value

        if fn_key is Select.BLANK:
            log.write_line("Error: no function selected.")
            return
        if not input_path.name:
            log.write_line("Error: no input file specified.")
            return
        if not input_path.exists():
            log.write_line(f"Error: {input_path} does not exist.")
            return
        if not output_path.name:
            log.write_line("Error: no output file specified.")
            return
        if save_plot and not plot_path.name:
            log.write_line("Error: no plot path specified.")
            return

        log.write_line(f"Loading  {input_path} …")
        waveform, sr = librosa.load(input_path, sr=None)
        log.write_line(
            f"Loaded   {len(waveform)} samples at {sr} Hz ({len(waveform) / sr:.2f}s)"
        )

        fn = getattr(distortion_tui, str(fn_key))
        fn_label = next(label for label, key, _ in _FUNCTIONS if key == str(fn_key))
        log.write_line(f"Applying {fn_label} …")
        processed = fn(waveform)
        log.write_line(f"Done     {len(processed)} samples output")

        output_path.parent.mkdir(parents=True, exist_ok=True)
        soundfile.write(output_path, processed, sr)
        log.write_line(f"Saved    {output_path}")

        if save_plot:
            payload = PlotPayload(
                waveform=waveform,
                processed=processed,
                label=fn_label,
                sr=sr,
            )
            self._save_plot(payload, plot_path=plot_path, log=log)

    @staticmethod
    def _save_plot(payload: PlotPayload, *, plot_path: Path, log: Log) -> None:
        log.write_line(f"Plotting {plot_path} …")
        fig, axes = plt.subplots(2, 1, figsize=(12, 6), sharex=False)

        librosa.display.waveshow(payload.waveform, sr=payload.sr, ax=axes[0])
        axes[0].set_title("Original")

        librosa.display.waveshow(payload.processed, sr=payload.sr, ax=axes[1])
        axes[1].set_title(payload.label)

        plt.tight_layout()
        plot_path.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(plot_path)
        plt.close(fig)
        log.write_line(f"Plot     {plot_path}")


def main() -> None:
    """Launch the TUI."""
    DistortionApp().run()


if __name__ == "__main__":
    main()
