use anyhow::{Context, Result};
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
use serde_json::{json, Value};
use sha2::{Digest as _, Sha256};
use std::io::Write;
use std::path::Path;
use std::time::{Duration, Instant};

use crate::client::{
    authenticated_or_local_trusted_request, list_task_events_if_available_with_session,
    JsonLineSession,
};
use crate::config::StoredCliSession;
use crate::errors::cli_error_summary_for_category;
use crate::process_events::{
    artifact_download_grant_disclosures, artifact_download_session_summary,
    artifact_export_plan_summary, artifact_response_machine_error, artifact_summaries,
};
use crate::{ArtifactDownloadArgs, ArtifactExportArgs, ArtifactListArgs};

pub(crate) const DEFAULT_ARTIFACT_EXPORT_MAX_BYTES: u64 = 64 * 1024 * 1024;
const ARTIFACT_DOWNLOAD_CONTROL_CHUNK_BYTES: u64 = 256 * 1024;

#[cfg(test)]
pub(crate) fn artifact_list_report(args: ArtifactListArgs) -> Result<Value> {
    artifact_list_report_with_session(args, None)
}

pub(crate) fn artifact_list_report_with_session(
    args: ArtifactListArgs,
    stored_session: Option<&StoredCliSession>,
) -> Result<Value> {
    let events = list_task_events_if_available_with_session(
        args.scope.coordinator.as_deref(),
        &args.scope,
        args.process.clone(),
        stored_session,
    )?;
    let artifacts = artifact_summaries(events.as_ref());
    Ok(json!({
        "command": "artifact list",
        "process": args.process,
        "source": "task_events",
        "artifacts": artifacts,
        "default_durable_store_assumed": false,
        "events": events,
    }))
}

#[cfg(test)]
pub(crate) fn artifact_download_report(args: ArtifactDownloadArgs) -> Result<Value> {
    artifact_download_report_with_session(args, None)
}

pub(crate) fn artifact_download_report_with_session(
    args: ArtifactDownloadArgs,
    stored_session: Option<&StoredCliSession>,
) -> Result<Value> {
    if let Some(coordinator) = &args.scope.coordinator {
        let mut session = JsonLineSession::connect(coordinator)?;
        let response = session.request(authenticated_or_local_trusted_request(
            coordinator,
            stored_session,
            json!({
                "type": "create_artifact_download_link",
                "artifact": args.artifact,
                "max_bytes": args.max_bytes,
            }),
            json!({
                "type": "create_artifact_download_link",
                "tenant": args.scope.tenant,
                "project": args.scope.project,
                "actor_user": args.scope.user,
                "artifact": args.artifact,
                "max_bytes": args.max_bytes,
            }),
        )?)?;
        let download_session = artifact_download_session_summary(&response);
        let grant_disclosures = artifact_download_grant_disclosures(&response);
        let local_download = match &args.to {
            Some(path)
                if response.get("type").and_then(Value::as_str)
                    == Some("artifact_download_link") =>
            {
                download_link_to_path(
                    &mut session,
                    coordinator,
                    &args.scope,
                    &args.artifact,
                    args.max_bytes,
                    path,
                    &response,
                    stored_session,
                )?
            }
            Some(path) => json!({
                "status": "download_link_failed",
                "local_path": path,
                "local_bytes_written_by_cli": false,
                "machine_error": artifact_response_machine_error(
                    &response,
                    "coordinator rejected artifact download",
                    "connectivity"
                ),
            }),
            None => Value::Null,
        };
        return Ok(json!({
            "command": "artifact download",
            "coordinator": coordinator,
            "artifact": args.artifact,
            "max_bytes": args.max_bytes,
            "to": args.to,
            "download_session": download_session,
            "local_download": local_download,
            "grant_disclosures": grant_disclosures,
            "response": response,
            "coordinator_session_requests": session.requests(),
        }));
    }
    Ok(json!({
        "command": "artifact download",
        "status": "requires_coordinator",
        "artifact": args.artifact,
        "max_bytes": args.max_bytes,
        "to": args.to,
        "download_session": {
            "status": "requires_coordinator",
            "link_issued": false,
            "explicit_user_action_required": true,
            "machine_error": cli_error_summary_for_category(
                "connectivity",
                "artifact download requires a coordinator"
            ),
        },
        "grant_disclosures": [],
    }))
}

#[cfg(test)]
pub(crate) fn artifact_export_report(args: ArtifactExportArgs) -> Result<Value> {
    artifact_export_report_with_session(args, None)
}

