#![allow(unused_imports)]

use super::*;

pub(crate) fn collect_history_ids(payload: &Value) -> Vec<String> {
    let mut ids = Vec::new();
    if let Some(id) = payload
        .get("history")
        .and_then(|history| history.get("job_id"))
        .and_then(Value::as_str)
        && !id.is_empty()
    {
        ids.push(id.to_string());
    }
    if let Some(job_ids) = payload
        .get("history")
        .and_then(|history| history.get("job_ids"))
        .and_then(Value::as_array)
    {
        for id in job_ids.iter().filter_map(Value::as_str) {
            if !id.is_empty() && !ids.iter().any(|existing| existing == id) {
                ids.push(id.to_string());
            }
        }
    }
    ids
}

pub(crate) fn output_files_from_payload(payload: &Value) -> Vec<Value> {
    let output = payload.get("output").cloned().unwrap_or_else(|| json!({}));
    let mut files = output
        .get("files")
        .and_then(Value::as_array)
        .cloned()
        .unwrap_or_default();
    if files.is_empty()
        && let Some(path) = output.get("path").and_then(Value::as_str)
    {
        files.push(json!({
            "index": 0,
            "path": path,
            "bytes": output.get("bytes").and_then(Value::as_u64).unwrap_or(0),
        }));
    }
    files
}

pub(crate) fn normalize_batch_output(files: Vec<Value>) -> Value {
    let indexed_files = files
        .into_iter()
        .enumerate()
        .map(|(index, mut file)| {
            if let Value::Object(object) = &mut file {
                object
                    .entry("index".to_string())
                    .or_insert_with(|| json!(index));
            }
            file
        })
        .collect::<Vec<_>>();
    let total_bytes = indexed_files
        .iter()
        .filter_map(|file| file.get("bytes").and_then(Value::as_u64))
        .sum::<u64>();
    let primary_path = indexed_files
        .iter()
        .find(|file| file.get("index").and_then(Value::as_u64) == Some(0))
        .and_then(|file| file.get("path"))
        .cloned()
        .unwrap_or(Value::Null);
    json!({
        "path": primary_path,
        "bytes": total_bytes,
        "files": indexed_files,
    })
}

pub(crate) fn batch_errors_json(errors: &[BatchItemError]) -> Value {
    Value::Array(
        errors
            .iter()
            .map(|error| {
                json!({
                    "index": error.index,
                    "message": error.message,
                })
            })
            .collect(),
    )
}

pub(crate) fn batch_error_summary(errors: &[BatchItemError]) -> Option<String> {
    if errors.is_empty() {
        return None;
    }
    let first = errors
        .first()
        .map(|error| error.message.as_str())
        .unwrap_or("Unknown batch error.");
    if errors.len() == 1 {
        Some(first.to_string())
    } else {
        Some(format!("{} 个子任务失败：{first}", errors.len()))
    }
}

pub(crate) fn merge_batch_payloads(
    command: &str,
    request_count: usize,
    payloads: Vec<(usize, Value)>,
    errors: Vec<BatchItemError>,
) -> Value {
    let first = payloads
        .first()
        .map(|(_, payload)| payload.clone())
        .unwrap_or_else(|| json!({}));
    let files = payloads
        .iter()
        .flat_map(|(batch_index, payload)| {
            output_files_from_payload(payload)
                .into_iter()
                .map(move |mut file| {
                    if let Value::Object(object) = &mut file {
                        object.insert("index".to_string(), json!(batch_index));
                    }
                    file
                })
        })
        .collect::<Vec<_>>();
    let mut history_job_ids = Vec::new();
    let mut revised_prompts = Vec::new();

    for (_, payload) in &payloads {
        history_job_ids.extend(collect_history_ids(payload));
        if let Some(prompts) = payload
            .get("response")
            .and_then(|response| response.get("revised_prompts"))
            .and_then(Value::as_array)
        {
            revised_prompts.extend(prompts.iter().cloned());
        }
    }

    history_job_ids.sort();
    history_job_ids.dedup();
    let primary_history_job_id = history_job_ids.first().cloned();
    let output = normalize_batch_output(files);
    let image_count = output
        .get("files")
        .and_then(Value::as_array)
        .map(Vec::len)
        .unwrap_or(0);
    let mut response = first.get("response").cloned().unwrap_or_else(|| json!({}));
    if let Value::Object(response) = &mut response {
        response.insert("image_count".to_string(), json!(image_count));
        response.insert("batch_count".to_string(), json!(request_count));
        response.insert("batch_request_count".to_string(), json!(request_count));
        response.insert("revised_prompts".to_string(), json!(revised_prompts));
    }
    let error_summary = batch_error_summary(&errors);
    let ok = image_count > 0;
    let status = if ok && errors.is_empty() {
        "completed"
    } else if ok {
        "partial_failed"
    } else {
        "failed"
    };

    let mut payload = json!({
        "ok": ok,
        "status": status,
        "command": command,
        "provider": first.get("provider").cloned().unwrap_or(Value::Null),
        "provider_selection": first.get("provider_selection").cloned().unwrap_or(Value::Null),
        "auth": first.get("auth").cloned().unwrap_or(Value::Null),
        "request": first.get("request").cloned().unwrap_or(Value::Null),
        "response": response,
        "output": output,
        "history": {
            "job_id": primary_history_job_id,
            "job_ids": history_job_ids,
        },
        "batch": {
            "mode": "parallel-single-output",
            "request_count": request_count,
            "success_count": image_count,
            "failure_count": errors.len(),
            "errors": batch_errors_json(&errors),
        },
        "events": {
            "count": request_count,
        }
    });
    if !errors.is_empty()
        && let Value::Object(object) = &mut payload
    {
        object.insert(
            "error".to_string(),
            json!({
                "code": if ok { "batch_partial_failed" } else { "batch_failed" },
                "message": error_summary.unwrap_or_else(|| "Batch request failed.".to_string()),
                "items": batch_errors_json(&errors),
            }),
        );
    }
    payload
}

