nolgia 0.1.0

//! Hand-written. Submitting a generation, waiting for it, and reading what it
//! made.

use std::{collections::HashSet, fmt, sync::Arc, time::Duration};

use reqwest::{StatusCode, header::HeaderValue};
use serde_json::{Value, json};
use tokio::time::Instant;
use uuid::Uuid;

use crate::{
    Client, Typed,
    api::types,
    client::decode,
    error::{Error, JobError, Result},
    transport::{Attempts, IDEMPOTENCY_KEY},
};

/// A generation request ready to send: the endpoint and its JSON body.
///
/// Every typed request in [`api::types`](crate::api::types) that answers with
/// a job converts into one ([`IntoGeneration`]), as does its builder, so
/// [`Client::submit`] takes either:
///
/// ```no_run
/// # async fn run(nolgia: nolgia::Client) -> nolgia::Result<()> {
/// use nolgia::types::GenerateVideoRequest;
///
/// let request = GenerateVideoRequest::builder()
///     .model("seedance-2.5")
///     .prompt("a red fox trots through fresh snow at dawn");
/// let handle = nolgia.submit(request).await?;
/// # Ok(()) }
/// ```
///
/// [`image`](crate::image), [`video`](crate::video) and [`audio`](crate::audio)
/// build the common case without the builder.
///
/// [`Generation::json`] takes any endpoint and body as JSON, for a field this
/// build of the SDK predates.
#[derive(Debug, Clone)]
pub struct Generation {
    endpoint: &'static str,
    operation: &'static str,
    kind: Option<&'static str>,
    body: Value,
}

impl Generation {
    /// A raw JSON request to one of the endpoints that answer `202` with a job:
    /// `/generate/image`, `/generate/video`, `/generate/audio`, `/generate/3d`,
    /// `/restore/video` or `/remove-background/video`.
    pub fn json(endpoint: &str, body: Value) -> Result<Self> {
        let (endpoint, operation, kind) = match endpoint.trim_end_matches('/') {
            "/generate/image" => ("/generate/image", "generate_image", Some("image")),
            "/generate/video" => ("/generate/video", "generate_video", Some("video")),
            "/generate/audio" => ("/generate/audio", "generate_audio", Some("audio")),
            "/generate/3d" => ("/generate/3d", "generate3_d", Some("3d")),
            "/restore/video" => ("/restore/video", "restore_video", None),
            "/remove-background/video" => {
                ("/remove-background/video", "remove_background_video", None)
            }
            // /generate/set answers an OutputSet polled through GET /sets/{id},
            // not a job: call `generate_set()` directly.
            other => {
                return Err(Error::InvalidRequest(format!(
                    "{other} does not answer with a job; use the generated method for it"
                )));
            }
        };
        Ok(Self {
            endpoint,
            operation,
            kind,
            body,
        })
    }

    /// The endpoint this request posts to.
    pub fn endpoint(&self) -> &str {
        self.endpoint
    }

    /// The JSON body.
    pub fn body(&self) -> &Value {
        &self.body
    }
}

/// An image generation: `POST /generate/image` with `model` and `prompt`.
/// Add any other field of `GenerateImageRequest` with [`Generation::set`].
///
/// ```
/// let request = nolgia::image("gpt-image-2", "a red fox in the snow").set("aspect_ratio", "16:9");
/// assert_eq!(request.body()["aspect_ratio"], "16:9");
/// ```
pub fn image(model: impl Into<String>, prompt: impl Into<String>) -> Generation {
    simple(
        "/generate/image",
        "generate_image",
        Some("image"),
        model,
        prompt,
    )
}

/// A video generation: `POST /generate/video` with `model` and `prompt`.
/// Add `duration_seconds`, `aspect_ratio`, reference assets and the rest with
/// [`Generation::set`].
pub fn video(model: impl Into<String>, prompt: impl Into<String>) -> Generation {
    simple(
        "/generate/video",
        "generate_video",
        Some("video"),
        model,
        prompt,
    )
}