pub(crate) fn artifact_export_report_with_session(
    args: ArtifactExportArgs,
    stored_session: Option<&StoredCliSession>,
) -> Result<Value> {
    if let Some(coordinator) = &args.scope.coordinator {
        let mut session = JsonLineSession::connect(coordinator)?;
        let response = session.request(authenticated_or_local_trusted_request(
            coordinator,
            stored_session,
            json!({
                "type": "export_artifact_to_node",
                "artifact": args.artifact,
                "receiver_node": args.receiver_node,
                "direct_connectivity": true,
                "failure_reason": "",
            }),
            json!({
                "type": "export_artifact_to_node",
                "tenant": args.scope.tenant,
                "project": args.scope.project,
                "actor_user": args.scope.user,
                "artifact": args.artifact,
                "receiver_node": args.receiver_node,
                "direct_connectivity": true,
                "failure_reason": "",
            }),
        )?)?;
        let mut export_plan = artifact_export_plan_summary(&response, &args.to);
        let local_export = if response.get("type").and_then(Value::as_str)
            == Some("artifact_export_plan")
        {
            artifact_export_local_write_followup(&mut session, coordinator, &args, stored_session)?
        } else {
            json!({
                "status": "skipped",
                "explicit_user_action": true,
                "local_path": &args.to,
                "local_bytes_written_by_cli": false,
                "machine_error": artifact_response_machine_error(
                    &response,
                    "coordinator rejected artifact export",
                    "connectivity"
                ),
            })
        };
        apply_local_export_summary(&mut export_plan, &local_export);
        let grant_disclosures = local_export
            .get("grant_disclosures")
            .cloned()
            .unwrap_or_else(|| json!([]));
        return Ok(json!({
            "command": "artifact export",
            "coordinator": coordinator,
            "artifact": args.artifact,
            "to": args.to,
            "receiver_node": args.receiver_node,
            "export_plan": export_plan,
            "local_export": local_export,
            "grant_disclosures": grant_disclosures,
            "response": response,
            "coordinator_session_requests": session.requests(),
        }));
    }
    Ok(json!({
        "command": "artifact export",
        "status": "requires_coordinator",
        "artifact": args.artifact,
        "to": args.to,
        "receiver_node": args.receiver_node,
        "export_plan": {
            "status": "requires_coordinator",
            "explicit_user_action": true,
            "local_bytes_written_by_cli": false,
            "default_durable_store_assumed": false,
            "machine_error": cli_error_summary_for_category(
                "connectivity",
                "artifact export requires a coordinator"
            ),
        },
        "grant_disclosures": [],
    }))
}

fn artifact_export_local_write_followup(
    session: &mut JsonLineSession,
    coordinator: &str,
    args: &ArtifactExportArgs,
    stored_session: Option<&StoredCliSession>,
) -> Result<Value> {
    let link_response = session.request(authenticated_or_local_trusted_request(
        coordinator,
        stored_session,
        json!({
            "type": "create_artifact_download_link",
            "artifact": args.artifact,
            "max_bytes": DEFAULT_ARTIFACT_EXPORT_MAX_BYTES,
        }),
        json!({
            "type": "create_artifact_download_link",
            "tenant": args.scope.tenant,
            "project": args.scope.project,
            "actor_user": args.scope.user,
            "artifact": args.artifact,
            "max_bytes": DEFAULT_ARTIFACT_EXPORT_MAX_BYTES,
        }),
    )?)?;
    if link_response.get("type").and_then(Value::as_str) != Some("artifact_download_link") {
        return Ok(json!({
            "status": "download_link_failed",
            "explicit_user_action": true,
            "local_path": &args.to,
            "local_bytes_written_by_cli": false,
            "content_bytes_available": false,
            "download_session": artifact_download_session_summary(&link_response),
            "grant_disclosures": [],
            "machine_error": artifact_response_machine_error(
                &link_response,
                "coordinator rejected artifact download for local export",
                "connectivity"
            ),
            "link_response": link_response,
        }));
    }
    download_link_to_path(
        session,
        coordinator,
        &args.scope,
        &args.artifact,
        DEFAULT_ARTIFACT_EXPORT_MAX_BYTES,
        &args.to,
        &link_response,
        stored_session,
    )
}

