mouth/src/overlay.rs

use std::num::NonZeroU32;
use tracing::{debug, error, info, warn};
use winit::application::ApplicationHandler;
use winit::dpi::{LogicalSize, PhysicalPosition};
use winit::event::WindowEvent;
use winit::event_loop::{ActiveEventLoop, EventLoop, EventLoopProxy};
use winit::window::{Window, WindowAttributes, WindowId, WindowLevel};

use crate::config::OverlayPosition;

const OVERLAY_WIDTH: u32 = 150;
const OVERLAY_HEIGHT: u32 = 18;

/// State of the overlay display.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum OverlayState {
    Hidden,
    Recording,
    Transcribing,
    Done,
    Error,
}

/// Events sent to the overlay from the coordinator.
#[derive(Debug, Clone)]
pub enum OverlayEvent {
    SetState(OverlayState),
    Shutdown,
}

/// The overlay application handler for winit.
struct OverlayApp {
    window: Option<std::rc::Rc<Window>>,
    surface: Option<softbuffer::Surface<std::rc::Rc<Window>, std::rc::Rc<Window>>>,
    state: OverlayState,
    position: OverlayPosition,
    _tray_icon: Option<tray_icon::TrayIcon>,
    tray_exit_id: Option<tray_icon::menu::MenuId>,
}

impl OverlayApp {
    fn draw(&mut self) {
        let Some(surface) = &mut self.surface else { return };
        let Some(window) = &self.window else { return };

        let size = window.inner_size();
        if size.width == 0 || size.height == 0 {
            return;
        }

        let Ok(w) = NonZeroU32::try_from(size.width) else { return };
        let Ok(h) = NonZeroU32::try_from(size.height) else { return };

        if surface.resize(w, h).is_err() {
            return;
        }

        let Ok(mut buffer) = surface.buffer_mut() else { return };

        let color = match self.state {
            OverlayState::Hidden => 0x00000000,
            OverlayState::Recording => 0xFFDD3333, // Red
            OverlayState::Transcribing => 0xFFDDAA33, // Amber
            OverlayState::Done => 0xFF33AA33, // Green
            OverlayState::Error => 0xFFDD3333, // Red
        };

        let width = size.width as usize;
        let height = size.height as usize;

        for y in 0..height {
            for x in 0..width {
                let radius = 8;
                let in_corner = (x < radius || x >= width - radius)
                    && (y < radius || y >= height - radius);

                let pixel = if in_corner {
                    let cx = if x < radius { radius } else { width - radius - 1 };
                    let cy = if y < radius { radius } else { height - radius - 1 };
                    let dx = x as i32 - cx as i32;
                    let dy = y as i32 - cy as i32;
                    if dx * dx + dy * dy <= (radius * radius) as i32 {
                        color
                    } else {
                        0x00000000
                    }
                } else {
                    color
                };

                buffer[y * width + x] = pixel;
            }
        }

        let _ = buffer.present();
    }

    fn update_visibility(&self) {
        if let Some(window) = &self.window {
            let visible = self.state != OverlayState::Hidden;
            window.set_visible(visible);
        }
    }

    fn create_tray_icon(&mut self) {
        use tray_icon::menu::{Menu, MenuItem};
        use tray_icon::TrayIconBuilder;

        let menu = Menu::new();
        let exit_item = MenuItem::new("Exit", true, None);
        let exit_id = exit_item.id().clone();
        if let Err(e) = menu.append(&exit_item) {
            warn!("Failed to add tray menu item: {e}");
            return;
        }

        let icon = match load_tray_icon() {
            Ok(i) => i,
            Err(e) => {
                warn!("Failed to load tray icon: {e}");
                return;
            }
        };

        match TrayIconBuilder::new()
            .with_menu(Box::new(menu))
            .with_tooltip("Mouth — Speech to Text")
            .with_icon(icon)
            .build()
        {
            Ok(tray) => {
                info!("System tray icon created");
                self._tray_icon = Some(tray);
                self.tray_exit_id = Some(exit_id);
            }
            Err(e) => {
                warn!("Failed to create tray icon: {e}");
            }
        }
    }
}

impl ApplicationHandler<OverlayEvent> for OverlayApp {
    fn resumed(&mut self, event_loop: &ActiveEventLoop) {
        if self.window.is_some() {
            return;
        }

        let attrs = WindowAttributes::default()
            .with_title("Mouth")
            .with_inner_size(LogicalSize::new(OVERLAY_WIDTH, OVERLAY_HEIGHT))
            .with_resizable(false)
            .with_decorations(false)
            .with_transparent(true)
            .with_window_level(WindowLevel::AlwaysOnTop)
            .with_visible(false);

        match event_loop.create_window(attrs) {
            Ok(window) => {
                let window = std::rc::Rc::new(window);

                // Position at top center of primary monitor
                if let Some(monitor) = window.current_monitor() {
                    let screen_size = monitor.size();
                    let pos = match self.position {
                        OverlayPosition::Top => PhysicalPosition::new(
                            (screen_size.width - OVERLAY_WIDTH) / 2,
                            10,
                        ),
                        OverlayPosition::Bottom => PhysicalPosition::new(
                            (screen_size.width - OVERLAY_WIDTH) / 2,
                            screen_size.height - OVERLAY_HEIGHT - 50,
                        ),
                        OverlayPosition::None => PhysicalPosition::new(0, 0),
                    };
                    window.set_outer_position(pos);
                }

                let context = softbuffer::Context::new(window.clone()).ok();
                let surface = context.and_then(|ctx| {
                    softbuffer::Surface::new(&ctx, window.clone()).ok()
                });

                if surface.is_none() {
                    warn!("Could not create softbuffer surface — overlay rendering disabled");
                }

                self.surface = surface;
                self.window = Some(window);
                info!("Overlay window created");
            }
            Err(e) => {
                error!("Failed to create overlay window: {e}");
            }
        }

        // Create tray icon (must be done on the main/event-loop thread)
        self.create_tray_icon();
    }

