//! HTTP client for the forge runner API (contract section 4).
//!
//! Runner calls (register, claim) authenticate with the runner secret. Job calls (logs,
//! heartbeat, artifacts, finish) authenticate with the per-job token.

use std::path::Path;
use std::time::Duration;

use forge_ci::proto::{
    ClaimRequest, FinishRequest, HeartbeatResponse, JobAssignment, RegisterRequest,
    RegisterResponse,
};
use reqwest::StatusCode;
use reqwest::header::{CONTENT_LENGTH, HeaderMap};
use serde::de::DeserializeOwned;
use sha2::{Digest, Sha256};
use tokio::io::AsyncWriteExt;

/// The claim long-poll is 25s on the forge side, so 35s leaves room for the response.
pub const CLAIM_TIMEOUT: Duration = Duration::from_secs(35);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
const LOG_TIMEOUT: Duration = Duration::from_secs(60);
const ARTIFACT_TIMEOUT: Duration = Duration::from_secs(60 * 60);
const BINARY_TIMEOUT: Duration = Duration::from_secs(10 * 60);
const HEADER_LOG_OFFSET: &str = "X-Forge-Log-Offset";
const HEADER_LOG_SIZE: &str = "X-Forge-Log-Size";
/// Sent on every claim: `<ARCH>-<OS>` of this runner, for example `x86_64-linux`.
const HEADER_RUNNER_TARGET: &str = "Forge-Runner-Target";
/// On a claim response: the sha256 (hex) of the forge's own binary, if the forge advertises one.
const HEADER_BINARY_SHA256: &str = "Forge-Binary-Sha256";
/// `GET`: the forge's binary, for a runner that is updating itself.
const BINARY_PATH: &str = "/api/v1/ci/runners/binary";

/// Errors from talking to the forge.
#[derive(Debug, thiserror::Error)]
pub enum ClientError {
    #[error("http: {0}")]
    Http(#[from] reqwest::Error),
    #[error("forge answered {status}: {body}")]
    Status { status: StatusCode, body: String },
    #[error("local io: {0}")]
    Io(#[from] std::io::Error),
    #[error("protocol: {0}")]
    Protocol(String),
}

/// The forge's answer to a log chunk.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LogAck {
    /// The chunk was stored (or was already stored, which is an idempotent retry).
    Accepted,
    /// The forge holds a different number of bytes. The runner must continue from this size.
    Resync(u64),
}

/// The forge's answer to a claim.
#[derive(Clone, Debug)]
pub struct ClaimReply {
    /// The job assigned to this runner, or `None` when there is nothing to run (204).
    pub job: Option<JobAssignment>,
    /// The sha256 (lowercase hex) of the forge's binary, when the response advertises one.
    pub binary_sha256: Option<String>,
}

/// A client for one forge base URL.
#[derive(Clone, Debug)]
pub struct ForgeClient {
    http: reqwest::Client,
    base: String,
}

impl ForgeClient {
    /// Builds a client for `base_url`, for example `https://forge.example`.
    pub fn new(base_url: &str) -> Result<Self, ClientError> {
        let http = reqwest::Client::builder()
            .user_agent(concat!("forge-runner/", env!("CARGO_PKG_VERSION")))
            .connect_timeout(Duration::from_secs(10))
            .build()?;
        Ok(Self {
            http,
            base: base_url.trim_end_matches('/').to_owned(),
        })
    }

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

    /// `POST /api/v1/ci/runners/register`.
    pub async fn register(&self, req: &RegisterRequest) -> Result<RegisterResponse, ClientError> {
        let resp = self
            .http
            .post(self.url("/api/v1/ci/runners/register"))
            .json(req)
            .timeout(REQUEST_TIMEOUT)
            .send()
            .await?;
        json_or_error(resp).await
    }

    /// `POST /api/v1/ci/runners/claim`. A 204 means nothing to run, and the reply has no job.
    pub async fn claim(
        &self,
        runner_secret: &str,
        labels: &[String],
    ) -> Result<ClaimReply, ClientError> {
        let req = ClaimRequest {
            labels: labels.to_vec(),
        };
        let target = format!("{}-{}", std::env::consts::ARCH, std::env::consts::OS);
        let resp = self
            .http
            .post(self.url("/api/v1/ci/runners/claim"))
            .bearer_auth(runner_secret)
            .header(HEADER_RUNNER_TARGET, target)
            .json(&req)
            .timeout(CLAIM_TIMEOUT)
            .send()
            .await?;
        let binary_sha256 = advertised_sha256(resp.headers());
        match resp.status() {
            StatusCode::NO_CONTENT => Ok(ClaimReply {
                job: None,
                binary_sha256,
            }),
            status if status.is_success() => Ok(ClaimReply {
                job: Some(resp.json().await?),
                binary_sha256,
            }),
            _ => Err(status_error(resp).await),
        }
    }

