use axum::body::Body;
use axum::extract::{Path, State};
use axum::http::{StatusCode, header};
use axum::response::Response;
use futures::TryStreamExt;
use jacquard::identity::resolver::IdentityResolver;
use jacquard_common::types::string::{Cid, Did};
use jacquard_common::xrpc::XrpcExt;
use polymodel_api::com_atproto::sync::get_blob::GetBlob;
use polymodel_api::space_polymodel::library::File;
use sha2::{Digest, Sha256};

use super::error::{AppResult, db, internal, invalid_request, not_found};
use super::state::AppState;

const MAX_DOWNLOAD_BYTES: i64 = 100 * 1024 * 1024;

pub(super) async fn download_part(
    State(state): State<AppState>,
    Path((did, rkey)): Path<(String, String)>,
) -> AppResult<Response> {
    let did: Did = Did::new_owned(did).map_err(|e| invalid_request(format!("invalid DID: {e}")))?;
    let rkey = rkey.trim();
    if rkey.is_empty() || rkey.contains('/') {
        return Err(invalid_request("invalid part rkey"));
    }

    #[derive(sqlx::FromRow)]
    struct PartDownloadRow {
        name: String,
        file_json: String,
    }

    let row = db(sqlx::query_as!(
        PartDownloadRow,
        r#"SELECT name AS "name!", file_json AS "file_json!" FROM parts WHERE did = ? AND rkey = ?"#,
        did.as_ref(),
        rkey,
    )
    .fetch_optional(&state.pool)
    .await)?
    .ok_or_else(not_found)?;

    let file: File = serde_json::from_str(&row.file_json)
        .map_err(|e| internal(format!("part file manifest parse failed: {e}")))?;
    if file.chunks.is_empty() {
        return Err(not_found());
    }
    validate_file_manifest(&file)?;

    let did_doc_response = state
        .resolver
        .resolve_did_doc(&did)
        .await
        .map_err(|e| internal(format!("DID document resolution failed: {e}")))?;
    let did_doc = did_doc_response
        .parse()
        .map_err(|e| internal(format!("DID document parse failed: {e}")))?;
    let pds_endpoint = did_doc
        .pds_endpoint()
        .ok_or_else(|| internal("DID document does not advertise a PDS endpoint"))?
        .to_owned();

    let mut chunks = file.chunks;
    chunks.sort_by_key(|chunk| chunk.offset);
    let mut bytes = Vec::with_capacity(file.size as usize);
    for chunk in chunks {
        let cid: Cid = Cid::new_owned(chunk.blob.blob().r#ref.as_str().as_bytes())
            .map_err(|e| internal(format!("part blob CID is invalid: {e}")))?;
        let request = GetBlob {
            did: did.clone(),
            cid,
        };
        let response = state
            .resolver
            .xrpc(pds_endpoint.borrow())
            .download(&request)
            .await
            .map_err(|e| internal(format!("blob fetch from {pds_endpoint} failed: {e}")))?;
        if !response.status().is_success() {
            return Err(internal(format!(
                "blob fetch from {pds_endpoint} failed with HTTP {}",
                response.status()
            )));
        }
        let chunk_start = bytes.len();
        let (_, body) = response.into_parts();
        let mut stream = body.into_inner();
        while let Some(chunk_bytes) = stream
            .try_next()
            .await
            .map_err(|e| internal(format!("blob stream from {pds_endpoint} failed: {e}")))?
        {
            bytes.extend_from_slice(&chunk_bytes);
        }
        if bytes.len() - chunk_start != chunk.size as usize {
            return Err(internal("blob chunk length did not match manifest"));
        }
    }
    if bytes.len() != file.size as usize {
        return Err(internal("downloaded file length did not match manifest"));
    }
    if let Some(digest) = file.digest.as_ref() {
        let actual = Sha256::digest(&bytes);
        if actual[..] != *digest.as_ref() {
            return Err(internal("downloaded file digest did not match manifest"));
        }
    }

    let filename = download_filename(&row.name);
    let content_type = file.mime_type.to_string();
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, content_type)
        .header(
            header::CONTENT_DISPOSITION,
            format!("attachment; filename=\"{filename}\""),
        )
        .body(Body::from(bytes))
        .map_err(|e| internal(format!("download response build failed: {e}")))
}

fn validate_file_manifest(file: &File) -> AppResult<()> {
    if file.size < 0 || file.size > MAX_DOWNLOAD_BYTES {
        return Err(invalid_request(
            "part file size is outside the download limit",
        ));
    }
    if file.chunks.len() > 4096 {
        return Err(invalid_request("part file has too many chunks"));
    }
    let mut chunks = file.chunks.iter().collect::<Vec<_>>();
    chunks.sort_by_key(|chunk| chunk.offset);
    let mut expected_offset = 0_i64;
    for chunk in chunks {
        if chunk.offset != expected_offset || chunk.size < 0 {
            return Err(invalid_request("part file chunks are not contiguous"));
        }
        expected_offset = expected_offset
            .checked_add(chunk.size)
            .ok_or_else(|| invalid_request("part file chunk sizes overflow"))?;
        if expected_offset > MAX_DOWNLOAD_BYTES {
            return Err(invalid_request(
                "part file size is outside the download limit",
            ));
        }
    }
    if expected_offset != file.size {
        return Err(invalid_request("part file chunks do not match file size"));
    }
    Ok(())
}

fn download_filename(name: &str) -> String {
    let cleaned = name
        .chars()
        .map(|ch| match ch {
            'a'..='z' | 'A'..='Z' | '0'..='9' | '.' | '-' | '_' => ch,
            _ => '_',
        })
        .collect::<String>();
    if cleaned.trim_matches('_').is_empty() {
        "polymodel-part.bin".to_string()
    } else {
        cleaned
    }
}

#[cfg(test)]
mod tests {
    use super::{download_filename, validate_file_manifest};
    use jacquard_common::deps::smol_str::SmolStr;
    use jacquard_common::types::blob::{Blob, BlobRef, MimeType};
    use jacquard_common::types::cid::CidLink;
    use polymodel_api::space_polymodel::library::{Chunk, File};

    fn blob_ref() -> BlobRef {
        BlobRef::from(Blob {
            r#ref: CidLink::new_owned(
                b"bafkreihdwdcefgh4dqkjv67uzcmw7ojee6xedzdetojuzjevtenxquvyku",
            )
            .unwrap(),
            mime_type: MimeType::new_owned("model/stl"),
            size: 4,
        })
    }

    fn file(chunks: Vec<Chunk>, size: i64) -> File {
        File {
            chunks,
            digest: None,
            mime_type: SmolStr::from("model/stl"),
            size,
            extra_data: None,
        }
    }

    #[test]
    fn download_filename_removes_header_unsafe_characters() {
        assert_eq!(download_filename("gear bracket.stl"), "gear_bracket.stl");
        assert_eq!(download_filename("../../x\r\n"), ".._.._x__");
        assert_eq!(download_filename("!!!"), "polymodel-part.bin");
    }

    #[test]
    fn file_manifest_validation_requires_contiguous_bounded_chunks() {
        let chunk = Chunk::builder().blob(blob_ref()).offset(0).size(4).build();
        assert!(validate_file_manifest(&file(vec![chunk.clone()], 4)).is_ok());

        let gapped = Chunk::builder().blob(blob_ref()).offset(2).size(4).build();
        assert!(validate_file_manifest(&file(vec![gapped], 6)).is_err());
        assert!(validate_file_manifest(&file(vec![chunk], 101 * 1024 * 1024)).is_err());
    }
}