    fn user_event(&mut self, event_loop: &ActiveEventLoop, event: OverlayEvent) {
        match event {
            OverlayEvent::SetState(state) => {
                debug!("Overlay state: {:?} -> {:?}", self.state, state);
                self.state = state;
                self.update_visibility();
                self.draw();
            }
            OverlayEvent::Shutdown => {
                info!("Overlay shutting down");
                event_loop.exit();
            }
        }
    }

    fn window_event(&mut self, _event_loop: &ActiveEventLoop, _id: WindowId, event: WindowEvent) {
        if let WindowEvent::RedrawRequested = event {
            self.draw();
        }
    }

    fn about_to_wait(&mut self, event_loop: &ActiveEventLoop) {
        // Poll tray menu events
        if let Some(exit_id) = &self.tray_exit_id {
            if let Ok(event) = tray_icon::menu::MenuEvent::receiver().try_recv() {
                if event.id() == exit_id {
                    info!("Exit requested via tray icon");
                    crate::ipc::cleanup();
                    event_loop.exit();
                }
            }
        }
    }
}

fn load_tray_icon() -> Result<tray_icon::Icon, Box<dyn std::error::Error>> {
    const S: u32 = 32;
    let mut pixels = vec![0u8; (S * S * 4) as usize];

    let cx = S as f32 / 2.0;

    for y in 0..S {
        for x in 0..S {
            let fx = x as f32 + 0.5;
            let fy = y as f32 + 0.5;
            let idx = ((y * S + x) * 4) as usize;

            let mut alpha: f32 = 0.0;

            // Microphone body: rounded rectangle (capsule shape)
            // Center x=16, from y=3 to y=18, radius 5
            let mic_top = 3.0;
            let mic_bot = 18.0;
            let mic_r = 5.5;
            let mic_cx = cx;
            {
                let dy = fy.clamp(mic_top + mic_r, mic_bot - mic_r);
                let dist = ((fx - mic_cx).powi(2) + (fy - dy).powi(2)).sqrt();
                if dist <= mic_r {
                    alpha = 1.0;
                } else if dist <= mic_r + 1.0 {
                    alpha = alpha.max(mic_r + 1.0 - dist); // anti-alias
                }
            }

            // Cradle arc: U-shape below mic, from y=14 to y=22
            {
                let arc_cy = 14.0;
                let arc_r = 8.5;
                let arc_thickness = 2.2;
                let dx = fx - cx;
                let dy = fy - arc_cy;
                let dist = (dx * dx + dy * dy).sqrt();
                if fy >= arc_cy && dist >= arc_r - arc_thickness / 2.0 && dist <= arc_r + arc_thickness / 2.0 {
                    let edge_outer = (arc_r + arc_thickness / 2.0 - dist).min(1.0).max(0.0);
                    let edge_inner = (dist - (arc_r - arc_thickness / 2.0)).min(1.0).max(0.0);
                    alpha = alpha.max(edge_outer.min(edge_inner));
                }
            }

            // Stem: vertical line from arc bottom to near bottom
            {
                let stem_top = 22.0;
                let stem_bot = 27.0;
                let stem_w = 1.2;
                if fy >= stem_top && fy <= stem_bot && (fx - cx).abs() <= stem_w {
                    let edge = (stem_w - (fx - cx).abs()).min(1.0);
                    alpha = alpha.max(edge);
                }
            }

            // Base: horizontal line at bottom
            {
                let base_y = 27.0;
                let base_h = 2.0;
                let base_hw = 5.0;
                if fy >= base_y && fy <= base_y + base_h && (fx - cx).abs() <= base_hw {
                    let edge = (base_hw - (fx - cx).abs()).min(1.0);
                    alpha = alpha.max(edge);
                }
            }

            let a = (alpha.clamp(0.0, 1.0) * 255.0) as u8;
            // White icon with alpha (looks good on both light and dark taskbars)
            pixels[idx] = 255;     // R
            pixels[idx + 1] = 255; // G
            pixels[idx + 2] = 255; // B
            pixels[idx + 3] = a;   // A
        }
    }

    let icon = tray_icon::Icon::from_rgba(pixels, S, S)?;
    Ok(icon)
}

/// Create an event loop and return the proxy for sending events.
pub fn create_event_loop() -> Result<(EventLoop<OverlayEvent>, EventLoopProxy<OverlayEvent>), winit::error::EventLoopError> {
    let event_loop: EventLoop<OverlayEvent> = EventLoop::with_user_event().build()?;
    let proxy = event_loop.create_proxy();
    Ok((event_loop, proxy))
}

/// Run the event loop with the given position config.
pub fn run_event_loop(
    event_loop: EventLoop<OverlayEvent>,
    position: OverlayPosition,
) -> Result<(), winit::error::EventLoopError> {
    let mut app = OverlayApp {
        window: None,
        surface: None,
        state: OverlayState::Hidden,
        position,
        _tray_icon: None,
        tray_exit_id: None,
    };

    event_loop.run_app(&mut app)
}