//! The async HTTP client: enqueue, poll, stream, list, and the text surface.

use std::time::Duration;

use futures_util::stream::Stream;
use serde::de::DeserializeOwned;
use serde::Serialize;

use crate::error::{ImaigeError, Result};
use crate::sse::SseDecoder;
use crate::types::{
    AnalyzeImageRequest, AnimateVideoRequest, GenerateImageRequest, GenerateTextRequest, JobEvent,
    JobFilter, JobHandle, JobList, JobResult, JobRow, JobStatus, ModelInfo,
    RemoveBackgroundRequest, TextResult, TextUsage, TransformImageRequest, UpscaleImageRequest,
};

/// Portable async sleep: tokio on native, gloo-timers on wasm.
#[cfg(not(target_arch = "wasm32"))]
async fn sleep(d: Duration) {
    tokio::time::sleep(d).await;
}

#[cfg(target_arch = "wasm32")]
async fn sleep(d: Duration) {
    gloo_timers::future::sleep(d).await;
}

/// Default API base URL.
pub const DEFAULT_BASE_URL: &str = "https://api.imaige.io";

/// The durable, 1-year-immutable media host. Media file IDs resolve to
/// `{MEDIA_BASE}/{fileId}`.
pub const MEDIA_BASE: &str = "https://api.imaige.io/media/media";

/// Resolve a media file id to its durable, 1-year-immutable URL.
pub fn media_url(file_id: &str) -> String {
    format!("{MEDIA_BASE}/{file_id}")
}

/// Configuration knobs for [`ImaigeClient`].
#[derive(Debug, Clone)]
pub struct ClientConfig {
    /// Per-request HTTP timeout. Does not bound `wait_for_job` polling.
    pub request_timeout: Duration,
    /// Interval between polls in `wait_for_job`.
    pub poll_interval: Duration,
    /// Overall deadline for `wait_for_job` before it returns `Timeout`.
    pub wait_timeout: Duration,
}

impl Default for ClientConfig {
    fn default() -> Self {
        Self {
            request_timeout: Duration::from_secs(60),
            poll_interval: Duration::from_secs(2),
            wait_timeout: Duration::from_secs(600),
        }
    }
}

/// Async, HTTP-only client for the Imaige API.
#[derive(Debug, Clone)]
pub struct ImaigeClient {
    http: reqwest::Client,
    base_url: String,
    api_key: String,
    config: ClientConfig,
}

impl ImaigeClient {
    /// Create a client against `base_url` (use [`DEFAULT_BASE_URL`] for prod)
    /// authenticating with `api_key` (an `imk_<prefix>_<secret>` key).
    pub fn new(base_url: impl Into<String>, api_key: impl Into<String>) -> Result<Self> {
        Self::with_config(base_url, api_key, ClientConfig::default())
    }

    /// Create a client with explicit [`ClientConfig`].
    pub fn with_config(
        base_url: impl Into<String>,
        api_key: impl Into<String>,
        config: ClientConfig,
    ) -> Result<Self> {
        let http = reqwest::Client::builder()
            .build()
            .map_err(ImaigeError::from)?;
        let base_url = base_url.into();
        let base_url = base_url.trim_end_matches('/').to_string();
        Ok(Self {
            http,
            base_url,
            api_key: api_key.into(),
            config,
        })
    }

    /// The configured base URL.
    pub fn base_url(&self) -> &str {
        &self.base_url
    }

    fn url(&self, path: &str) -> String {
        format!("{}{}", self.base_url, path)
    }

    /// Parse a `Retry-After` header into milliseconds. Supports the
    /// delta-seconds form (an integer count of seconds).
    fn retry_after_ms(resp: &reqwest::Response) -> Option<u64> {
        resp.headers()
            .get(reqwest::header::RETRY_AFTER)
            .and_then(|v| v.to_str().ok())
            .and_then(|s| s.trim().parse::<u64>().ok())
            .map(|secs| secs.saturating_mul(1000))
    }

