use std::time::Duration;

use anyhow::Context;
use axum::extract::ws::{Message, WebSocket};
use base64::Engine;
use futures_util::{Sink, SinkExt, StreamExt};
use serde::Deserialize;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::task::JoinHandle;
use tokio_serial::{ClearBuffer, SerialPort, SerialPortBuilderExt};

use crate::{
    config::BoardConfig,
    power::{PowerAction, PowerActionError},
    serial::discovery::resolve_serial_config,
    session::SessionState,
    state::AppState,
};

const SERIAL_READ_BUFFER_SIZE: usize = 64;
const SERIAL_READ_TIMEOUT: Duration = Duration::from_millis(20);

#[derive(Debug, Deserialize)]
struct ClientControlMessage {
    #[serde(rename = "type")]
    kind: String,
    encoding: Option<String>,
    data: Option<String>,
}

pub async fn run_serial_ws(
    socket: WebSocket,
    state: AppState,
    session: std::sync::Arc<SessionState>,
) {
    let result = run_serial_ws_inner(socket, &state, session.clone()).await;
    session.clear_serial_connected();
    if let Err(err) = result {
        log::warn!("serial websocket ended with error: {err:#}");
    }
}

async fn run_serial_ws_inner(
    socket: WebSocket,
    state: &AppState,
    session: std::sync::Arc<SessionState>,
) -> anyhow::Result<()> {
    let session_id = session.snapshot().await.id;
    let board = session.board().clone();
    let serial = board
        .serial
        .as_ref()
        .ok_or_else(|| anyhow::anyhow!("board has no serial configuration"))?;
    let resolved_serial = resolve_serial_config(serial)?;
    let port = tokio_serial::new(&resolved_serial.current_device_path, serial.baud_rate)
        .timeout(SERIAL_READ_TIMEOUT)
        .open_native_async()
        .with_context(|| {
            format!(
                "failed to open serial port {}",
                resolved_serial.current_device_path
            )
        })?;

    let (mut ws_sender, mut ws_receiver) = socket.split();
    let (mut serial_rx, mut serial_tx) = tokio::io::split(port);
    let mut serial_buffer = [0u8; SERIAL_READ_BUFFER_SIZE];
    let mut power_on_task = Some(spawn_power_action_task(
        state.clone(),
        board.clone(),
        PowerAction::On,
    ));
    let power_linked = true;
    let mut shutdown_rx = session.subscribe_shutdown();

    ws_sender
        .send(Message::Text(r#"{"type":"opened"}"#.to_string().into()))
        .await
        .ok();
    let result = async {
        loop {
            if let Some(task) = power_on_task.as_mut() {
                tokio::select! {
                    power_result = task => {
                        power_on_task = None;
                        match power_result {
                            Ok(Ok(_)) => {}
                            Ok(Err(err)) => {
                                let message = format!("automatic power-on failed: {err}");
                                log::warn!("session `{session_id}` {message}");
                                send_power_on_failure_and_close(&mut ws_sender, &message).await;
                                break;
                            }
                            Err(err) => {
                                let message = format!("automatic power-on task join failed: {err}");
                                log::warn!("session `{session_id}` {message}");
                                send_power_on_failure_and_close(&mut ws_sender, &message).await;
                                break;
                            }
                        }
                    }
                    changed = shutdown_rx.changed() => {
                        if changed.is_ok() && *shutdown_rx.borrow() {
                            let _ = ws_sender
                                .send(Message::Text(r#"{"type":"closed"}"#.to_string().into()))
                                .await;
                            break;
                        }
                    }
                    maybe_message = ws_receiver.next() => {
                        let Some(message) = maybe_message else {
                            break;
                        };
                        match message {
                            Ok(Message::Binary(bytes)) => {
                                write_serial_payload(&mut serial_tx, &bytes).await?;
                            }
                            Ok(Message::Text(text)) => {
                                let control: ClientControlMessage = serde_json::from_str(&text)?;
                                match control.kind.as_str() {
                                    "close" => {
                                        let _ = ws_sender
                                            .send(Message::Text(r#"{"type":"closed"}"#.to_string().into()))
                                            .await;
                                        break;
                                    }
                                    "tx" => {
                                        let Some(data) = control.data.as_deref() else {
                                            anyhow::bail!("missing tx data");
                                        };
                                        let payload = match control.encoding.as_deref() {
                                            Some("base64") => base64::engine::general_purpose::STANDARD
                                                .decode(data)
                                                .context("invalid base64 payload")?,
                                            Some("utf8") | None => data.as_bytes().to_vec(),
                                            Some(other) => anyhow::bail!("unsupported encoding `{other}`"),
                                        };
                                        write_serial_payload(&mut serial_tx, &payload).await?;
                                    }
                                    other => anyhow::bail!("unsupported websocket control type `{other}`"),
                                }
                            }
                            Ok(Message::Close(_)) => break,
                            Ok(Message::Ping(payload)) => {
                                ws_sender.send(Message::Pong(payload)).await.ok();
                            }
                            Ok(Message::Pong(_)) => {}
                            Err(err) => return Err(err.into()),
                        }
                        let _ = session.heartbeat().await;
                    }
                    read = serial_rx.read(&mut serial_buffer) => {
                        let read = read.context("serial read failed")?;
                        if read == 0 {
                            break;
                        }
                        ws_sender
                            .send(Message::Binary(serial_buffer[..read].to_vec().into()))
                            .await
                            .context("failed to send serial output over websocket")?;
                        let _ = session.heartbeat().await;
                    }
                }
            } else {
                tokio::select! {
                    changed = shutdown_rx.changed() => {
                        if changed.is_ok() && *shutdown_rx.borrow() {
                            let _ = ws_sender
                                .send(Message::Text(r#"{"type":"closed"}"#.to_string().into()))
                                .await;
                            break;
                        }
                    }
                    maybe_message = ws_receiver.next() => {
                        let Some(message) = maybe_message else {
                            break;
                        };
                        match message {
                            Ok(Message::Binary(bytes)) => {
                                write_serial_payload(&mut serial_tx, &bytes).await?;
                            }
                            Ok(Message::Text(text)) => {
                                let control: ClientControlMessage = serde_json::from_str(&text)?;
                                match control.kind.as_str() {
                                    "close" => {
                                        let _ = ws_sender
                                            .send(Message::Text(r#"{"type":"closed"}"#.to_string().into()))
                                            .await;
                                        break;
                                    }
                                    "tx" => {
                                        let Some(data) = control.data.as_deref() else {
                                            anyhow::bail!("missing tx data");
                                        };
                                        let payload = match control.encoding.as_deref() {
                                            Some("base64") => base64::engine::general_purpose::STANDARD
                                                .decode(data)
                                                .context("invalid base64 payload")?,
                                            Some("utf8") | None => data.as_bytes().to_vec(),
                                            Some(other) => anyhow::bail!("unsupported encoding `{other}`"),
                                        };
                                        write_serial_payload(&mut serial_tx, &payload).await?;
                                    }
                                    other => anyhow::bail!("unsupported websocket control type `{other}`"),
                                }
                            }
                            Ok(Message::Close(_)) => break,
                            Ok(Message::Ping(payload)) => {
                                ws_sender.send(Message::Pong(payload)).await.ok();
                            }
                            Ok(Message::Pong(_)) => {}
                            Err(err) => return Err(err.into()),
                        }
                        let _ = session.heartbeat().await;
                    }
                    read = serial_rx.read(&mut serial_buffer) => {
                        let read = read.context("serial read failed")?;
                        if read == 0 {
                            break;
                        }
                        ws_sender
                            .send(Message::Binary(serial_buffer[..read].to_vec().into()))
                            .await
                            .context("failed to send serial output over websocket")?;
                        let _ = session.heartbeat().await;
                    }
                }
            }
        }

        Ok(())
    }
    .await;

    let result =
        finalize_power_linked_session(state, &board, power_linked, power_on_task, result).await;
    let mut port = serial_rx.unsplit(serial_tx);
    let result = preserve_result_after_serial_cleanup(&session_id, result, &mut port).await;
    let _ = state
        .request_session_stop(&session_id, crate::session::SessionStopReason::SerialClosed)
        .await;
    let _ = ws_sender.send(Message::Close(None)).await;
    result
}

fn spawn_power_action_task(
    state: AppState,
    board: BoardConfig,
    action: PowerAction,
) -> JoinHandle<Result<String, PowerActionError>> {
    tokio::spawn(async move { state.execute_board_power_action(&board, action).await })
}

async fn cleanup_power_link(
    board: &BoardConfig,
    power_linked: bool,
    power_on_task: Option<JoinHandle<Result<String, PowerActionError>>>,
) {
    if !power_linked {
        return;
    }

    if let Some(task) = power_on_task {
        match task.await {
            Ok(Ok(_)) => {}
            Ok(Err(err)) => {
                log::warn!(
                    "session `{}` power-on task ended with error: {err}",
                    board.id
                )
            }
            Err(err) => log::warn!("session `{}` power-on task join failed: {err}", board.id),
        }
    }
}

async fn finalize_power_linked_session<T>(
    _state: &AppState,
    board: &BoardConfig,
    power_linked: bool,
    power_on_task: Option<JoinHandle<Result<String, PowerActionError>>>,
    result: anyhow::Result<T>,
) -> anyhow::Result<T> {
    cleanup_power_link(board, power_linked, power_on_task).await;
    result
}

async fn send_power_on_failure_and_close<S>(ws_sender: &mut S, message: &str)
where
    S: Sink<Message> + Unpin,
{
    let payload = serde_json::json!({
        "type": "error",
        "message": message,
    })
    .to_string();
    let _ = ws_sender.send(Message::Text(payload.into())).await;
    let _ = ws_sender
        .send(Message::Text(r#"{"type":"closed"}"#.to_string().into()))
        .await;
    let _ = ws_sender.send(Message::Close(None)).await;
}

async fn write_serial_payload(
    port: &mut tokio::io::WriteHalf<tokio_serial::SerialStream>,
    payload: &[u8],
) -> anyhow::Result<()> {
    port.write_all(payload).await?;
    port.flush().await?;
    Ok(())
}

#[async_trait::async_trait]
trait SerialQueueCleanup {
    async fn flush_output(&mut self) -> std::io::Result<()>;
    fn clear_all_buffers(&mut self) -> std::io::Result<()>;
}

#[async_trait::async_trait]
impl SerialQueueCleanup for tokio_serial::SerialStream {
    async fn flush_output(&mut self) -> std::io::Result<()> {
        AsyncWriteExt::flush(self).await
    }

    fn clear_all_buffers(&mut self) -> std::io::Result<()> {
        self.clear(ClearBuffer::All).map_err(std::io::Error::from)
    }
}

async fn cleanup_serial_queue_before_close<T>(port: &mut T) -> anyhow::Result<()>
where
    T: SerialQueueCleanup + ?Sized,
{
    port.flush_output()
        .await
        .context("failed to flush serial output before close")?;
    port.clear_all_buffers()
        .context("failed to clear serial buffers before close")?;
    Ok(())
}

async fn preserve_result_after_serial_cleanup<T, P>(
    session_id: &str,
    result: anyhow::Result<T>,
    port: &mut P,
) -> anyhow::Result<T>
where
    P: SerialQueueCleanup + ?Sized,
{
    if let Err(err) = cleanup_serial_queue_before_close(port).await {
        log::warn!("session `{session_id}` failed to clean serial queue before close: {err:#}");
    }
    result
}

#[cfg(test)]
mod tests {
    use std::{
        fs, io,
        pin::Pin,
        sync::{Arc, Mutex},
        task::{Context, Poll},
        time::Duration,
    };

    use axum::extract::ws::Message;
    use futures_util::Sink;
    use tempfile::tempdir;

    use super::{
        ClientControlMessage, SerialQueueCleanup, cleanup_power_link,
        cleanup_serial_queue_before_close, finalize_power_linked_session,
        preserve_result_after_serial_cleanup, send_power_on_failure_and_close,
    };
    use crate::{
        build_app_state,
        config::{
            BoardConfig, BootConfig, BuiltinTftpConfig, CustomPowerManagement,
            PowerManagementConfig, PxeProfile, ServerConfig, TftpConfig,
        },
        power::PowerActionError,
        tftp::service::{TftpManager, build_tftp_manager},
    };

    #[derive(Default)]
    struct VecSink {
        messages: Vec<Message>,
    }

    #[derive(Debug, Clone, PartialEq, Eq)]
    enum CleanupEvent {
        Flush,
        ClearAll,
    }

    struct RecordingSerialCleanup {
        events: Arc<Mutex<Vec<CleanupEvent>>>,
        clear_result: io::Result<()>,
    }

    #[async_trait::async_trait]
    impl SerialQueueCleanup for RecordingSerialCleanup {
        async fn flush_output(&mut self) -> io::Result<()> {
            self.events.lock().unwrap().push(CleanupEvent::Flush);
            Ok(())
        }

        fn clear_all_buffers(&mut self) -> io::Result<()> {
            self.events.lock().unwrap().push(CleanupEvent::ClearAll);
            self.clear_result
                .as_ref()
                .map(|_| ())
                .map_err(|err| io::Error::new(err.kind(), err.to_string()))
        }
    }

    impl Sink<Message> for VecSink {
        type Error = ();

        fn poll_ready(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
            self.get_mut().messages.push(item);
            Ok(())
        }

        fn poll_flush(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn poll_close(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }
    }

    #[test]
    fn control_message_parses_close_type() {
        let message: ClientControlMessage = serde_json::from_str(r#"{"type":"close"}"#).unwrap();
        assert_eq!(message.kind, "close");
    }

    #[tokio::test]
    async fn serial_cleanup_flushes_before_clearing_all_buffers() {
        let events = Arc::new(Mutex::new(Vec::new()));
        let mut cleanup = RecordingSerialCleanup {
            events: events.clone(),
            clear_result: Ok(()),
        };

        cleanup_serial_queue_before_close(&mut cleanup)
            .await
            .unwrap();

        assert_eq!(
            events.lock().unwrap().as_slice(),
            &[CleanupEvent::Flush, CleanupEvent::ClearAll]
        );
    }

    #[tokio::test]
    async fn serial_cleanup_reports_clear_failures() {
        let events = Arc::new(Mutex::new(Vec::new()));
        let mut cleanup = RecordingSerialCleanup {
            events: events.clone(),
            clear_result: Err(io::Error::new(io::ErrorKind::Other, "clear failed")),
        };

        let err = cleanup_serial_queue_before_close(&mut cleanup)
            .await
            .unwrap_err();

        assert!(err.to_string().contains("failed to clear serial buffers"));
        assert_eq!(
            events.lock().unwrap().as_slice(),
            &[CleanupEvent::Flush, CleanupEvent::ClearAll]
        );
    }

    #[tokio::test]
    async fn serial_cleanup_failure_preserves_original_session_error() {
        let events = Arc::new(Mutex::new(Vec::new()));
        let mut cleanup = RecordingSerialCleanup {
            events,
            clear_result: Err(io::Error::new(io::ErrorKind::Other, "clear failed")),
        };

        let err = preserve_result_after_serial_cleanup::<(), _>(
            "session-1",
            Err(anyhow::anyhow!("websocket failed")),
            &mut cleanup,
        )
        .await
        .unwrap_err();

        assert_eq!(err.to_string(), "websocket failed");
    }

    async fn test_state(root: &std::path::Path) -> crate::AppState {
        let config_path = root.join(".ostool-server.toml");
        let config = ServerConfig {
            listen_addr: "127.0.0.1:0".parse().unwrap(),
            data_dir: root.join("data"),
            board_dir: root.join("boards"),
            dtb_dir: root.join("dtbs"),
            tftp: TftpConfig::Builtin(BuiltinTftpConfig::default_with_root(root.join("tftp"))),
            ..ServerConfig::default()
        };
        let manager: std::sync::Arc<dyn TftpManager> = build_tftp_manager(&config.tftp);
        build_app_state(config_path, config, manager).await.unwrap()
    }

    #[tokio::test]
    async fn cleanup_waits_for_power_on_task_before_power_off() {
        let dir = tempdir().unwrap();
        let output_path = dir.path().join("power.log");
        let board = BoardConfig {
            id: "demo".into(),
            board_type: "demo".into(),
            tags: vec![],
            serial: None,
            power_management: PowerManagementConfig::Custom(CustomPowerManagement {
                power_on_cmd: String::new(),
                power_off_cmd: format!("printf 'off\\n' >> {}", output_path.display()),
            }),
            boot: BootConfig::Pxe(PxeProfile::default()),
            notes: None,
            disabled: false,
        };

        let power_on_task = tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(20)).await;
            fs::write(&output_path, "on\n").unwrap();
            Ok::<String, PowerActionError>("executed".into())
        });

        cleanup_power_link(&board, true, Some(power_on_task)).await;

        let content = fs::read_to_string(dir.path().join("power.log")).unwrap();
        assert_eq!(content, "on\n");
    }

    #[tokio::test]
    async fn finalize_runs_power_off_even_when_session_errors() {
        let dir = tempdir().unwrap();
        let output_path = dir.path().join("power.log");
        let state = test_state(dir.path()).await;
        let board = BoardConfig {
            id: "demo".into(),
            board_type: "demo".into(),
            tags: vec![],
            serial: None,
            power_management: PowerManagementConfig::Custom(CustomPowerManagement {
                power_on_cmd: String::new(),
                power_off_cmd: format!("printf 'off\\n' >> {}", output_path.display()),
            }),
            boot: BootConfig::Pxe(PxeProfile::default()),
            notes: None,
            disabled: false,
        };

        let power_on_task =
            tokio::spawn(async { Ok::<String, PowerActionError>("executed".into()) });
        let result = finalize_power_linked_session::<()>(
            &state,
            &board,
            true,
            Some(power_on_task),
            Err(anyhow::anyhow!("websocket send failed")),
        )
        .await;

        assert!(result.is_err());
        assert!(!output_path.exists());
    }

    #[tokio::test]
    async fn power_on_failure_sends_error_then_close_messages() {
        let mut sender = VecSink::default();
        send_power_on_failure_and_close(&mut sender, "automatic power-on failed").await;
        let mut messages = sender.messages.into_iter();
        let first = messages.next().unwrap();
        let second = messages.next().unwrap();
        let third = messages.next().unwrap();

        match first {
            Message::Text(text) => assert!(text.contains(r#""type":"error""#)),
            other => panic!("unexpected first message: {other:?}"),
        }
        match second {
            Message::Text(text) => assert_eq!(text, r#"{"type":"closed"}"#),
            other => panic!("unexpected second message: {other:?}"),
        }
        assert!(matches!(third, Message::Close(_)));
    }
}
