#!/usr/bin/env python3

#---------------------------------
# Author: Ankit Anand
# Date: 06/11/25
# Email: ankit0.anand0@gmail.com
#---------------------------------

import sounddevice as sd
import numpy as np
import ipywidgets as widgets
from IPython.display import display, clear_output, Audio

class MicRecorder:
    
    def __init__(self):
        pass
        
    def record():
        """
        Create a UI to record audio in jupyter notebook, the 
        recorded signal is available as array.
    
        .. code-block:: python
            
            import modusa as ms
            result = ms.record()
            y, sr, title = result() # Keep it in the next cell
    
        Returns
        -------
        Callable
            A lambda function that returns y(audio signal), sr(sampling rate), title(title set in the UI)
        """
        import sounddevice as sd
        import numpy as np
        import ipywidgets as widgets
        from IPython.display import display, Audio, clear_output
        
        devices = sd.query_devices()
        device_options = [(d['name'][:30], i) for i, d in enumerate(devices) if d['max_input_channels'] > 0]
        
        # Controls
        device_dropdown = widgets.Dropdown(
            options=device_options,
            description="Microphone:",
            layout=widgets.Layout(width="300px")
        )
        
        sr_dropdown = widgets.Dropdown(
            options=[('16 Khz', 16000), ('22.05 Khz', 22050), ('44.1 Khz', 44100), ('48 Khz', 48000)],
            value=22050,
            description="Sample Rate:",
            layout=widgets.Layout(width="300px")
        )
        
        title_box = widgets.Text(
            placeholder="Title",
            description="Title:",
            layout=widgets.Layout(width="300px")
        )
        
        toggle_button = widgets.Button(description="Record", button_style="")
        status = widgets.HTML(value="")
        out = widgets.Output()
        
        # State
        result_container = {"result": None}
        stream = {"obj": None, "frames": [], "recording": False}
        
        def callback(indata, frames, time, status):
            if not status:
                stream["frames"].append(indata.copy())
                
        def on_toggle(b):
            if not stream["recording"]:
                stream["frames"].clear()
                sr = sr_dropdown.value
                device_id = device_dropdown.value
                
                stream["obj"] = sd.InputStream(callback=callback, channels=1, samplerate=sr, device=device_id)
                stream["obj"].start()
                stream["recording"] = True
                
                toggle_button.description = "Stop"
                toggle_button.button_style = "danger"
                status.value = "Recording..."
            else:
                stream["obj"].stop()
                stream["obj"].close()
                
                sr = sr_dropdown.value
                y = np.concatenate(stream["frames"], axis=0).flatten()
                title = title_box.value.strip() or "Recording"
                
                result_container["result"] = (y, sr, title)
                stream["recording"] = False
                
                toggle_button.description = "Record"
                toggle_button.button_style = "success"
                
                with out:
                    clear_output()
                    display(Audio(y, rate=sr))
                    
        toggle_button.on_click(on_toggle)
        
        # Layout
        ui = widgets.VBox([
            device_dropdown,
            sr_dropdown,
            title_box,
            widgets.HBox([toggle_button]),
            out
        ])
        
        display(ui)
        
        return lambda: result_container["result"]