    /// Issue a JSON request with auth + optional attribution headers, then
    /// deserialize a successful response into `T`. Non-2xx maps to the
    /// matching [`ImaigeError`] variant.
    async fn send_json<B, T>(
        &self,
        method: reqwest::Method,
        path: &str,
        body: Option<&B>,
        idempotency_key: Option<&str>,
        customer_id: Option<&str>,
    ) -> Result<T>
    where
        B: Serialize + ?Sized,
        T: DeserializeOwned,
    {
        let mut req = self
            .http
            .request(method, self.url(path))
            .timeout(self.config.request_timeout)
            .bearer_auth(&self.api_key);
        if let Some(b) = body {
            req = req.json(b);
        }
        if let Some(k) = idempotency_key {
            req = req.header("Idempotency-Key", k);
        }
        if let Some(c) = customer_id {
            req = req.header("X-Imaige-Customer-Id", c);
        }

        let resp = req.send().await.map_err(ImaigeError::from)?;
        let status = resp.status();
        if status.is_success() {
            let text = resp.text().await.map_err(ImaigeError::from)?;
            return serde_json::from_str::<T>(&text).map_err(|e| ImaigeError::Decode {
                message: format!("{e}; body was: {text}"),
            });
        }
        let retry = Self::retry_after_ms(&resp);
        let body = resp.text().await.unwrap_or_default();
        Err(ImaigeError::from_status(status.as_u16(), retry, body))
    }

    // ---- generation enqueues ------------------------------------------

    /// Enqueue any generation job by raw mode path + JSON body. Most callers
    /// should prefer the typed convenience wrappers below.
    pub async fn enqueue<B: Serialize>(
        &self,
        route: &str,
        body: &B,
        idempotency_key: Option<&str>,
        customer_id: Option<&str>,
    ) -> Result<JobHandle> {
        self.send_json(
            reqwest::Method::POST,
            route,
            Some(body),
            idempotency_key,
            customer_id,
        )
        .await
    }

    /// Text-to-image (`POST /v1/images/generate`).
    pub async fn generate_image(&self, req: &GenerateImageRequest) -> Result<JobHandle> {
        self.enqueue(
            "/v1/images/generate",
            req,
            req.attribution.idempotency_key.as_deref(),
            req.attribution.customer_id.as_deref(),
        )
        .await
    }

    /// Image-to-image transform (`POST /v1/images/transform`).
    pub async fn transform_image(&self, req: &TransformImageRequest) -> Result<JobHandle> {
        self.enqueue(
            "/v1/images/transform",
            req,
            req.attribution.idempotency_key.as_deref(),
            req.attribution.customer_id.as_deref(),
        )
        .await
    }

    /// AI upscale (`POST /v1/images/upscale`).
    pub async fn upscale_image(&self, req: &UpscaleImageRequest) -> Result<JobHandle> {
        self.enqueue(
            "/v1/images/upscale",
            req,
            req.attribution.idempotency_key.as_deref(),
            req.attribution.customer_id.as_deref(),
        )
        .await
    }

    /// Background removal (`POST /v1/images/remove-background`).
    pub async fn remove_background(&self, req: &RemoveBackgroundRequest) -> Result<JobHandle> {
        self.enqueue(
            "/v1/images/remove-background",
            req,
            req.attribution.idempotency_key.as_deref(),
            req.attribution.customer_id.as_deref(),
        )
        .await
    }

    /// VLM scene analysis (`POST /v1/images/analyze`).
    pub async fn analyze_image(&self, req: &AnalyzeImageRequest) -> Result<JobHandle> {
        self.enqueue(
            "/v1/images/analyze",
            req,
            req.attribution.idempotency_key.as_deref(),
            req.attribution.customer_id.as_deref(),
        )
        .await
    }

    /// Image-to-video animation (`POST /v1/videos/animate`).
    pub async fn generate_video_animate(&self, req: &AnimateVideoRequest) -> Result<JobHandle> {
        self.enqueue(
            "/v1/videos/animate",
            req,
            req.attribution.idempotency_key.as_deref(),
            req.attribution.customer_id.as_deref(),
        )
        .await
    }

    // ---- text generation ----------------------------------------------

    /// Enqueue a text generation job (`POST /v1/text/generate`), returning a
    /// [`JobHandle`] like every other mode.
    pub async fn generate_text(&self, req: &GenerateTextRequest) -> Result<JobHandle> {
        self.enqueue(
            "/v1/text/generate",
            req,
            req.attribution.idempotency_key.as_deref(),
            req.attribution.customer_id.as_deref(),
        )
        .await
    }

    /// Generate text and block until the answer is ready, returning the
    /// `{text, model, usage}` triple. Polls to completion under the hood.
    pub async fn generate_text_sync(&self, req: &GenerateTextRequest) -> Result<TextResult> {
        let handle = self.generate_text(req).await?;
        let result = self.wait_for_job(&handle.job_id).await?;
        if let Some(err) = result.error {
            return Err(ImaigeError::Http {
                status: 500,
                message: err,
            });
        }
        Self::extract_text(&result, &req.model)
    }

