use super::*;

#[derive(Debug, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum PadServerMessage {
    Authenticated {
        title: String,
        content: String,
        owner_map: String,
        revision_id: i64,
        access_level: String,
        catchup_operations: Vec<AppliedOperation>,
        acknowledged_update_ids: Vec<u64>,
        resync_required: bool,
    },
    Document {
        base_revision_id: i64,
        revision_id: i64,
        updated_at: String,
        author: Option<String>,
        client_id: String,
        update_id: u64,
        operation: TextOperation,
        owner_replacements: Vec<OwnerReplacement>,
    },
    Resync {
        content: String,
        revision_id: i64,
        owner_map: String,
        acknowledged_update_ids: Vec<u64>,
    },
    Presence {
        users: Vec<PresenceUser>,
    },
    Chat {
        sender: String,
        text: String,
    },
    Pong {
        nonce: u64,
    },
    Diagnostics {
        diagnostics: ConnectionDiagnostics,
    },
    PasswordRequired,
    PasswordChanged,
    Error {
        message: String,
    },
}
pub async fn upgrade_pad(
    ws: WebSocketUpgrade,
    headers: HeaderMap,
    Path(slug): Path<String>,
    State(state): State<SharedState>,
) -> Response {
    if !crate::security::websocket_origin_allowed(&headers) {
        warn!(%slug, "pad websocket rejected: invalid origin");
        return (StatusCode::FORBIDDEN, "Invalid WebSocket origin").into_response();
    }
    let account_token = crate::security::session_token(&headers).map(str::to_owned);
    let share_session_token =
        crate::security::share_session_token(&headers, "pad", &slug).map(str::to_owned);
    let resource_token = crate::security::resource_token(&headers, "pad", &slug).map(str::to_owned);
    let client_key = crate::security::client_key(&headers);
    let client_context = RequestClientContext::from_headers(&headers, &client_key);
    ws.on_upgrade(move |socket| {
        handle_pad_socket(
            socket,
            state,
            slug,
            account_token,
            share_session_token,
            resource_token,
            client_key,
            client_context,
        )
    })
}
async fn handle_pad_socket(
    mut socket: WebSocket,
    state: SharedState,
    slug: String,
    cookie_session_token: Option<String>,
    cookie_share_session_token: Option<String>,
    cookie_password_token: Option<String>,
    client_key: String,
    client_context: RequestClientContext,
) {
    info!(%slug, "pad websocket connected");
    let Some(pad) = db::find_pad(&state.db, &slug).await.ok().flatten() else {
        warn!(%slug, "pad websocket rejected: pad not found");
        let _ = send_pad(
            &mut socket,
            &PadServerMessage::Error {
                message: "Note not found".into(),
            },
        )
        .await;
        return;
    };
    let (
        password,
        access_token,
        nickname,
        guest_id,
        color,
        client_diagnostics,
        collaboration_client_id,
        known_revision_id,
    ) = match socket.recv().await {
        Some(Ok(Message::Text(text))) => match serde_json::from_str::<ClientMessage>(&text) {
            Ok(ClientMessage::Authenticate {
                password,
                access_token,
                nickname,
                guest_id,
                color,
                diagnostics,
                client_id,
                known_revision_id,
            }) => (
                password,
                access_token,
                clean_nickname(nickname),
                clean_guest_id(guest_id),
                clean_color(color),
                diagnostics,
                clean_collaboration_client_id(client_id)
                    .unwrap_or_else(|| format!("legacy_{}", db::random_suffix(24))),
                known_revision_id.filter(|revision_id| *revision_id >= 0),
            ),
            _ => {
                let _ = send_pad(
                    &mut socket,
                    &PadServerMessage::Error {
                        message: "Wymagane uwierzytelnienie".into(),
                    },
                )
                .await;
                return;
            }
        },
        _ => return,
    };
    let session_token = cookie_session_token;
    let explicit_access_token = access_token
        .as_deref()
        .map(str::trim)
        .filter(|value| !value.is_empty() && *value != "cookie")
        .map(str::to_owned);
    let external_tokens = [
        explicit_access_token.as_deref(),
        cookie_share_session_token.as_deref(),
        cookie_password_token.as_deref(),
    ];
    let nickname = match auth::authorize_nickname(&state, nickname, session_token.clone()).await {
        Ok(value) => value,
        Err(message) => {
            let _ = send_pad(&mut socket, &PadServerMessage::Error { message }).await;
            return;
        }
    };
    let presence_identity = match session_token.as_deref() {
        Some(token) => auth::user_from_token(&state, token)
            .await
            .ok()
            .flatten()
            .map(|user| format!("user:{}", user.id)),
        None => guest_id.as_ref().and_then(|id| {
            nickname
                .as_ref()
                .map(|name| format!("guest:{id}:{}", name.to_lowercase()))
        }),
    };
    let permission = resource_permission_from_tokens(
        &state,
        "pad",
        &slug,
        &external_tokens,
        session_token.as_deref(),
    )
    .await;
    let anonymous_token_ok =
        password_access_from_tokens(&state, "pad", &slug, &external_tokens).await;
    let password_limit_key = format!("resource-password:{client_key}:pad:{slug}");
    let password_attempted = password
        .as_deref()
        .map(str::trim)
        .is_some_and(|value| !value.is_empty());
    if pad.password_hash.is_some()
        && permission.as_deref() != Some("rw")
        && !anonymous_token_ok
        && password_attempted
    {
        let window = std::time::Duration::from_secs(15 * 60);
        if let Err(seconds) = state
            .check_rate_limit(format!("resource-password-client:{client_key}"), 50, window)
            .await
        {
            let _ = send_pad(
                &mut socket,
                &PadServerMessage::Error {
                    message: format!("Too many password attempts. Try again in {seconds} seconds."),
                },
            )
            .await;
            return;
        }
        if let Err(seconds) = state
            .check_rate_limit(password_limit_key.clone(), 10, window)
            .await
        {
            let _ = send_pad(
                &mut socket,
                &PadServerMessage::Error {
                    message: format!("Too many password attempts. Try again in {seconds} seconds."),
                },
            )
            .await;
            return;
        }
    }
    let password_ok = db::verify_pad_password(&pad, password.as_deref());
    if password_ok {
        state.clear_rate_limit(&password_limit_key).await;
    }
    if pad.is_private != 0 && permission.is_none() && !anonymous_token_ok {
        let _ = send_pad(
            &mut socket,
            &PadServerMessage::Error {
                message: "Note not found".into(),
            },
        )
        .await;
        return;
    }
    if pad.password_hash.is_some() && !password_ok && permission.is_none() && !anonymous_token_ok {
        warn!(pad_id = pad.id, "pad websocket rejected: invalid password");
        let _ = send_pad(
            &mut socket,
            &PadServerMessage::Error {
                message: "Invalid password".into(),
            },
        )
        .await;
        return;
    }
    let (_, write_allowed) = current_resource_access(
        &state,
        "pad",
        &slug,
        &external_tokens,
        session_token.as_deref(),
        password_ok,
    )
    .await;
    info!(pad_id = pad.id, nickname = ?nickname, "pad websocket authenticated");
    let room_key = AppState::pad_room_key(&slug);
    let collaboration_snapshot = match db::pad_collaboration_snapshot(&state.db, pad.id).await {
        Ok(snapshot) => snapshot,
        Err(error) => {
            warn!(%error, pad_id = pad.id, "failed to load collaborative document");
            let _ = send_pad(
                &mut socket,
                &PadServerMessage::Error {
                    message: "Failed to load the document".into(),
                },
            )
            .await;
            return;
        }
    };
    let collaborative_document = state
        .collaborative_document(
            &room_key,
            collaboration_snapshot.content,
            collaboration_snapshot.owner_map,
            collaboration_snapshot.revision_id,
        )
        .await;
    // Subscribe before taking the authentication snapshot. Updates committed after
    // the snapshot are then queued for this connection instead of falling into a gap.
    let channel = state.pad_channel(&slug).await;
    let mut updates = channel.subscribe();
    let persisted_acknowledged_update_id =
        match db::latest_pad_collaboration_update_id(&state.db, pad.id, &collaboration_client_id)
            .await
        {
            Ok(update_id) => update_id,
            Err(error) => {
                warn!(%error, pad_id = pad.id, "failed to load collaborative acknowledgement");
                let _ = send_pad(
                    &mut socket,
                    &PadServerMessage::Error {
                        message: "Failed to load the document".into(),
                    },
                )
                .await;
                return;
            }
        };
    let (
        authenticated_content,
        authenticated_owner_map,
        authenticated_revision_id,
        catchup_operations,
        acknowledged_update_ids,
        resync_required,
    ) = {
        let mut document = collaborative_document.lock().await;
        if let Some(update_id) = persisted_acknowledged_update_id {
            document.acknowledge(&collaboration_client_id, update_id);
        }
        let (catchup_operations, resync_required) = match known_revision_id {
            Some(revision_id) => match document.operations_after(revision_id) {
                Some(operations) => (operations, false),
                None => (Vec::new(), revision_id != document.revision_id),
            },
            None => (Vec::new(), false),
        };
        (
            document.content.clone(),
            document.owner_map.clone(),
            document.revision_id,
            catchup_operations,
            document.acknowledged_updates(&collaboration_client_id),
            resync_required,
        )
    };
    if send_pad(
        &mut socket,
        &PadServerMessage::Authenticated {
            title: pad.title.clone(),
            content: authenticated_content,
            owner_map: authenticated_owner_map,
            revision_id: authenticated_revision_id,
            catchup_operations,
            acknowledged_update_ids,
            resync_required,
            access_level: if write_allowed {
                "full".into()
            } else {
                "read_only".into()
            },
        },
    )
    .await
    .is_err()
    {
        return;
    }
    let display_name = nickname.clone().unwrap_or_else(|| "Guest".into());
    let (connection_id, users) = state
        .join_room(&room_key, display_name.clone(), color, presence_identity)
        .await;
    let _ = channel.send(RoomEvent::Presence(users));
    let mut last_chat = Instant::now() - Duration::from_secs(1);
    let mut access_refresh = tokio::time::interval(Duration::from_secs(10));
    access_refresh.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
    access_refresh.tick().await;
    let (mut sender, mut receiver) = socket.split();
    if send_pad_split(
        &mut sender,
        &PadServerMessage::Diagnostics {
            diagnostics: connection_diagnostics(&client_context, client_diagnostics),
        },
    )
    .await
    .is_err()
    {
        let users = state.leave_room(&room_key, connection_id).await;
        let _ = channel.send(RoomEvent::Presence(users));
        return;
    }
    loop {
        tokio::select! {
            incoming=receiver.next()=>match incoming{
                Some(Ok(Message::Text(text)))=>match serde_json::from_str::<ClientMessage>(&text){
                    Ok(ClientMessage::Update{base_revision_id,update_id,operation,owner_replacements})=>{
                        let (read_allowed, current_write_allowed) = current_resource_access(
                            &state,
                            "pad",
                            &slug,
                            &external_tokens,
                            session_token.as_deref(),
                            password_ok,
                        ).await;
                        if !read_allowed { let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Access expired or revoked".into()}).await;break; }
                        if !current_write_allowed{let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Read-only access".into()}).await;continue;}
                        if update_id == 0
                            || update_id > i64::MAX as u64
                            || !valid_owner_replacements(&owner_replacements)
                        {
                            let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Invalid collaborative update".into()}).await;
                            continue;
                        }
                        let mut document = collaborative_document.lock().await;
                        if document.has_applied_update(&collaboration_client_id, update_id) {
                            let snapshot = (
                                document.content.clone(),
                                document.revision_id,
                                document.owner_map.clone(),
                                document.acknowledged_updates(&collaboration_client_id),
                            );
                            drop(document);
                            let _ = send_pad_split(&mut sender, &PadServerMessage::Resync {
                                content: snapshot.0,
                                revision_id: snapshot.1,
                                owner_map: snapshot.2,
                                acknowledged_update_ids: snapshot.3,
                            }).await;
                            continue;
                        }
                        let transformed = match document.transform_from(
                            base_revision_id,
                            &operation,
                            &collaboration_client_id,
                            update_id,
                        ) {
                            Ok(operation) => operation,
                            Err(collab::OperationError::RevisionUnavailable) => {
                                let snapshot = (
                                    document.content.clone(),
                                    document.revision_id,
                                    document.owner_map.clone(),
                                    document.acknowledged_updates(&collaboration_client_id),
                                );
                                drop(document);
                                let _=send_pad_split(&mut sender,&PadServerMessage::Resync{
                                    content:snapshot.0,
                                    revision_id:snapshot.1,
                                    owner_map:snapshot.2,
                                    acknowledged_update_ids:snapshot.3,
                                }).await;
                                continue;
                            }
                            Err(error) => {
                                drop(document);
                                warn!(%error, pad_id = pad.id, "invalid collaborative operation");
                                let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Invalid collaborative update".into()}).await;
                                continue;
                            }
                        };
                        let applied_base_revision_id = document.revision_id;
                        let (content,owner_map)=match collab::apply_operation_to_document(
                            &document.content,
                            &document.owner_map,
                            &transformed,
                            &owner_replacements,
                        ) {
                            Ok(document) => document,
                            Err(error) => {
                                drop(document);
                                warn!(%error, pad_id = pad.id, "failed to apply collaborative operation");
                                let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Invalid collaborative update".into()}).await;
                                continue;
                            }
                        };
                        if content.len()>2_000_000 {
                            drop(document);
                            let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"The document is too large".into()}).await;
                            continue;
                        }
                        match db::save_collaborative_pad_revision(
                            &state.db,
                            pad.id,
                            &content,
                            nickname.as_deref(),
                            &owner_map,
                            &collaboration_client_id,
                            update_id as i64,
                        ).await{
                            Ok((revision_id,updated_at))=>{
                                document.content=content.clone();
                                document.owner_map=owner_map.clone();
                                document.revision_id=revision_id;
                                document.record(AppliedOperation{
                                    base_revision_id:applied_base_revision_id,
                                    revision_id,
                                    client_id:collaboration_client_id.clone(),
                                    update_id,
                                    operation:transformed.clone(),
                                    owner_replacements:owner_replacements.clone(),
                                });
                                let _=channel.send(RoomEvent::Document(NoteUpdate{
                                    base_revision_id:applied_base_revision_id,
                                    revision_id,
                                    updated_at,
                                    author:nickname.clone(),
                                    client_id:collaboration_client_id.clone(),
                                    update_id,
                                    operation:transformed,
                                    owner_replacements,
                                }));
                                drop(document);
                            }
                            Err(error)=>{
                                drop(document);
                                warn!(%error, pad_id = pad.id, "failed to save revision");
                                let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Failed to save the document".into()}).await;
                            }
                        }
                    }
                    Ok(ClientMessage::Ping{nonce})=>{ let _=send_pad_split(&mut sender,&PadServerMessage::Pong{nonce}).await; },
                    Ok(ClientMessage::Chat{text})=>{
                        let text=clean_chat(text);
                        if !text.is_empty() && last_chat.elapsed() >= Duration::from_millis(500) { last_chat=Instant::now(); let _=channel.send(RoomEvent::Chat{sender:display_name.clone(),text}); }
                    }
                    Ok(ClientMessage::SetColor{color})=>{ let users=state.update_room_color(&room_key,connection_id,clean_color(color)).await; let _=channel.send(RoomEvent::Presence(users)); },
                    Ok(ClientMessage::Authenticate{..})=>{},
                    Err(error)=>warn!(%error,"invalid pad websocket message"),
                },
                Some(Ok(Message::Close(_)))|None=>break,
                Some(Ok(_))=>{},
                Some(Err(error))=>{debug!(%error,"pad websocket receive error");break;}
            },
            _=access_refresh.tick()=>{
                let (read_allowed, _) = current_resource_access(
                    &state,
                    "pad",
                    &slug,
                    &external_tokens,
                    session_token.as_deref(),
                    password_ok,
                ).await;
                if !read_allowed {
                    let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Access expired or revoked".into()}).await;
                    break;
                }
            },
            update=updates.recv()=>{
                if let Ok(RoomEvent::PasswordRequired { except_client_id }) = &update {
                    if password_event_excludes_connection(
                        except_client_id.as_deref(),
                        &collaboration_client_id,
                    ) {
                        let _=send_pad_split(&mut sender,&PadServerMessage::PasswordChanged).await;
                        break;
                    }
                    let (read_allowed, _) = current_resource_access(
                        &state,
                        "pad",
                        &slug,
                        &external_tokens,
                        session_token.as_deref(),
                        false,
                    ).await;
                    if !read_allowed {
                        let _=send_pad_split(&mut sender,&PadServerMessage::PasswordRequired).await;
                        break;
                    }
                    continue;
                }
                let (read_allowed, _) = current_resource_access(
                    &state,
                    "pad",
                    &slug,
                    &external_tokens,
                    session_token.as_deref(),
                    password_ok,
                ).await;
                if !read_allowed {
                    let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Access expired or revoked".into()}).await;
                    break;
                }
                match update {
                    Ok(RoomEvent::Document(u))=>if send_pad_split(&mut sender,&PadServerMessage::Document{base_revision_id:u.base_revision_id,revision_id:u.revision_id,updated_at:u.updated_at,author:u.author,client_id:u.client_id,update_id:u.update_id,operation:u.operation,owner_replacements:u.owner_replacements}).await.is_err(){break;},
                    Ok(RoomEvent::Presence(users))=>if send_pad_split(&mut sender,&PadServerMessage::Presence{users}).await.is_err(){break;},
                    Ok(RoomEvent::Chat{sender:chat_sender,text})=>if send_pad_split(&mut sender,&PadServerMessage::Chat{sender:chat_sender,text}).await.is_err(){break;},
                    Ok(RoomEvent::PasswordRequired { .. })=>{},
                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_))=>{
                        let document=collaborative_document.lock().await;
                        let snapshot=(document.content.clone(),document.revision_id,document.owner_map.clone(),document.acknowledged_updates(&collaboration_client_id));
                        drop(document);
                        if send_pad_split(&mut sender,&PadServerMessage::Resync{content:snapshot.0,revision_id:snapshot.1,owner_map:snapshot.2,acknowledged_update_ids:snapshot.3}).await.is_err(){break;}
                    },
                    Err(tokio::sync::broadcast::error::RecvError::Closed)=>break,
                }
            }
        }
    }
    let users = state.leave_room(&room_key, connection_id).await;
    let _ = channel.send(RoomEvent::Presence(users));
    info!(pad_id = pad.id, "pad websocket disconnected");
}
async fn send_pad(socket: &mut WebSocket, message: &PadServerMessage) -> Result<(), axum::Error> {
    socket
        .send(Message::Text(
            serde_json::to_string(message).unwrap().into(),
        ))
        .await
}
async fn send_pad_split(
    sender: &mut futures_util::stream::SplitSink<WebSocket, Message>,
    message: &PadServerMessage,
) -> Result<(), axum::Error> {
    sender
        .send(Message::Text(
            serde_json::to_string(message).unwrap().into(),
        ))
        .await
}