#![allow(unused_imports)]

use super::*;

pub(crate) fn run_generate_request(
    mut request: GenerateRequest,
    fallback_id: String,
    dir: PathBuf,
    stream: Option<StreamContext>,
) -> Result<Value, String> {
    if request.prompt.trim().is_empty() {
        return Err("Prompt is required.".to_string());
    }
    let output_count = requested_n(request.n)?;
    if request.n.is_some() {
        request.n = Some(output_count);
    }
    let provider_supports_n = provider_supports_n(request.provider.as_deref());
    let payload = if provider_supports_n || output_count == 1 {
        let out = dir.join(format!(
            "out.{}",
            output_extension(request.format.as_deref())
        ));
        cli_json_result(&generate_args_with_recovery(
            &request,
            &out,
            provider_supports_n,
            Some((&fallback_id, &dir)),
        ))?
    } else {
        let recovery_targets = (0..output_count)
            .map(|index| {
                (
                    batch_recovery_job_id(&fallback_id, index),
                    batch_recovery_job_dir(&dir, index),
                )
            })
            .collect::<Vec<_>>();
        let arg_sets = recovery_targets
            .iter()
            .enumerate()
            .map(|(index, (recovery_job_id, recovery_job_dir))| {
                generate_args_with_recovery(
                    &request,
                    &batch_output_path(&dir, request.format.as_deref(), index as u8),
                    false,
                    Some((recovery_job_id.as_str(), recovery_job_dir.as_path())),
                )
            })
            .collect::<Vec<_>>();
        let partials = Arc::new(Mutex::new(Vec::<Value>::new()));
        let partials_for_cb = partials.clone();
        let stream_for_cb = stream.clone();
        let batch = run_payloads_concurrently_streaming(arg_sets, move |index, payload| {
            if let Some(ctx) = &stream_for_cb {
                let mut list = partials_for_cb
                    .lock()
                    .unwrap_or_else(|poisoned| poisoned.into_inner());
                apply_partial_output(ctx, &mut list, index, payload);
            }
        });
        let child_dirs = recovery_targets
            .iter()
            .map(|(_, recovery_dir)| recovery_dir.clone())
            .collect::<Vec<_>>();
        let merged = merge_batch_payloads(
            "images generate",
            output_count.into(),
            batch.payloads,
            batch.errors,
        );
        let generation_slots =
            generation_slots_from_batch_payload(output_count.into(), &merged, &child_dirs);
        let outputs_present = merged
            .get("output")
            .and_then(|output| output.get("files"))
            .and_then(Value::as_array)
            .map(Vec::len)
            .unwrap_or(0);
        let failures = merged
            .get("batch")
            .and_then(|batch| batch.get("failure_count"))
            .and_then(Value::as_u64)
            .unwrap_or(0) as usize;
        write_batch_recovery_summary(
            &fallback_id,
            &dir,
            &child_dirs,
            outputs_present,
            failures,
            generation_slots,
        )
        .map_err(app_error)?;
        merged
    };
    let request_meta = serde_json::to_value(&request).unwrap_or_else(|_| json!({}));
    let job = job_from_payload(&payload, &fallback_id, "images generate", request_meta);
    let event_type = terminal_event_type(job.get("status").and_then(Value::as_str));
    Ok(json!({
        "job_id": job.get("id").cloned().unwrap_or(Value::Null),
        "job": job,
        "events": [{
            "seq": 1,
            "kind": "local",
            "type": event_type,
            "data": {"status": job.get("status"), "output": payload.get("output"), "error": payload.get("error")}
        }],
        "payload": payload,
    }))
}