    /// Pull `{text, model, usage}` out of a completed text job's `result`
    /// payload, tolerant of several common server shapes.
    fn extract_text(result: &JobResult, fallback_model: &str) -> Result<TextResult> {
        let payload = result.result.clone().unwrap_or(serde_json::Value::Null);

        // text: top-level `text`, or OpenAI-style `choices[0].message.content`.
        let text = payload
            .get("text")
            .and_then(|v| v.as_str())
            .map(str::to_string)
            .or_else(|| {
                payload
                    .get("choices")
                    .and_then(|c| c.get(0))
                    .and_then(|c| c.get("message"))
                    .and_then(|m| m.get("content"))
                    .and_then(|v| v.as_str())
                    .map(str::to_string)
            })
            .ok_or_else(|| ImaigeError::Decode {
                message: format!("text job result missing `text`: {payload}"),
            })?;

        let model = payload
            .get("model")
            .and_then(|v| v.as_str())
            .map(str::to_string)
            .unwrap_or_else(|| fallback_model.to_string());

        let usage = payload
            .get("usage")
            .and_then(|u| serde_json::from_value::<TextUsage>(u.clone()).ok())
            .unwrap_or_default();

        Ok(TextResult { text, model, usage })
    }

    // ---- consumption: poll, wait, list --------------------------------

    /// Fetch one job row (`GET /v1/jobs/{id}`).
    pub async fn get_job(&self, job_id: &str) -> Result<JobRow> {
        self.send_json::<(), JobRow>(
            reqwest::Method::GET,
            &format!("/v1/jobs/{job_id}"),
            None,
            None,
            None,
        )
        .await
    }

    /// Poll a job to completion and return its [`JobResult`] with durable
    /// media URLs resolved. Returns [`ImaigeError::Timeout`] if the
    /// configured `wait_timeout` elapses first.
    pub async fn wait_for_job(&self, job_id: &str) -> Result<JobResult> {
        // Track elapsed time by accumulating the poll interval rather than
        // reading a monotonic clock, so the loop is portable to wasm32
        // (where `std::time::Instant` is unavailable).
        let mut elapsed = Duration::ZERO;
        loop {
            let row = self.get_job(job_id).await?;
            if row.status.is_terminal() {
                return Ok(Self::row_to_result(row));
            }
            if elapsed >= self.config.wait_timeout {
                return Err(ImaigeError::Timeout {
                    message: format!("job {job_id} did not complete within wait_timeout"),
                });
            }
            sleep(self.config.poll_interval).await;
            elapsed = elapsed.saturating_add(self.config.poll_interval);
        }
    }

    /// Convenience: enqueue a text-to-image job and wait for the result.
    pub async fn generate_image_and_wait(&self, req: &GenerateImageRequest) -> Result<JobResult> {
        let handle = self.generate_image(req).await?;
        self.wait_for_job(&handle.job_id).await
    }

    /// Convenience: enqueue a transform job and wait for the result.
    pub async fn transform_image_and_wait(&self, req: &TransformImageRequest) -> Result<JobResult> {
        let handle = self.transform_image(req).await?;
        self.wait_for_job(&handle.job_id).await
    }

    /// Convenience: enqueue an upscale job and wait for the result.
    pub async fn upscale_image_and_wait(&self, req: &UpscaleImageRequest) -> Result<JobResult> {
        let handle = self.upscale_image(req).await?;
        self.wait_for_job(&handle.job_id).await
    }

    /// Convenience: enqueue a background-removal job and wait for the result.
    pub async fn remove_background_and_wait(
        &self,
        req: &RemoveBackgroundRequest,
    ) -> Result<JobResult> {
        let handle = self.remove_background(req).await?;
        self.wait_for_job(&handle.job_id).await
    }

    /// Convenience: enqueue an analysis job and wait for the result.
    pub async fn analyze_image_and_wait(&self, req: &AnalyzeImageRequest) -> Result<JobResult> {
        let handle = self.analyze_image(req).await?;
        self.wait_for_job(&handle.job_id).await
    }

    /// Convenience: enqueue an animate job and wait for the result.
    pub async fn generate_video_animate_and_wait(
        &self,
        req: &AnimateVideoRequest,
    ) -> Result<JobResult> {
        let handle = self.generate_video_animate(req).await?;
        self.wait_for_job(&handle.job_id).await
    }

