use std::collections::HashMap;
use std::thread;
use std::time::Duration;

use crate::api::ApiClient;
use crate::error::BioLibError;

const MIN_CHUNK_SIZE: usize = 10_000_000; // 10 MB
const MAX_CHUNK_COUNT: usize = 9_000;
const MAX_CHUNK_RETRIES: u32 = 20;

pub fn upload_module_input(
    api_client: &mut ApiClient,
    job_uuid: &str,
    job_auth_token: &str,
    module_input: &[u8],
) -> crate::Result<()> {
    let mut headers = HashMap::new();
    headers.insert("job-auth-token".to_string(), job_auth_token.to_string());

    api_client.start_multipart_upload(
        &format!("/jobs/{job_uuid}/storage/input/start_upload/"),
        Some(&headers),
    )?;

    let chunk_size = calculate_chunk_size(module_input.len());
    let mut parts = Vec::new();
    let mut offset = 0usize;
    let mut part_number = 1u32;

    while offset < module_input.len() {
        let end = std::cmp::min(offset + chunk_size, module_input.len());
        let chunk = &module_input[offset..end];

        let (etag, _chunk_len) = upload_chunk(
            api_client,
            &format!("/jobs/{job_uuid}/storage/input/presigned_upload_url/"),
            part_number,
            chunk,
            &headers,
        )?;

        parts.push(serde_json::json!({
            "ETag": etag,
            "PartNumber": part_number,
        }));

        offset = end;
        part_number += 1;
    }

    api_client.complete_multipart_upload(
        &format!("/jobs/{job_uuid}/storage/input/complete_upload/"),
        &parts,
        Some(&headers),
    )?;

    Ok(())
}

fn upload_chunk(
    api_client: &mut ApiClient,
    presigned_url_path: &str,
    part_number: u32,
    chunk: &[u8],
    headers: &HashMap<String, String>,
) -> crate::Result<(String, usize)> {
    let mut last_error: Option<BioLibError> = None;

    for attempt in 0..MAX_CHUNK_RETRIES {
        if attempt > 0 {
            // Quadratic backoff matching Python: sleep(attempt² + 2)
            let delay_secs = (attempt as u64) * (attempt as u64) + 2;
            crate::logging::warning(&format!(
                "Retrying upload of part {part_number} (attempt {attempt}, waiting {delay_secs}s)..."
            ));
            thread::sleep(Duration::from_secs(delay_secs));
        }

        let presigned_url = match api_client.get_presigned_upload_url(
            presigned_url_path,
            part_number,
            Some(headers),
        ) {
            Ok(url) => url,
            Err(err) => {
                crate::logging::warning(&format!(
                    "Error getting URL for part {part_number}: {err}"
                ));
                last_error = Some(err);
                continue;
            }
        };

        match api_client.upload_to_presigned_url(&presigned_url, chunk) {
            Ok(etag) => return Ok((etag, chunk.len())),
            Err(err) => {
                crate::logging::warning(&format!("Error uploading part {part_number}: {err}"));
                last_error = Some(err);
            }
        }
    }

    Err(last_error.unwrap_or_else(|| {
        BioLibError::General(format!("Max retries hit when uploading part {part_number}"))
    }))
}

fn calculate_chunk_size(total_size: usize) -> usize {
    std::cmp::max(MIN_CHUNK_SIZE, total_size / MAX_CHUNK_COUNT)
}