/// Audio: `POST /generate/audio` with `model` and `prompt` (the text to speak
/// for a voice model, the description for music and sound). Pick a voice with
/// `.set("voice", "...")`; `GET /voices` and `GET /models` list them.
pub fn audio(model: impl Into<String>, prompt: impl Into<String>) -> Generation {
    simple(
        "/generate/audio",
        "generate_audio",
        Some("audio"),
        model,
        prompt,
    )
}

fn simple(
    endpoint: &'static str,
    operation: &'static str,
    kind: Option<&'static str>,
    model: impl Into<String>,
    prompt: impl Into<String>,
) -> Generation {
    Generation {
        endpoint,
        operation,
        kind,
        body: json!({ "model": model.into(), "prompt": prompt.into() }),
    }
}

impl Generation {
    /// Set one more field of the request body.
    pub fn set(mut self, field: &str, value: impl serde::Serialize) -> Self {
        let value = serde_json::to_value(value).unwrap_or(Value::Null);
        if let Value::Object(map) = &mut self.body {
            map.insert(field.to_owned(), value);
        }
        self
    }
}

/// Anything [`Client::submit`] accepts. Implemented for the typed generation
/// requests, their builders, and [`Generation`].
pub trait IntoGeneration {
    /// Convert, validating a builder's required fields.
    fn into_generation(self) -> Result<Generation>;
}

impl IntoGeneration for Generation {
    fn into_generation(self) -> Result<Generation> {
        Ok(self)
    }
}

impl IntoGeneration for Result<Generation> {
    fn into_generation(self) -> Result<Generation> {
        self
    }
}

macro_rules! generation {
    ($ty:ident, $endpoint:literal) => {
        impl IntoGeneration for types::$ty {
            fn into_generation(self) -> Result<Generation> {
                let body = serde_json::to_value(&self)
                    .map_err(|err| Error::InvalidRequest(err.to_string()))?;
                Generation::json($endpoint, body)
            }
        }
        impl IntoGeneration for types::builder::$ty {
            fn into_generation(self) -> Result<Generation> {
                types::$ty::try_from(self)
                    .map_err(|err| Error::InvalidRequest(err.to_string()))?
                    .into_generation()
            }
        }
    };
}

generation!(GenerateImageRequest, "/generate/image");
generation!(GenerateVideoRequest, "/generate/video");
generation!(GenerateAudioRequest, "/generate/audio");
generation!(Generate3DRequest, "/generate/3d");
generation!(RestoreVideoRequest, "/restore/video");
generation!(RemoveBackgroundVideoRequest, "/remove-background/video");

/// Options for [`Client::submit_with`].
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct SubmitOptions {
    /// `Idempotency-Key` for this submit. Without one, the API refuses an
    /// identical submit within five minutes with `409` naming the existing
    /// job ([`ApiError::is_duplicate_submission`](crate::ApiError::is_duplicate_submission));
    /// a distinct key declares a deliberate new generation.
    pub idempotency_key: Option<String>,
}

impl SubmitOptions {
    /// Submit with this `Idempotency-Key`.
    pub fn idempotency_key(mut self, key: impl Into<String>) -> Self {
        self.idempotency_key = Some(key.into());
        self
    }
}

/// Called with each new state of a job while waiting.
pub type StatusCallback = Arc<dyn Fn(&types::Job) + Send + Sync>;

/// Options for [`Client::wait_job`].
#[derive(Clone)]
#[non_exhaustive]
pub struct WaitOptions {
    /// The whole wait. Default 30 minutes. Running out returns
    /// [`Error::Timeout`]; the job is NOT canceled and is still billed.
    pub timeout: Duration,
    /// One long-poll window (`GET /jobs/{id}/wait?timeout_seconds=`), at most
    /// 900 s. Default 300 s.
    pub poll_window: Duration,
    /// Called whenever the job's status, progress or message changes.
    pub on_status: Option<StatusCallback>,
}