    /// List jobs scoped to the calling key (`GET /v1/jobs?…`).
    pub async fn list_jobs(&self, filter: &JobFilter) -> Result<JobList> {
        let mut params: Vec<(String, String)> = Vec::new();
        if let Some(c) = &filter.customer_id {
            params.push(("customerId".into(), c.clone()));
        }
        if let Some(s) = filter.status {
            let v = serde_json::to_value(s)
                .ok()
                .and_then(|v| v.as_str().map(str::to_string))
                .unwrap_or_default();
            params.push(("status".into(), v));
        }
        if let Some(m) = &filter.mode {
            params.push(("mode".into(), m.clone()));
        }
        if let Some(f) = &filter.from {
            params.push(("from".into(), f.clone()));
        }
        if let Some(t) = &filter.to {
            params.push(("to".into(), t.clone()));
        }
        if let Some(l) = filter.limit {
            params.push(("limit".into(), l.to_string()));
        }
        if let Some(cur) = &filter.cursor {
            params.push(("cursor".into(), cur.clone()));
        }

        let resp = self
            .http
            .get(self.url("/v1/jobs"))
            .timeout(self.config.request_timeout)
            .bearer_auth(&self.api_key)
            .query(&params)
            .send()
            .await
            .map_err(ImaigeError::from)?;
        let status = resp.status();
        if !status.is_success() {
            let retry = Self::retry_after_ms(&resp);
            let body = resp.text().await.unwrap_or_default();
            return Err(ImaigeError::from_status(status.as_u16(), retry, body));
        }
        let text = resp.text().await.map_err(ImaigeError::from)?;
        serde_json::from_str::<JobList>(&text).map_err(|e| ImaigeError::Decode {
            message: format!("{e}; body was: {text}"),
        })
    }

    /// List every job tagged with `customer_id` (convenience over
    /// [`list_jobs`](Self::list_jobs) with the filter pre-set).
    pub async fn get_jobs_by_custom_identifier(&self, customer_id: &str) -> Result<JobList> {
        self.list_jobs(&JobFilter {
            customer_id: Some(customer_id.to_string()),
            ..Default::default()
        })
        .await
    }

    /// List available models (`GET /models`). The model catalog (incl. pricing)
    /// is fetched from the Imaige API. Tolerates either a bare array
    /// or `{models: [...]}` / `{items: [...]}` envelopes.
    pub async fn list_models(&self) -> Result<Vec<ModelInfo>> {
        let resp = self
            .http
            .get(self.url("/models"))
            .timeout(self.config.request_timeout)
            .bearer_auth(&self.api_key)
            .send()
            .await
            .map_err(ImaigeError::from)?;
        let status = resp.status();
        if !status.is_success() {
            let retry = Self::retry_after_ms(&resp);
            let body = resp.text().await.unwrap_or_default();
            return Err(ImaigeError::from_status(status.as_u16(), retry, body));
        }
        let value: serde_json::Value = resp.json().await.map_err(ImaigeError::from)?;
        let arr = if value.is_array() {
            value
        } else {
            value
                .get("models")
                .or_else(|| value.get("items"))
                .cloned()
                .unwrap_or(serde_json::Value::Array(vec![]))
        };
        serde_json::from_value::<Vec<ModelInfo>>(arr).map_err(ImaigeError::from)
    }

    // ---- SSE streaming -------------------------------------------------

    /// Stream a job's progressive [`JobEvent`]s from
    /// `GET /v1/jobs/{id}/events`. The returned stream yields events until a
    /// terminal one (completion or error) or the underlying body ends.
    ///
    /// Each item is a `Result<JobEvent>`; a transport fault surfaces as an
    /// `Err(ImaigeError::Stream { .. })` and ends the stream.
    ///
    /// On native targets the returned stream is `Send` (so it can be driven
    /// from a multi-threaded runtime); on wasm32 it is not (the browser
    /// event loop is single-threaded and reqwest's fetch body is `!Send`).
    #[cfg(not(target_arch = "wasm32"))]
    pub async fn stream_job(
        &self,
        job_id: &str,
    ) -> Result<impl Stream<Item = Result<JobEvent>> + Send> {
        self.stream_job_impl(job_id).await
    }

