use anyhow::{Context, Result};
use bytes::BytesMut;
use futures_util::{SinkExt, StreamExt};
use pulse_core::{
    codec::{Encoding, SessionCodec},
    control::{AcceptPayload, ConnectPayload, StreamStatus},
    plp::{Header, PLP_VERSION, Packet, PacketFlags, PacketType},
    plp_conn::{ConnError, ConnEvent, PlpConn},
};
use std::collections::HashMap;
use std::{net::SocketAddr, sync::Arc};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::oneshot;
use tokio_tungstenite::{accept_async, tungstenite::Message};
use tracing::{debug, error, info, warn};

/// Minimal browser gateway: accepts PLP binary frames over WebSocket.
///
/// This is intentionally small and stateful-per-connection:
/// - CONNECT -> ACCEPT (JSON payload)
/// - PING -> PONG
/// - STREAM_DATA -> apply WASM transforms -> echo STREAM_DATA
pub async fn run(bind_addr: SocketAddr, state: Arc<crate::AppState>) -> Result<()> {
    let listener = TcpListener::bind(bind_addr)
        .await
        .with_context(|| format!("failed to bind WebSocket gateway on tcp://{bind_addr}"))?;

    info!("[WS] WebSocket gateway listening on tcp://{bind_addr}");

    // Never shuts down in production path.
    let (_tx, rx) = oneshot::channel::<()>();
    run_with_listener(listener, state, rx).await
}

pub async fn run_with_listener(
    listener: TcpListener,
    state: Arc<crate::AppState>,
    mut shutdown: oneshot::Receiver<()>,
) -> Result<()> {
    loop {
        tokio::select! {
            _ = &mut shutdown => {
                info!("[WS] WebSocket gateway shutting down");
                return Ok(());
            }
            accept = listener.accept() => {
                let (stream, peer) = accept?;
                let state = Arc::clone(&state);
                tokio::spawn(async move {
                    if let Err(e) = handle_conn(stream, peer, state).await {
                        warn!("[WS] connection {peer} ended with error: {e:#}");
                    }
                });
            }
        }
    }
}