    /// `GET /api/v1/ci/runners/binary`. Streams the forge's binary into `dest` and returns the
    /// sha256 (lowercase hex) of the bytes written. The caller checks it against the advertised
    /// value and syncs the file.
    pub async fn download_binary(
        &self,
        runner_secret: &str,
        dest: &Path,
    ) -> Result<String, ClientError> {
        let mut resp = self
            .http
            .get(self.url(BINARY_PATH))
            .bearer_auth(runner_secret)
            .timeout(BINARY_TIMEOUT)
            .send()
            .await?;
        if !resp.status().is_success() {
            return Err(status_error(resp).await);
        }
        let mut file = tokio::fs::File::create(dest).await?;
        let mut hasher = Sha256::new();
        while let Some(chunk) = resp.chunk().await? {
            hasher.update(&chunk);
            file.write_all(&chunk).await?;
        }
        file.flush().await?;
        Ok(hex::encode(hasher.finalize()))
    }

    /// `POST /api/v1/ci/jobs/{id}/logs` with the byte offset of `body` in the job log.
    pub async fn post_logs(
        &self,
        job_id: &str,
        job_token: &str,
        offset: u64,
        body: Vec<u8>,
    ) -> Result<LogAck, ClientError> {
        let resp = self
            .http
            .post(self.url(&format!("/api/v1/ci/jobs/{job_id}/logs")))
            .bearer_auth(job_token)
            .header(HEADER_LOG_OFFSET, offset.to_string())
            .header(reqwest::header::CONTENT_TYPE, "text/plain; charset=utf-8")
            .body(body)
            .timeout(LOG_TIMEOUT)
            .send()
            .await?;
        match resp.status() {
            StatusCode::CONFLICT => {
                let size = resp
                    .headers()
                    .get(HEADER_LOG_SIZE)
                    .and_then(|v| v.to_str().ok())
                    .and_then(|v| v.parse::<u64>().ok())
                    .ok_or_else(|| {
                        ClientError::Protocol("409 without a valid X-Forge-Log-Size".into())
                    })?;
                Ok(LogAck::Resync(size))
            }
            status if status.is_success() => Ok(LogAck::Accepted),
            _ => Err(status_error(resp).await),
        }
    }

    /// `POST /api/v1/ci/jobs/{id}/heartbeat`.
    pub async fn heartbeat(
        &self,
        job_id: &str,
        job_token: &str,
    ) -> Result<HeartbeatResponse, ClientError> {
        let resp = self
            .http
            .post(self.url(&format!("/api/v1/ci/jobs/{job_id}/heartbeat")))
            .bearer_auth(job_token)
            .timeout(REQUEST_TIMEOUT)
            .send()
            .await?;
        json_or_error(resp).await
    }

    /// `PUT /api/v1/ci/jobs/{id}/artifacts/{name}`, streaming the file from disk.
    pub async fn upload_artifact(
        &self,
        job_id: &str,
        job_token: &str,
        name: &str,
        path: &Path,
    ) -> Result<(), ClientError> {
        let file = tokio::fs::File::open(path).await?;
        let len = file.metadata().await?.len();
        let stream = tokio_util::io::ReaderStream::new(file);
        let resp = self
            .http
            .put(self.url(&format!("/api/v1/ci/jobs/{job_id}/artifacts/{name}")))
            .bearer_auth(job_token)
            .header(CONTENT_LENGTH, len)
            .body(reqwest::Body::wrap_stream(stream))
            .timeout(ARTIFACT_TIMEOUT)
            .send()
            .await?;
        if resp.status().is_success() {
            Ok(())
        } else {
            Err(status_error(resp).await)
        }
    }

    /// `POST /api/v1/ci/jobs/{id}/finish`.
    pub async fn finish(
        &self,
        job_id: &str,
        job_token: &str,
        req: &FinishRequest,
    ) -> Result<(), ClientError> {
        let resp = self
            .http
            .post(self.url(&format!("/api/v1/ci/jobs/{job_id}/finish")))
            .bearer_auth(job_token)
            .json(req)
            .timeout(REQUEST_TIMEOUT)
            .send()
            .await?;
        if resp.status().is_success() {
            Ok(())
        } else {
            Err(status_error(resp).await)
        }
    }
}

/// The `Forge-Binary-Sha256` header, trimmed and lowercased. An empty value counts as absent.
fn advertised_sha256(headers: &HeaderMap) -> Option<String> {
    headers
        .get(HEADER_BINARY_SHA256)?
        .to_str()
        .ok()
        .map(|v| v.trim().to_ascii_lowercase())
        .filter(|v| !v.is_empty())
}

async fn json_or_error<T: DeserializeOwned>(resp: reqwest::Response) -> Result<T, ClientError> {
    if resp.status().is_success() {
        Ok(resp.json().await?)
    } else {
        Err(status_error(resp).await)
    }
}

async fn status_error(resp: reqwest::Response) -> ClientError {
    let status = resp.status();
    let body: String = resp
        .text()
        .await
        .unwrap_or_default()
        .chars()
        .take(512)
        .collect();
    ClientError::Status { status, body }
}