    /// wasm32 variant of [`stream_job`](Self::stream_job) without the `Send`
    /// bound.
    #[cfg(target_arch = "wasm32")]
    pub async fn stream_job(&self, job_id: &str) -> Result<impl Stream<Item = Result<JobEvent>>> {
        self.stream_job_impl(job_id).await
    }

    /// Shared body of `stream_job` across targets. The concrete future type
    /// is `Send` on native and `!Send` on wasm purely from the underlying
    /// reqwest body stream, so both public signatures above are satisfied.
    async fn stream_job_impl(&self, job_id: &str) -> Result<impl Stream<Item = Result<JobEvent>>> {
        let resp = self
            .http
            .get(self.url(&format!("/v1/jobs/{job_id}/events")))
            .header(reqwest::header::ACCEPT, "text/event-stream")
            .bearer_auth(&self.api_key)
            .send()
            .await
            .map_err(ImaigeError::from)?;

        let status = resp.status();
        if !status.is_success() {
            let retry = Self::retry_after_ms(&resp);
            let body = resp.text().await.unwrap_or_default();
            return Err(ImaigeError::from_status(status.as_u16(), retry, body));
        }

        let job_id = job_id.to_string();
        let stream = async_stream::try_stream! {
            let mut decoder = SseDecoder::new();
            let mut byte_stream = resp.bytes_stream();
            use futures_util::StreamExt;

            while let Some(chunk) = byte_stream.next().await {
                let chunk = chunk.map_err(|e| ImaigeError::Stream { message: e.to_string() })?;
                for frame in decoder.push(&chunk) {
                    if let Some(event) = Self::frame_to_event(&frame, &job_id)? {
                        let terminal = event.terminal;
                        yield event;
                        if terminal {
                            return;
                        }
                    }
                }
            }
            if let Some(frame) = decoder.finish() {
                if let Some(event) = Self::frame_to_event(&frame, &job_id)? {
                    yield event;
                }
            }
        };
        Ok(stream)
    }

    /// Convert one SSE frame into a [`JobEvent`]. Empty/keepalive frames map
    /// to `None`. The frame's `data` is JSON; we enrich it with the job id
    /// and infer terminality from the status / event name.
    fn frame_to_event(frame: &crate::sse::SseFrame, job_id: &str) -> Result<Option<JobEvent>> {
        if frame.data.trim().is_empty() {
            return Ok(None);
        }
        let mut value: serde_json::Value =
            serde_json::from_str(&frame.data).map_err(|e| ImaigeError::Stream {
                message: format!("malformed SSE data frame: {e}; data was: {}", frame.data),
            })?;

        // Inject the job id if the payload omitted it.
        if let serde_json::Value::Object(map) = &mut value {
            map.entry("job_id")
                .or_insert_with(|| serde_json::Value::String(job_id.to_string()));

            // Resolve media file ids to URLs if the server sent ids rather
            // than URLs.
            if !map.contains_key("media_urls") {
                if let Some(ids) = map.get("mediaFileIds").and_then(|v| v.as_array()) {
                    let urls: Vec<serde_json::Value> = ids
                        .iter()
                        .filter_map(|v| v.as_str())
                        .map(|id| serde_json::Value::String(media_url(id)))
                        .collect();
                    if !urls.is_empty() {
                        map.insert("media_urls".into(), serde_json::Value::Array(urls));
                    }
                }
            }
        }

        let mut event: JobEvent =
            serde_json::from_value(value).map_err(|e| ImaigeError::Stream {
                message: format!("could not decode JobEvent: {e}"),
            })?;

        // Terminality: explicit `terminal`, a terminal status, an error, or a
        // `complete`/`done`/`error` event name.
        let event_terminal = matches!(
            frame.event.as_deref(),
            Some("complete") | Some("completed") | Some("done") | Some("error") | Some("failed")
        );
        let status_terminal = event.status.map(JobStatus::is_terminal).unwrap_or(false);
        event.terminal =
            event.terminal || event_terminal || status_terminal || event.error.is_some();

        Ok(Some(event))
    }

    /// Map a terminal job row into a [`JobResult`] with media URLs resolved.
    fn row_to_result(row: JobRow) -> JobResult {
        let media_urls = row.media_file_ids.iter().map(|id| media_url(id)).collect();
        JobResult {
            job_id: row.job_id,
            mode: row.mode,
            status: row.status,
            media_urls,
            media_file_ids: row.media_file_ids,
            result: row.result,
            actual_cost_cents: row.actual_cost_cents,
            error: row.error,
        }
    }
}