impl Default for WaitOptions {
    fn default() -> Self {
        Self {
            timeout: Duration::from_secs(30 * 60),
            poll_window: Duration::from_secs(300),
            on_status: None,
        }
    }
}

impl fmt::Debug for WaitOptions {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("WaitOptions")
            .field("timeout", &self.timeout)
            .field("poll_window", &self.poll_window)
            .field("on_status", &self.on_status.is_some())
            .finish()
    }
}

impl WaitOptions {
    /// Wait at most `timeout` in total.
    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    /// Long-poll in windows of `window` (clamped to 1 to 900 s).
    pub fn poll_window(mut self, window: Duration) -> Self {
        self.poll_window = window;
        self
    }

    /// Call `callback` on every status change.
    pub fn on_status(mut self, callback: impl Fn(&types::Job) + Send + Sync + 'static) -> Self {
        self.on_status = Some(Arc::new(callback));
        self
    }
}

/// One file a finished job produced.
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct Media {
    /// The asset in your library.
    pub asset_id: Uuid,
    /// A signed, expiring download URL (see [`Client::download`]).
    pub url: String,
    /// `image`, `video`, `audio` or `3d`.
    pub modality: String,
    /// The MIME type, when known.
    pub content_type: Option<String>,
    /// When `url` stops working. Fetch the asset again for a fresh one.
    pub expires_at: chrono::DateTime<chrono::Utc>,
}

impl Media {
    fn from_asset(asset: &types::Asset) -> Self {
        Self {
            asset_id: asset.id,
            url: asset.signed_url.clone(),
            modality: asset.modality.to_string(),
            content_type: asset.mime_type.clone(),
            expires_at: asset.expires_at,
        }
    }
}

/// A finished, successful generation.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct GenerationResult {
    /// The terminal job.
    pub job: types::Job,
    /// Every asset the job produced, in the server's order.
    pub media: Vec<Media>,
}

impl GenerationResult {
    /// The first output's URL, when there is one.
    pub fn url(&self) -> Option<&str> {
        self.media.first().map(|m| m.url.as_str())
    }
}

/// A submitted job. Cheap to clone.
#[derive(Debug, Clone)]
pub struct JobHandle {
    client: Client,
    job: types::Job,
}

impl JobHandle {
    /// The job id.
    pub fn id(&self) -> Uuid {
        self.job.id
    }

    /// The job as last seen (the submit response, until [`refresh`](Self::refresh)).
    pub fn job(&self) -> &types::Job {
        &self.job
    }

    /// Fetch the job's current state once.
    pub async fn refresh(&mut self) -> Result<&types::Job> {
        self.job = self.client.get_job().id(self.job.id).send().typed().await?;
        Ok(&self.job)
    }

    /// Cancel the job on the server (`POST /jobs/{id}/cancel`).
    ///
    /// Returns the canceled job; its `cancellation` says what the provider did
    /// and whether the credits were refunded (`settlement` may read `pending`
    /// until the provider answers). A job that already finished is refused
    /// with `409` and code `job_not_cancellable`.
    pub async fn cancel(&self) -> Result<types::Job> {
        self.client
            .cancel_job()
            .id(self.job.id)
            .send()
            .typed()
            .await
    }

    /// Wait until the job succeeds (default [`WaitOptions`]).
    pub async fn wait(&self) -> Result<types::Job> {
        self.client
            .wait_job(self.job.id, WaitOptions::default())
            .await
    }

    /// Wait with options.
    pub async fn wait_with(&self, options: WaitOptions) -> Result<types::Job> {
        self.client.wait_job(self.job.id, options).await
    }

    /// Wait, then list what the job made.
    pub async fn result(&self) -> Result<GenerationResult> {
        self.result_with(WaitOptions::default()).await
    }

    /// [`result`](Self::result) with wait options.
    pub async fn result_with(&self, options: WaitOptions) -> Result<GenerationResult> {
        let job = self.wait_with(options).await?;
        let media = self.client.job_media(&job).await;
        Ok(GenerationResult { job, media })
    }