#[allow(clippy::too_many_arguments)]
fn download_link_to_path(
    session: &mut JsonLineSession,
    coordinator: &str,
    scope: &crate::CliScopeArgs,
    artifact: &str,
    max_bytes: u64,
    destination: &Path,
    link_response: &Value,
    stored_session: Option<&StoredCliSession>,
) -> Result<Value> {
    let download_session = artifact_download_session_summary(link_response);
    let grant_disclosures = artifact_download_grant_disclosures(link_response);
    let token_digest = link_response
        .pointer("/link/scoped_token_digest")
        .cloned()
        .context("artifact download link did not include a scoped token digest")?;
    let expected_digest: clusterflux_core::Digest = serde_json::from_value(
        link_response
            .pointer("/link/artifact_digest")
            .cloned()
            .context("artifact download link omitted the published digest")?,
    )?;
    let expected_size = link_response
        .pointer("/link/artifact_size_bytes")
        .and_then(Value::as_u64)
        .context("artifact download link omitted the published size")?;
    if expected_size > max_bytes {
        return Ok(json!({
            "status": "download_limit_failed",
            "local_path": destination,
            "local_bytes_written_by_cli": false,
            "content_bytes_available": false,
            "download_session": download_session,
            "grant_disclosures": grant_disclosures,
            "machine_error": cli_error_summary_for_category(
                "quota",
                &format!("artifact size {expected_size} exceeds requested limit {max_bytes}")
            ),
        }));
    }

    let parent = destination
        .parent()
        .filter(|parent| !parent.as_os_str().is_empty())
        .unwrap_or_else(|| Path::new("."));
    std::fs::create_dir_all(parent)
        .with_context(|| format!("failed to create {}", parent.display()))?;
    let mut temporary = tempfile::Builder::new()
        .prefix(".clusterflux-download-")
        .tempfile_in(parent)
        .with_context(|| {
            format!(
                "failed to create a bounded download in {}",
                parent.display()
            )
        })?;
    let started = Instant::now();
    let mut received_bytes = 0_u64;
    let mut hasher = Sha256::new();
    let stream_response = loop {
        let response = session.request(authenticated_or_local_trusted_request(
            coordinator,
            stored_session,
            json!({
                "type": "open_artifact_download_stream",
                "artifact": artifact,
                "max_bytes": max_bytes,
                "token_digest": token_digest,
                "chunk_bytes": expected_size.clamp(1, ARTIFACT_DOWNLOAD_CONTROL_CHUNK_BYTES),
            }),
            json!({
                "type": "open_artifact_download_stream",
                "tenant": scope.tenant,
                "project": scope.project,
                "actor_user": scope.user,
                "artifact": artifact,
                "max_bytes": max_bytes,
                "token_digest": token_digest,
                "chunk_bytes": expected_size.clamp(1, ARTIFACT_DOWNLOAD_CONTROL_CHUNK_BYTES),
            }),
        )?)?;
        if response.get("type").and_then(Value::as_str) != Some("artifact_download_stream") {
            break response;
        }
        if response
            .get("content_bytes_available")
            .and_then(Value::as_bool)
            == Some(true)
        {
            let offset = response
                .get("content_offset")
                .and_then(Value::as_u64)
                .unwrap_or(u64::MAX);
            if offset != received_bytes {
                break json!({
                    "type": "error",
                    "message": format!(
                        "artifact reverse stream returned offset {offset}, expected {received_bytes}"
                    ),
                });
            }
            let Some(content_base64) = response.get("content_base64").and_then(Value::as_str)
            else {
                break json!({
                    "type": "error",
                    "message": "artifact reverse stream marked bytes available without content",
                });
            };
            let content = BASE64_STANDARD
                .decode(content_base64)
                .context("artifact download stream returned invalid base64 content")?;
            received_bytes = received_bytes
                .checked_add(content.len() as u64)
                .context("artifact download byte count overflowed")?;
            if received_bytes > expected_size || received_bytes > max_bytes {
                break json!({
                    "type": "error",
                    "message": "artifact reverse stream exceeded its published size or CLI limit",
                });
            }
            temporary
                .write_all(&content)
                .context("write bounded artifact download")?;
            hasher.update(&content);
            if response.get("content_eof").and_then(Value::as_bool) == Some(true) {
                break response;
            }
            continue;
        }
        if started.elapsed() >= Duration::from_secs(120) {
            break json!({
                "type": "error",
                "message": "timed out waiting for the retaining node to reverse-stream artifact bytes",
            });
        }
        std::thread::sleep(Duration::from_millis(25));
    };
    if stream_response.get("type").and_then(Value::as_str) != Some("artifact_download_stream") {
        return Ok(json!({
            "status": "download_stream_failed",
            "explicit_user_action": true,
            "local_path": destination,
            "local_bytes_written_by_cli": false,
            "content_bytes_available": false,
            "download_session": download_session,
            "grant_disclosures": grant_disclosures,
            "machine_error": artifact_response_machine_error(
                &stream_response,
                "coordinator rejected artifact download stream",
                "connectivity"
            ),
            "stream": artifact_stream_summary(&stream_response),
        }));
    }
    let actual_digest_hex = format!("{:x}", hasher.finalize());
    let actual_digest = clusterflux_core::Digest::from_sha256_hex(&actual_digest_hex)
        .map_err(anyhow::Error::msg)?;
    if received_bytes != expected_size || actual_digest != expected_digest {
        return Ok(json!({
            "status": "download_integrity_failed",
            "explicit_user_action": true,
            "local_path": destination,
            "local_bytes_written_by_cli": false,
            "content_bytes_available": false,
            "download_session": download_session,
            "grant_disclosures": grant_disclosures,
            "machine_error": cli_error_summary_for_category(
                "program",
                &format!(
                    "artifact download digest/size mismatch: received {received_bytes} bytes with {actual_digest}, expected {expected_size} bytes with {expected_digest}"
                )
            ),
            "stream": artifact_stream_summary(&stream_response),
        }));
    }
    temporary
        .as_file()
        .sync_all()
        .context("sync verified artifact download")?;
    temporary.persist_noclobber(destination).map_err(|error| {
        anyhow::anyhow!(
            "refusing to replace artifact download destination {}: {}",
            destination.display(),
            error.error
        )
    })?;

    Ok(json!({
        "status": "local_bytes_written",
        "explicit_user_action": true,
        "local_path": destination,
        "local_bytes_written_by_cli": true,
        "bytes_written": received_bytes,
        "verified_digest": actual_digest,
        "published_digest": expected_digest,
        "content_bytes_available": true,
        "content_source": stream_response.get("content_source").cloned().unwrap_or(Value::Null),
        "download_session": download_session,
        "grant_disclosures": grant_disclosures,
        "stream": artifact_stream_summary(&stream_response),
    }))
}

