joecalsend/src/transfer/upload.rs

use std::{collections::HashMap, net::SocketAddr, path::PathBuf, sync::Arc};

use axum::{
    Extension, Json,
    body::Bytes,
    extract::{ConnectInfo, Query, State},
    http::StatusCode,
    response::IntoResponse,
};
use julid::Julid;
use native_dialog::MessageDialogBuilder;
use serde::{Deserialize, Serialize};

use crate::{
    Client, JoecalState,
    error::{LocalSendError, Result},
    models::{device::DeviceInfo, file::FileMetadata},
    transfer::session::{Session, SessionStatus},
};

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PrepareUploadResponse {
    pub session_id: String,
    pub files: HashMap<String, String>,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PrepareUploadRequest {
    pub info: DeviceInfo,
    pub files: HashMap<String, FileMetadata>,
}

impl Client {
    pub async fn prepare_upload(
        &self,
        peer: String,
        files: HashMap<String, FileMetadata>,
    ) -> Result<PrepareUploadResponse> {
        if !self.peers.lock().await.contains_key(&peer) {
            return Err(LocalSendError::PeerNotFound);
        }

        let peer = self.peers.lock().await.get(&peer).unwrap().clone();
        println!("Peer: {peer:?}");

        let response = self
            .http_client
            .post(&format!(
                "{}://{}/api/localsend/v2/prepare-upload",
                peer.1.protocol,
                peer.0.clone()
            ))
            .json(&PrepareUploadRequest {
                info: self.device.clone(),
                files: files.clone(),
            })
            .send()
            .await?;

        println!("Response: {response:?}");

        let response: PrepareUploadResponse = response.json().await?;

        let session = Session {
            session_id: response.session_id.clone(),
            files,
            file_tokens: response.files.clone(),
            receiver: peer.1,
            sender: self.device.clone(),
            status: SessionStatus::Active,
            addr: peer.0,
        };

        self.sessions
            .lock()
            .await
            .insert(response.session_id.clone(), session);

        Ok(response)
    }

    pub async fn upload(
        &self,
        session_id: String,
        file_id: String,
        token: String,
        body: Bytes,
    ) -> Result<()> {
        let sessions = self.sessions.lock().await;
        let session = sessions.get(&session_id).unwrap();

        if session.status != SessionStatus::Active {
            return Err(LocalSendError::SessionInactive);
        }

        if session.file_tokens.get(&file_id) != Some(&token) {
            return Err(LocalSendError::InvalidToken);
        }

        let request = self
            .http_client
            .post(format!(
                "{}://{}/api/localsend/v2/upload?sessionId={session_id}&fileId={file_id}&token={token}",
                session.receiver.protocol, session.addr))
            .body(body);

        println!("Uploading file: {request:?}");
        let response = request.send().await?;

        if response.status() != 200 {
            println!("Upload failed: {response:?}");
            return Err(LocalSendError::UploadFailed);
        }

        Ok(())
    }

    pub async fn send_file(&self, peer: String, file_path: PathBuf) -> Result<()> {
        // Generate file metadata
        let file_metadata = FileMetadata::from_path(&file_path)?;

        // Prepare files map
        let mut files = HashMap::new();
        files.insert(file_metadata.id.clone(), file_metadata.clone());

        // Prepare upload
        let prepare_response = self.prepare_upload(peer, files).await?;

        // Get file token
        let token = prepare_response
            .files
            .get(&file_metadata.id)
            .ok_or(LocalSendError::InvalidToken)?;

        // Read file contents
        let file_contents = tokio::fs::read(&file_path).await?;
        let bytes = Bytes::from(file_contents);

        // Upload file
        self.upload(
            prepare_response.session_id,
            file_metadata.id,
            token.clone(),
            bytes,
        )
        .await?;

        Ok(())
    }

    pub async fn cancel_upload(&self, session_id: String) -> Result<()> {
        let sessions = self.sessions.lock().await;
        let session = sessions.get(&session_id).unwrap();

        let request = self
            .http_client
            .post(format!(
                "{}://{}/api/localsend/v2/cancel?sessionId={}",
                session.receiver.protocol, session.addr, session_id
            ))
            .send()
            .await?;

        if request.status() != 200 {
            return Err(LocalSendError::CancelFailed);
        }

        Ok(())
    }
}

pub async fn register_prepare_upload(
    State(state): State<Arc<JoecalState>>,
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    Json(req): Json<PrepareUploadRequest>,
) -> impl IntoResponse {
    println!("Received upload request from alias: {}", req.info.alias);

    let result = MessageDialogBuilder::default()
        .set_title(&req.info.alias)
        .set_text("Do you want to receive files from this device?")
        .confirm()
        .show()
        .unwrap();

    if result {
        let session_id = Julid::new().to_string();

        let file_tokens: HashMap<String, String> = req
            .files
            .keys()
            .map(|id| (id.clone(), Julid::new().to_string())) // Replace with actual token logic
            .collect();

        let session = Session {
            session_id: session_id.clone(),
            files: req.files.clone(),
            file_tokens: file_tokens.clone(),
            receiver: state.device.clone(),
            sender: req.info.clone(),
            status: SessionStatus::Active,
            addr,
        };

        state
            .sessions
            .lock()
            .await
            .insert(session_id.clone(), session);

        (
            StatusCode::OK,
            Json(PrepareUploadResponse {
                session_id,
                files: file_tokens,
            }),
        )
            .into_response()
    } else {
        StatusCode::FORBIDDEN.into_response()
    }
}

pub async fn register_upload(
    Query(params): Query<UploadParams>,
    State(state): State<Arc<JoecalState>>,
    Extension(download_dir): Extension<String>,
    body: Bytes,
) -> impl IntoResponse {
    // Extract query parameters
    let session_id = &params.session_id;
    let file_id = &params.file_id;
    let token = &params.token;

    // Get session and validate
    let mut sessions_lock = state.sessions.lock().await;
    let session = match sessions_lock.get_mut(session_id) {
        Some(session) => session,
        None => return StatusCode::BAD_REQUEST.into_response(),
    };

    if session.status != SessionStatus::Active {
        return StatusCode::BAD_REQUEST.into_response();
    }

    // Validate token
    if session.file_tokens.get(file_id) != Some(&token.to_string()) {
        return StatusCode::FORBIDDEN.into_response();
    }

    // Get file metadata
    let file_metadata = match session.files.get(file_id) {
        Some(metadata) => metadata,
        None => {
            return (
                StatusCode::INTERNAL_SERVER_ERROR,
                "File not found".to_string(),
            )
                .into_response();
        }
    };

    // Create directory if it doesn't exist
    if let Err(e) = tokio::fs::create_dir_all(&*download_dir).await {
        return (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to create directory: {e}"),
        )
            .into_response();
    }

    // Create file path
    let file_path = format!("{}/{}", download_dir, file_metadata.file_name);

    // Write file
    if let Err(e) = tokio::fs::write(&file_path, body).await {
        return (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Failed to write file: {e}"),
        )
            .into_response();
    }

    StatusCode::OK.into_response()
}

// Query parameters struct
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UploadParams {
    session_id: String,
    file_id: String,
    token: String,
}

pub async fn register_cancel(
    Query(params): Query<CancelParams>,
    State(state): State<JoecalState>,
) -> impl IntoResponse {
    let mut sessions_lock = state.sessions.lock().await;
    let session = match sessions_lock.get_mut(&params.session_id) {
        Some(session) => session,
        None => return StatusCode::BAD_REQUEST.into_response(),
    };
    session.status = SessionStatus::Cancelled;
    StatusCode::OK.into_response()
}

// Cancel parameters struct
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CancelParams {
    session_id: String,
}