use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message;

use super::WebSocket;
use tokio::sync::{mpsc, oneshot};

use super::frames::write_frames;
use crate::Frame;
use crate::engine::{self, Data, Packet as EnginePacket};
use crate::session::Input;

pub(super) async fn run_upgrade(
    input: mpsc::Sender<Input>,
    mut socket: WebSocket,
    out_tx: mpsc::Sender<Vec<Frame>>,
    mut outbound: mpsc::Receiver<Vec<Frame>>,
) {
    let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10);
    let mut out_tx = Some(out_tx);
    let mut probed = false;
    let mut upgraded = false;

    loop {
        tokio::select! {
            _ = input.closed() => break,
            _ = tokio::time::sleep_until(deadline), if !upgraded => break,
            incoming = socket.next() => {
                let Some(Ok(message)) = incoming else {
                    break;
                };
                match message {
                    Message::Text(text) => {
                        let Some(packet) = engine::decode_packet(&text) else {
                            break;
                        };
                        match packet {
                            EnginePacket::Ping(payload) if payload == "probe" && !upgraded => {
                                if !write_frames(&mut socket, vec![Frame::Text("3probe".to_string())]).await {
                                    break;
                                }
                                probed = true;
                            }
                            EnginePacket::Upgrade(_) if probed && !upgraded => {
                                let Some(tx) = out_tx.take() else { break };
                                let (result, receiver) = oneshot::channel();
                                if input.try_send(Input::Upgrade { tx, result }).is_err() {
                                    break;
                                }
                                if !matches!(tokio::time::timeout_at(deadline, receiver).await, Ok(Ok(true))) {
                                    break;
                                }
                                upgraded = true;
                            }
                            other if upgraded => {
                                if input.try_send(Input::Packets(vec![other])).is_err() {
                                    break;
                                }
                            }
                            _ => {}
                        }
                    }
                    Message::Binary(data) if upgraded => {
                        if input
                            .try_send(Input::Packets(vec![EnginePacket::Message(Data::Binary(
                                data.to_vec(),
                            ))]))
                            .is_err()
                        {
                            break;
                        }
                    }
                    Message::Binary(_) => {}
                    Message::Ping(_) => {
                        if !matches!(tokio::time::timeout(std::time::Duration::from_secs(5), socket.flush()).await, Ok(Ok(()))) {
                            break;
                        }
                    }
                    Message::Pong(_) | Message::Frame(_) => {}
                    Message::Close(_) => break,
                }
            }
            frames = outbound.recv() => {
                match frames {
                    Some(frames) => {
                        if !write_frames(&mut socket, frames).await {
                            break;
                        }
                    }
                    None => break,
                }
            }
        }
    }

    if upgraded {
        let _ =
            tokio::time::timeout(std::time::Duration::from_secs(5), input.send(Input::Close)).await;
    }
    let _ = tokio::time::timeout(
        std::time::Duration::from_secs(5),
        SinkExt::close(&mut socket),
    )
    .await;
}