pub(crate) fn artifact_stream_summary(response: &Value) -> Value {
    let status = response
        .get("type")
        .and_then(Value::as_str)
        .unwrap_or("coordinator_response");
    let mut summary = json!({
        "status": response.get("type").and_then(Value::as_str).unwrap_or("coordinator_response"),
        "streamed_bytes": response.get("streamed_bytes").cloned().unwrap_or_else(|| json!(0)),
        "charged_download_bytes": response.get("charged_download_bytes").cloned().unwrap_or_else(|| json!(0)),
        "content_bytes_available": response.get("content_bytes_available").cloned().unwrap_or(json!(false)),
        "content_offset": response.get("content_offset").cloned().unwrap_or(Value::Null),
        "content_eof": response.get("content_eof").cloned().unwrap_or(json!(false)),
        "content_source": response.get("content_source").cloned().unwrap_or(Value::Null),
        "content_material_returned_in_report": false,
        "error": response.get("message").cloned().unwrap_or(Value::Null),
    });
    if status != "artifact_download_stream" {
        if let Some(object) = summary.as_object_mut() {
            object.insert(
                "machine_error".to_owned(),
                artifact_response_machine_error(
                    response,
                    "coordinator rejected artifact download stream",
                    "connectivity",
                ),
            );
        }
    }
    summary
}

fn apply_local_export_summary(export_plan: &mut Value, local_export: &Value) {
    let Some(object) = export_plan.as_object_mut() else {
        return;
    };
    object.insert(
        "local_bytes_written_by_cli".to_owned(),
        local_export
            .get("local_bytes_written_by_cli")
            .cloned()
            .unwrap_or(json!(false)),
    );
    object.insert(
        "local_export_status".to_owned(),
        local_export
            .get("status")
            .cloned()
            .unwrap_or_else(|| json!("unknown")),
    );
    object.insert(
        "content_bytes_available".to_owned(),
        local_export
            .get("content_bytes_available")
            .cloned()
            .unwrap_or(json!(false)),
    );
    if let Some(bytes_written) = local_export.get("bytes_written") {
        object.insert("bytes_written".to_owned(), bytes_written.clone());
        object.insert(
            "writes_require_data_plane_followup".to_owned(),
            json!(false),
        );
    }
}