async fn handle_conn(
    stream: TcpStream,
    peer: SocketAddr,
    state: Arc<crate::AppState>,
) -> Result<()> {
    // Protocol detection: peek at the first bytes to determine HTTP vs WebSocket.
    //
    // Both WebSocket and plain HTTP start with an HTTP request. The difference is
    // the presence of the "Upgrade: websocket" header. We peek instead of reading
    // so the bytes remain in the buffer for the actual parser.
    if crate::http_listener::is_http_request(&stream).await {
        if !crate::http_listener::is_websocket_upgrade(&stream).await {
            // Plain HTTP request — handle webhooks, health checks, etc.
            debug!("[WS] detected plain HTTP from {peer}, routing to HTTP handler");
            let registry = Arc::clone(&state.service_state.webhook_registry);
            crate::http_listener::handle_http_request(stream, registry).await;
            return Ok(());
        }
    }

    // WebSocket upgrade path (existing behavior).
    let mut ws = accept_async(stream)
        .await
        .context("websocket accept failed")?;

    debug!("[WS] accepted websocket from {peer}");

    let mut buf = BytesMut::with_capacity(16 * 1024);
    let mut conn_sm = PlpConn::new_server();
    let mut session_id = String::new();

    let mut sfu_rx: Option<tokio::sync::mpsc::Receiver<crate::sfu::ForwardPacket>> = None;
    let mut name_to_id = HashMap::<String, u16>::new();
    let mut id_to_name = HashMap::<u16, String>::new();

    loop {
        tokio::select! {
            Some(fp) = async {
                if let Some(rx) = sfu_rx.as_mut() {
                    rx.recv().await
                } else {
                    futures_util::future::pending().await
                }
            } => {
                // Forwarded packet from another connection!
                if let Some(&local_stream_id) = name_to_id.get(&fp.stream_id) {
                    let resp = Packet {
                        header: Header {
                            version: PLP_VERSION,
                            packet_type: PacketType::StreamData,
                            flags: PacketFlags::empty(),
                            stream_id: local_stream_id,
                            sequence_number: fp.sequence,
                            payload_length: fp.payload.len() as u64,
                        },
                        payload: bytes::Bytes::from(fp.payload),
                    };
                    ws.send(Message::Binary(encode_packet(&resp).into())).await?;
                }
            }
            msg_res = ws.next() => {
                let Some(msg_res) = msg_res else { break };
                let msg = msg_res.context("websocket read failed")?;
                match msg {
            Message::Binary(b) => {
                buf.extend_from_slice(&b);
                while let Some(pkt) = decode_one(&mut buf)? {
                    match pkt.header.packet_type {
                        PacketType::Connect => {
                            match conn_sm.on_inbound(&pkt) {
                                Ok(Some(ConnEvent::Handshake(
                                    pulse_core::plp_handshake::HandshakeEvent::ConnectReceived,
                                ))) => {}
                                Ok(_) => continue,
                                Err(_) => continue,
                            }

                            let connect: ConnectPayload = serde_json::from_slice(&pkt.payload)
                                .context("failed to parse CONNECT payload JSON")?;

                            // Determine the encoding negotiated for this session.
                            let negotiated = connect
                                .encoding
                                .as_deref()
                                .map(Encoding::from_wire_name)
                                .unwrap_or(Encoding::Json);

                            session_id = format!("{peer}-ws");

                            info!(
                                "[WS] CONNECT from {peer} — encoding={negotiated}, session={session_id}"
                            );

                            // ── Auth verification ───────────────────────────
                            let auth_result = state.auth_verifier.verify(connect.auth.as_ref()).await;
                            let auth_user = match auth_result {
                                crate::auth::AuthResult::Authenticated(user) => {
                                    info!("[WS] CONNECT auth OK: user={}", user.id);
                                    Some(user)
                                }
                                crate::auth::AuthResult::Skipped => None,
                                crate::auth::AuthResult::Rejected(reason) => {
                                    warn!("[WS] CONNECT auth REJECTED from {peer}: {reason}");
                                    let reject = pulse_core::control::RejectPayload {
                                        reason: format!("authentication failed: {reason}"),
                                    };
                                    let payload = serde_json::to_vec(&reject)
                                        .context("failed to encode REJECT payload JSON")?;
                                    let resp = Packet {
                                        header: Header {
                                            version: PLP_VERSION,
                                            packet_type: PacketType::Reject,
                                            flags: PacketFlags::empty(),
                                            stream_id: pkt.header.stream_id,
                                            sequence_number: pkt.header.sequence_number,
                                            payload_length: payload.len() as u64,
                                        },
                                        payload: bytes::Bytes::from(payload),
                                    };
                                    ws.send(Message::Binary(encode_packet(&resp).into()))
                                        .await?;
                                    break; // close connection after reject
                                }
                            };

                            let streams: Vec<StreamStatus> = connect
                                .streams
                                .iter()
                                .enumerate()
                                .map(|(i, s)| StreamStatus {
                                    id: (i + 1) as u16,
                                    stream_type: s.stream_type.clone(),
                                    accepted: true,
                                })
                                .collect();

                            let accept = AcceptPayload {
                                session_id: session_id.clone(),
                                streams,
                                relay_path: vec!["pulse-relay".to_string()],
                                metrics_endpoint: "/api/v1/metrics".to_string(),
                                encoding: connect.encoding.clone(),
                                user: auth_user,
                            };

                            let payload = serde_json::to_vec(&accept)
                                .context("failed to encode ACCEPT payload JSON")?;

                            let resp = Packet {
                                header: Header {
                                    version: PLP_VERSION,
                                    packet_type: PacketType::Accept,
                                    flags: PacketFlags::empty(),
                                    stream_id: pkt.header.stream_id,
                                    sequence_number: pkt.header.sequence_number,
                                    payload_length: payload.len() as u64,
                                },
                                payload: bytes::Bytes::from(payload),
                            };

                            ws.send(Message::Binary(encode_packet(&resp).into()))
                                .await?;
                            let _ = conn_sm.on_outbound(PacketType::Accept, 0);

                            info!("[WS] CONNECT accepted for session {session_id}");

                            // Initialize SFU participant for this new session
                            sfu_rx = Some(state.sfu.add_participant(session_id.clone()).await);
                        }
                        PacketType::StreamOpen => {
                            match conn_sm.on_inbound(&pkt) {
                                Ok(_) => {}
                                Err(ConnError::HandshakeNotAccepted) => continue,
                                Err(_) => continue,
                            }
                            // Parse payload to get stream name
                            let payload_text = String::from_utf8_lossy(&pkt.payload);
                            if let Ok(json) = serde_json::from_str::<serde_json::Value>(&payload_text) {
                                if let Some(name) = json.get("name").and_then(|n| n.as_str()) {
                                    let stream_name = name.to_string();
                                    name_to_id.insert(stream_name.clone(), pkt.header.stream_id);
                                    id_to_name.insert(pkt.header.stream_id, stream_name.clone());
                                    
                                    // Subscribe and publish to the SFU
                                    state.sfu.publish(&session_id, stream_name.clone()).await;
                                    state.sfu.subscribe(&session_id, stream_name).await;
                                }
                            }
                        }
                        PacketType::Ping => {
                            let resp = Packet {
                                header: Header {
                                    version: PLP_VERSION,
                                    packet_type: PacketType::Pong,
                                    flags: PacketFlags::empty(),
                                    stream_id: pkt.header.stream_id,
                                    sequence_number: pkt.header.sequence_number,
                                    payload_length: 0,
                                },
                                payload: bytes::Bytes::new(),
                            };
                            ws.send(Message::Binary(encode_packet(&resp).into()))
                                .await?;
                        }
                        PacketType::StreamData => {
                            match conn_sm.on_inbound(&pkt) {
                                Ok(_) => {}
                                Err(ConnError::HandshakeNotAccepted) => {
                                    warn!("[WS] STREAM_DATA before CONNECT from {peer}; dropping");
                                    continue;
                                }
                                Err(ConnError::Stream(
                                    pulse_core::plp_stream::StreamError::NotOpen,
                                )) => {
                                    warn!(
                                        "[WS] STREAM_DATA before STREAM_OPEN from {peer} stream_id={}; dropping",
                                        pkt.header.stream_id
                                    );
                                    continue;
                                }
                                Err(_) => continue,
                            }

                            let fp = crate::sfu::ForwardPacket {
                                source: peer.to_string(),
                                stream_id: pkt.header.stream_id.to_string(),
                                sequence: pkt.header.sequence_number,
                                payload: pkt.payload.to_vec(),
                                lossless: false,
                            };

                            let mut final_fp = fp;
                            match state.media_pipeline.process_frame(final_fp.clone()).await {
                                Ok(processed) => {
                                    final_fp.payload = processed.payload.clone();
                                    let resp = Packet {
                                        header: Header {
                                            version: PLP_VERSION,
                                            packet_type: PacketType::StreamData,
                                            flags: pkt.header.flags,
                                            stream_id: pkt.header.stream_id,
                                            sequence_number: pkt.header.sequence_number,
                                            payload_length: processed.payload.len() as u64,
                                        },
                                        payload: bytes::Bytes::from(processed.payload),
                                    };

                                    // Echo back to sender
                                    ws.send(Message::Binary(encode_packet(&resp).into()))
                                        .await?;
                                }
                                Err(e) => {
                                    error!("[WS] WASM transform failed for {peer}: {e}");
                                }
                            }
                            
                            // Broadcast to other subscribers!
                            if let Some(stream_name) = id_to_name.get(&pkt.header.stream_id) {
                                let mut bcast_fp = final_fp.clone();
                                bcast_fp.source = session_id.clone();
                                bcast_fp.stream_id = stream_name.clone();
                                bcast_fp.lossless = true; // Use lossless for reliable data
                                state.sfu.forward(bcast_fp).await;
                            }
                        }
                        PacketType::StreamClose | PacketType::StreamReset => {
                            match conn_sm.on_inbound(&pkt) {
                                Ok(_) => {}
                                Err(ConnError::HandshakeNotAccepted) => continue,
                                Err(_) => continue,
                            }
                            
                            if let Some(stream_name) = id_to_name.remove(&pkt.header.stream_id) {
                                name_to_id.remove(&stream_name);
                            }

                            let resp = Packet {
                                header: Header {
                                    version: PLP_VERSION,
                                    packet_type: pkt.header.packet_type,
                                    flags: PacketFlags::empty(),
                                    stream_id: pkt.header.stream_id,
                                    sequence_number: pkt.header.sequence_number,
                                    payload_length: pkt.payload.len() as u64,
                                },
                                payload: pkt.payload.clone(),
                            };
                            ws.send(Message::Binary(encode_packet(&resp).into()))
                                .await?;
                        }
                        PacketType::Metrics => {
                            if conn_sm.handshake_state()
                                != pulse_core::plp_handshake::HandshakeState::Accepted
                            {
                                continue;
                            }

                            let v: serde_json::Value = match serde_json::from_slice(&pkt.payload) {
                                Ok(v) => v,
                                Err(_) => continue,
                            };

                            let rtt_ms =
                                v.get("rtt_ms").and_then(|x| x.as_u64()).unwrap_or(0) as u32;
                            let packet_loss_pct =
                                v.get("packet_loss_pct")
                                    .and_then(|x| x.as_f64())
                                    .unwrap_or(0.0) as f32;

                            let mut bitrates: HashMap<String, u32> = HashMap::new();
                            if let Some(obj) = v.get("bitrates").and_then(|x| x.as_object()) {
                                for (k, vv) in obj.iter() {
                                    if let Some(n) = vv.as_u64() {
                                        bitrates.insert(k.clone(), n as u32);
                                    }
                                }
                            }

                            let control_url = state.control_url.clone();
                            let session_id = session_id.clone();
                            tokio::spawn(async move {
                                let client = reqwest::Client::new();
                                let url = format!("{}/api/v1/metrics", control_url);
                                let body = serde_json::json!({
                                    "session_id": session_id,
                                    "rtt_ms": rtt_ms,
                                    "packet_loss_pct": packet_loss_pct,
                                    "bitrates": bitrates,
                                });
                                let _ = client.post(url).json(&body).send().await;
                            });
                        }
                        PacketType::ServiceRequest => {
                            if conn_sm.handshake_state()
                                != pulse_core::plp_handshake::HandshakeState::Accepted
                            {
                                continue;
                            }

                            let req: crate::service::ServiceRequest =
                                match serde_json::from_slice(&pkt.payload) {
                                    Ok(r) => r,
                                    Err(_) => continue,
                                };

                            let service_state = Arc::clone(&state.service_state);
                            let seq = pkt.header.sequence_number;
                            let resp =
                                crate::service::handle_service_request(req, &service_state).await;
                            let resp_pkt = crate::service::encode_service_response(&resp, seq);
                            ws.send(Message::Binary(encode_packet(&resp_pkt).into()))
                                .await?;
                        }
                        PacketType::HttpOutbound => {
                            if conn_sm.handshake_state()
                                != pulse_core::plp_handshake::HandshakeState::Accepted
                            {
                                continue;
                            }

                            let req: crate::http_bridge::HttpOutboundRequest =
                                match serde_json::from_slice(&pkt.payload) {
                                    Ok(r) => r,
                                    Err(_) => continue,
                                };

                            let config = Arc::clone(&state.http_proxy_config);
                            let seq = pkt.header.sequence_number;
                            let resp =
                                crate::http_bridge::handle_http_outbound(req, &config, &session_id)
                                    .await;
                            let resp_pkt = crate::http_bridge::encode_http_response(&resp, seq);
                            ws.send(Message::Binary(encode_packet(&resp_pkt).into()))
                                .await?;
                        }
                        other => {
                            debug!("[WS] ignoring packet type {:?} from {peer}", other);
                        }
                    }
                }
            }
            Message::Text(_) => {
                // Browser clients should use binary PLP frames.
                continue;
            }
            Message::Close(_) => break,
            Message::Ping(payload) => {
                ws.send(Message::Pong(payload)).await?;
            }
            Message::Pong(_) => {}
            _ => {}
        }
    }
    }
    }
    
    // Clean up SFU participant
    if !session_id.is_empty() {
        state.sfu.remove_participant(&session_id).await;
    }

    info!(
        "[WS] session ended: {}",
        if session_id.is_empty() {
            peer.to_string()
        } else {
            session_id
        }
    );
    Ok(())
}

