//! `/asr` — browser-facing WS that proxies to the upstream ASR server.
//!
//! Upstream protocol (Qwen3-ASR streaming server on i7:9000):
//!   client → server: text JSON  `{"type":"init","language":"...","chunk_size_sec":N}`,
//!                    binary     Int16 LE PCM, 16 kHz mono,
//!                    text JSON  `{"type":"flush"}`
//!   server → client: text JSON  `{"text":"..."}`, `{"text":"...","final":true}`
//!
//! Browser only pushes PCM binary + optional flush. The proxy injects the
//! upstream `init` from server-side config.

use axum::{
    extract::{
        ws::{Message, WebSocket},
        State, WebSocketUpgrade,
    },
    response::Response,
};
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::{
    connect_async,
    tungstenite::{client::IntoClientRequest, Message as TgMessage},
};
use tracing::{debug, error, info, warn};

use crate::state::AppState;

pub async fn ws_handler(State(state): State<AppState>, ws: WebSocketUpgrade) -> Response {
    ws.on_upgrade(move |socket| handle_socket(socket, state))
}

async fn handle_socket(client_ws: WebSocket, state: AppState) {
    info!(upstream = %state.asr_upstream, "ASR proxy: dialing upstream");
    let request = match state.asr_upstream.as_str().into_client_request() {
        Ok(r) => r,
        Err(e) => {
            error!(error = %e, "failed to build upstream request");
            return;
        }
    };

    let (upstream_ws, _resp) = match connect_async(request).await {
        Ok(pair) => pair,
        Err(e) => {
            error!(error = %e, "ASR upstream connection failed");
            let _ = close_with_msg(client_ws, "ASR upstream unreachable").await;
            return;
        }
    };

    let (mut up_tx, mut up_rx) = upstream_ws.split();
    let (mut cl_tx, mut cl_rx) = client_ws.split();

    let init = serde_json::json!({
        "type": "init",
        "language": state.asr_language,
        "chunk_size_sec": state.asr_chunk_size_sec,
    })
    .to_string();
    if let Err(e) = up_tx.send(TgMessage::Text(init)).await {
        warn!(error = %e, "ASR upstream init send failed");
        return;
    }

    let to_upstream = async move {
        while let Some(msg) = cl_rx.next().await {
            match msg {
                Ok(Message::Binary(b)) => {
                    if up_tx.send(TgMessage::Binary(b.to_vec())).await.is_err() {
                        break;
                    }
                }
                Ok(Message::Text(t)) => {
                    if up_tx.send(TgMessage::Text(t.to_string())).await.is_err() {
                        break;
                    }
                }
                Ok(Message::Close(_)) => {
                    let _ = up_tx.send(TgMessage::Close(None)).await;
                    break;
                }
                Ok(Message::Ping(p)) => {
                    let _ = up_tx.send(TgMessage::Ping(p)).await;
                }
                Ok(_) => {}
                Err(e) => {
                    debug!(error = %e, "client ws recv error");
                    break;
                }
            }
        }
    };

    let to_client = async move {
        while let Some(msg) = up_rx.next().await {
            match msg {
                Ok(TgMessage::Text(t)) => {
                    if cl_tx
                        .send(Message::Text(t.as_str().to_string().into()))
                        .await
                        .is_err()
                    {
                        break;
                    }
                }
                Ok(TgMessage::Binary(b)) => {
                    if cl_tx
                        .send(Message::Binary(b.to_vec().into()))
                        .await
                        .is_err()
                    {
                        break;
                    }
                }
                Ok(TgMessage::Close(_)) => {
                    let _ = cl_tx.send(Message::Close(None)).await;
                    break;
                }
                Ok(_) => {}
                Err(e) => {
                    debug!(error = %e, "upstream ws recv error");
                    break;
                }
            }
        }
    };

    tokio::select! {
        () = to_upstream => warn!("ASR proxy: client side closed"),
        () = to_client => warn!("ASR proxy: upstream side closed"),
    }
}

async fn close_with_msg(mut ws: WebSocket, reason: &str) -> Result<(), axum::Error> {
    let _ = ws
        .send(Message::Text(
            format!("{{\"error\":\"{reason}\"}}").into(),
        ))
        .await;
    ws.send(Message::Close(None)).await
}