    /// Live status events for this job (see [`Client::job_events`]).
    pub async fn events(
        &self,
    ) -> Result<impl futures_util::Stream<Item = Result<crate::JobEvent>> + Send + 'static> {
        self.client.job_events(self.job.id).await
    }
}

/// What a status callback compares: status, progress, detail, message.
type Seen = (String, Option<f64>, Option<String>, Option<String>);

const TERMINAL_FAILURES: [&str; 2] = ["failed", "canceled"];

impl Client {
    /// Submit a generation and return at once with its job.
    ///
    /// Generation is asynchronous: the API answers `202` with a queued job.
    /// Wait with [`JobHandle::wait`] or [`JobHandle::result`], or use
    /// [`subscribe`](Self::subscribe) to do both.
    pub async fn submit(&self, request: impl IntoGeneration) -> Result<JobHandle> {
        self.submit_with(request, SubmitOptions::default()).await
    }

    /// [`submit`](Self::submit) with options (an `Idempotency-Key`).
    pub async fn submit_with(
        &self,
        request: impl IntoGeneration,
        options: SubmitOptions,
    ) -> Result<JobHandle> {
        let generation = request.into_generation()?;
        let mut builder = self
            .http()
            .post(self.url(generation.endpoint))
            .json(&generation.body);
        if let Some(key) = &options.idempotency_key {
            let value = HeaderValue::from_str(key).map_err(|_| {
                Error::InvalidRequest(format!(
                    "idempotency key {key:?} is not a valid header value"
                ))
            })?;
            builder = builder.header(IDEMPOTENCY_KEY, value);
        }
        let request = builder.build().map_err(Error::transport)?;
        let response =
            crate::transport::execute(self.http(), self.inner(), request, generation.operation)
                .await
                .map_err(Error::transport)?;
        let retried = response
            .extensions()
            .get::<Attempts>()
            .is_some_and(|attempts| attempts.0 > 0);
        if response.status().is_success() {
            let job: types::Job = decode(response).await?;
            return Ok(self.handle(job));
        }
        let error = Error::from_response(response).await;
        // Our own retry of a submit the server had in fact accepted: the API's
        // duplicate guard answers 409 naming that job. It is this call's job.
        if retried
            && let Error::Api(api) = &error
            && api.is_duplicate_submission()
            && let Some(job_id) = api.job_id
        {
            let job = self.get_job().id(job_id).send().typed().await?;
            return Ok(self.handle(job));
        }
        Err(error)
    }

    /// Submit, wait, and list the outputs.
    ///
    /// ```no_run
    /// # async fn run(nolgia: nolgia::Client) -> nolgia::Result<()> {
    /// let done = nolgia.subscribe(nolgia::image("gpt-image-2", "a red fox in the snow")).await?;
    /// nolgia.download(done.url().unwrap(), "fox.png").await?;
    /// # Ok(()) }
    /// ```
    ///
    /// A timeout only stops waiting: the job keeps running and is billed.
    pub async fn subscribe(&self, request: impl IntoGeneration) -> Result<GenerationResult> {
        self.submit(request).await?.result().await
    }

    /// A handle for a job submitted elsewhere.
    pub async fn job(&self, job_id: Uuid) -> Result<JobHandle> {
        let job = self.get_job().id(job_id).send().typed().await?;
        Ok(self.handle(job))
    }

    fn handle(&self, job: types::Job) -> JobHandle {
        JobHandle {
            client: self.clone(),
            job,
        }
    }

    /// What a generation would cost, without submitting it (`POST /jobs/cost`).
    /// Free: nothing is held or charged.
    ///
    /// The quote's `credits` is produced by the submit path's own pricer, so a
    /// quote and a submit of the same body agree to the credit.
    pub async fn quote(&self, request: impl IntoGeneration) -> Result<types::JobCostQuote> {
        let generation = request.into_generation()?;
        let kind = generation.kind.ok_or_else(|| {
            Error::InvalidRequest(format!("{} cannot be quoted", generation.endpoint))
        })?;
        let field = if kind == "3d" { "three_d" } else { kind };
        let body = json!({ "kind": kind, field: generation.body });
        self.send_json(
            self.http().post(self.url("/jobs/cost")).json(&body),
            "quote_job_cost",
        )
        .await
    }