fn encode_packet(pkt: &Packet) -> Vec<u8> {
    let mut out = BytesMut::with_capacity((pkt.header.payload_length as usize) + 32);
    pkt.encode(&mut out);
    out.to_vec()
}

fn decode_one(buf: &mut BytesMut) -> Result<Option<Packet>> {
    let mut slice: &[u8] = &buf[..];
    match Packet::decode(&mut slice) {
        Ok(pkt) => {
            let consumed = buf.len() - slice.len();
            let _ = buf.split_to(consumed);
            Ok(Some(pkt))
        }
        Err(pulse_core::plp::Error::BufferTooShort) => Ok(None),
        Err(e) => Err(anyhow::anyhow!(e)).context("invalid PLP frame"),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;
    use pulse_core::plp::Packet as PlpPacket;
    use std::time::Duration;
    use tokio::sync::RwLock;
    use tokio_tungstenite::connect_async;

    fn encode_plp(pkt: &PlpPacket) -> Bytes {
        let mut out = BytesMut::with_capacity((pkt.header.payload_length as usize) + 32);
        pkt.encode(&mut out);
        out.freeze()
    }

    fn decode_plp(buf: &[u8]) -> PlpPacket {
        let mut slice: &[u8] = buf;
        PlpPacket::decode(&mut slice).expect("decode plp")
    }

    #[tokio::test]
    async fn ws_gateway_connect_accept_and_streamdata_echo() {
        // Build minimal AppState.
        let gpu_manager = Arc::new(RwLock::new(crate::gpu::GpuManager::new()));
        let ai = Arc::new(crate::ai_processor::AiProcessor::new(
            crate::ai_processor::AiConfig {
                stt_endpoint: None,
                vision_endpoint: None,
                tts_endpoint: None,
            },
        ));
        let engine = Arc::new(crate::edge_wasm::EdgeWasmEngine::new().expect("wasm engine"));
        let media_pipeline = Arc::new(crate::media::MediaPipeline::new(engine));

        let state = Arc::new(crate::AppState {
            gpu_manager,
            ai_processor: ai,
            media_pipeline,
            control_url: "http://localhost:4002".to_string(),
            sfu: Arc::new(crate::sfu::Sfu::new()),
            max_payload_bytes: 1024 * 1024,
            service_state: Arc::new(crate::service::ServiceState::new(
                Arc::new(pulse_core::queue_storage::InMemoryBackend::new()),
            )),
            http_proxy_config: Arc::new(crate::http_bridge::HttpProxyConfig::default()),
            queue_backend: Arc::new(pulse_core::queue_storage::InMemoryBackend::new()),
            auth_verifier: crate::auth::AuthVerifier::new(crate::auth::AuthConfig::default()),
        });

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();

        let server_state = Arc::clone(&state);
        let server = tokio::spawn(async move {
            run_with_listener(listener, server_state, shutdown_rx)
                .await
                .expect("ws gateway run");
        });

        let url = format!("ws://{}", addr);
        let (mut ws, _resp) = connect_async(&url).await.expect("connect_async");

        // Send CONNECT.
        let connect_json = serde_json::json!({
            "version": 1,
            "token": "",
            "capabilities": ["plp/1"],
            "streams": [],
        })
        .to_string();

        let connect_pkt = PlpPacket {
            header: Header {
                version: PLP_VERSION,
                packet_type: PacketType::Connect,
                flags: PacketFlags::empty(),
                stream_id: 0,
                sequence_number: 0,
                payload_length: connect_json.as_bytes().len() as u64,
            },
            payload: Bytes::from(connect_json.into_bytes()),
        };
        ws.send(Message::Binary(encode_plp(&connect_pkt)))
            .await
            .unwrap();

        // Expect ACCEPT.
        let msg = tokio::time::timeout(Duration::from_secs(2), ws.next())
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        let accept_bin = match msg {
            Message::Binary(b) => b,
            other => panic!("expected binary accept, got {other:?}"),
        };
        let accept = decode_plp(&accept_bin);
        assert_eq!(accept.header.packet_type, PacketType::Accept);

        // Send STREAM_OPEN then STREAM_DATA and expect echo STREAM_DATA.
        let open_json = serde_json::json!({
            "stream_id": 1,
            "type": "data",
            "reliability": "full",
            "priority": 0,
        })
        .to_string();
        let mut open_payload = vec![0x01]; // RPC Mode
        open_payload.extend_from_slice(open_json.as_bytes());

        let open_pkt = PlpPacket {
            header: Header {
                version: PLP_VERSION,
                packet_type: PacketType::StreamOpen,
                flags: PacketFlags::empty(),
                stream_id: 1,
                sequence_number: 0,
                payload_length: open_payload.len() as u64,
            },
            payload: Bytes::from(open_payload),
        };
        ws.send(Message::Binary(encode_plp(&open_pkt)))
            .await
            .unwrap();

        let payload = Bytes::from_static(b"abc");
        let data_pkt = PlpPacket {
            header: Header {
                version: PLP_VERSION,
                packet_type: PacketType::StreamData,
                flags: PacketFlags::empty(),
                stream_id: 1,
                sequence_number: 1,
                payload_length: payload.len() as u64,
            },
            payload,
        };
        ws.send(Message::Binary(encode_plp(&data_pkt)))
            .await
            .unwrap();

        let msg = tokio::time::timeout(Duration::from_secs(2), ws.next())
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        let echoed_bin = match msg {
            Message::Binary(b) => b,
            other => panic!("expected binary echo, got {other:?}"),
        };
        let echoed = decode_plp(&echoed_bin);
        assert_eq!(echoed.header.packet_type, PacketType::StreamData);
        assert_eq!(echoed.header.stream_id, 1);
        assert_eq!(echoed.payload.as_ref(), b"abc");

        // Close the stream and expect STREAM_CLOSE echo.
        let close_pkt = PlpPacket {
            header: Header {
                version: PLP_VERSION,
                packet_type: PacketType::StreamClose,
                flags: PacketFlags::empty(),
                stream_id: 1,
                sequence_number: 2,
                payload_length: 0,
            },
            payload: Bytes::new(),
        };
        ws.send(Message::Binary(encode_plp(&close_pkt)))
            .await
            .unwrap();

        let msg = tokio::time::timeout(Duration::from_secs(2), ws.next())
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        let close_bin = match msg {
            Message::Binary(b) => b,
            other => panic!("expected binary close, got {other:?}"),
        };
        let close_resp = decode_plp(&close_bin);
        assert_eq!(close_resp.header.packet_type, PacketType::StreamClose);
        assert_eq!(close_resp.header.stream_id, 1);

        let _ = shutdown_tx.send(());
        let _ = server.await;
    }
}