pub(crate) fn cleanup_child_history(payload: &Value, app_job_id: &str) {
    for id in collect_history_ids(payload) {
        if id != app_job_id {
            let _ = delete_history_job(&id);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn payload(path: &str) -> Value {
        json!({
            "ok": true,
            "provider": "mock",
            "output": {
                "path": path,
                "bytes": 10,
            },
            "history": {
                "job_id": format!("child-{path}"),
            },
            "response": {
                "revised_prompts": [],
            },
        })
    }

    #[test]
    fn merge_batch_payloads_keeps_successful_outputs_with_failed_items() {
        let merged = merge_batch_payloads(
            "images generate",
            3,
            vec![(0, payload("/tmp/a.png")), (2, payload("/tmp/c.png"))],
            vec![BatchItemError {
                index: 1,
                message: "upstream rejected candidate B".to_string(),
            }],
        );

        assert_eq!(merged["status"], "partial_failed");
        assert_eq!(merged["ok"], true);
        let files = merged["output"]["files"].as_array().unwrap();
        assert_eq!(files.len(), 2);
        assert_eq!(files[0]["index"], 0);
        assert_eq!(files[0]["path"], "/tmp/a.png");
        assert_eq!(files[1]["index"], 2);
        assert_eq!(files[1]["path"], "/tmp/c.png");
        assert_eq!(merged["output"]["path"], "/tmp/a.png");
        assert_eq!(merged["batch"]["request_count"], 3);
        assert_eq!(merged["batch"]["success_count"], 2);
        assert_eq!(merged["batch"]["failure_count"], 1);
        assert_eq!(merged["batch"]["errors"][0]["index"], 1);
        assert_eq!(merged["error"]["message"], "upstream rejected candidate B");
    }

    #[test]
    fn merge_batch_payloads_marks_total_failure_not_ok() {
        let merged = merge_batch_payloads(
            "images generate",
            2,
            vec![],
            vec![
                BatchItemError {
                    index: 0,
                    message: "candidate A failed".to_string(),
                },
                BatchItemError {
                    index: 1,
                    message: "candidate B failed".to_string(),
                },
            ],
        );

        assert_eq!(merged["ok"], false);
        assert_eq!(merged["status"], "failed");
        assert_eq!(merged["output"]["files"].as_array().unwrap().len(), 0);
        assert!(merged["output"]["path"].is_null());
        assert_eq!(merged["batch"]["success_count"], 0);
        assert_eq!(merged["batch"]["failure_count"], 2);
        assert_eq!(merged["error"]["code"], "batch_failed");
        assert_eq!(merged["error"]["items"][0]["index"], 0);

        let job = job_from_payload(&merged, "job-1", "images generate", json!({}));
        assert_eq!(job["status"], "failed");
        assert_eq!(terminal_event_type(job["status"].as_str()), "job.failed");
        assert!(!terminal_status_runs_storage_upload(job["status"].as_str()));
    }
}