    /// Wait for a job to finish, long-polling `GET /jobs/{id}/wait`.
    ///
    /// Returns the job once it `succeeded`; a `failed` or `canceled` job is
    /// [`Error::Job`] (with the typed code and the refund state), and running
    /// out of time is [`Error::Timeout`], which does NOT cancel the job.
    pub async fn wait_job(&self, job_id: Uuid, options: WaitOptions) -> Result<types::Job> {
        let deadline = Instant::now() + options.timeout;
        let mut last: Option<Seen> = None;
        loop {
            let remaining = deadline.saturating_duration_since(Instant::now());
            if remaining.is_zero() {
                return Err(Error::Timeout {
                    job_id: Some(job_id),
                });
            }
            let window = options
                .poll_window
                .min(Duration::from_secs(900))
                .min(remaining)
                .as_secs()
                .max(1);
            let request = self
                .http()
                .get(self.url(&format!("/jobs/{job_id}/wait")))
                .query(&[("timeout_seconds", window)])
                .timeout(Duration::from_secs(window + 30))
                .build()
                .map_err(Error::transport)?;
            let response =
                crate::transport::execute(self.http(), self.inner(), request, "wait_for_job")
                    .await
                    .map_err(Error::transport)?;
            if response.status() == StatusCode::REQUEST_TIMEOUT {
                continue; // The window closed with the job still running.
            }
            if !response.status().is_success() {
                return Err(Error::from_response(response).await);
            }
            let job: types::Job = decode(response).await?;
            let seen = (
                job.status.clone(),
                job.progress,
                job.status_detail.clone(),
                job.status_message.clone(),
            );
            if last.as_ref() != Some(&seen) {
                if let Some(callback) = &options.on_status {
                    callback(&job);
                }
                last = Some(seen);
            }
            match job.status.as_str() {
                "succeeded" => return Ok(job),
                status if TERMINAL_FAILURES.contains(&status) => {
                    return Err(Error::Job(Box::new(JobError::from_job(&job))));
                }
                _ => {}
            }
        }
    }

    /// The outputs of a finished job, in the server's order.
    ///
    /// One job can make several assets, and `job.asset` carries only one, so
    /// this lists by job id. A failed listing does not fail a paid, finished
    /// generation: it falls back to the inline asset (as the TypeScript and
    /// Python SDKs do).
    pub async fn job_media(&self, job: &types::Job) -> Vec<Media> {
        let listed = self
            .list_assets()
            .job_id(job.id.to_string())
            .limit(100u64)
            .send()
            .typed()
            .await
            .map(|page| page.items)
            .unwrap_or_default();
        let assets: Vec<&types::Asset> = if listed.is_empty() {
            job.asset.iter().collect()
        } else {
            listed.iter().collect()
        };
        let mut seen = HashSet::new();
        assets
            .into_iter()
            .filter(|asset| seen.insert(asset.id))
            .map(Media::from_asset)
            .collect()
    }
}

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

    #[test]
    fn builders_convert_and_validate() {
        let ok = types::GenerateVideoRequest::builder()
            .model("seedance-2.5")
            .prompt("a red fox")
            .into_generation()
            .unwrap();
        assert_eq!(ok.endpoint(), "/generate/video");
        assert_eq!(ok.body()["model"], "seedance-2.5");
        let missing = types::GenerateAudioRequest::builder()
            .model("eleven-v3")
            .into_generation();
        assert!(matches!(missing, Err(Error::InvalidRequest(_))));
    }

    #[test]
    fn raw_json_rejects_non_job_endpoints() {
        assert!(Generation::json("/generate/video", json!({})).is_ok());
        assert!(Generation::json("/generate/set", json!({})).is_err());
    }
}