//! /store and /stream handlers.

use std::io::Error as IoError;

use axum::{
    Extension,
    body::{Body, Bytes},
    extract::{Path, Query, RawQuery, State},
    http::{
        HeaderMap, HeaderValue, StatusCode,
        header::{ACCEPT, CONTENT_TYPE, ETAG, LOCATION},
    },
    response::{IntoResponse, Response},
};
use serde::Deserialize;
use serde_json::Value as JsonValue;
use sha2::{Digest, Sha256};
use tokio::sync::broadcast;
use tracing::{debug, error, instrument, warn};

use super::auth::AuthenticatedUser;
use super::{ApiError, media_type, store_status};
use crate::document::ChangeHash;
use crate::document::meta;
use crate::document::store::{self, StoreError, paths_related};
use crate::state::{AppState, ReadSource};
use crate::transport::Hex;

#[derive(Debug, Default, Deserialize)]
pub(crate) struct ReadQuery {
    #[serde(default)]
    conflicts: bool,
    /// Return the document's heads at the time of the read (`x-continuum-heads`), e.g.
    /// to base a replacing remake on exactly what was read.
    #[serde(default)]
    heads: bool,
    /// What's there, one level deep (`store::read_shallow`), to browse a large store a
    /// level at a time.
    #[serde(default)]
    shallow: bool,
    /// Read as the document was at these heads (comma-separated hex, as
    /// `x-continuum-heads` gives them), not as it is now (`document::at`).
    at: Option<String>,
}

impl ReadQuery {
    /// The heads `at` names, if any; `400` if it doesn't name them.
    fn at(&self) -> Result<Option<Vec<ChangeHash>>, ApiError> {
        let Some(at) = &self.at else {
            return Ok(None);
        };
        at.split(',')
            .map(|head| head.trim().parse::<ChangeHash>())
            .collect::<Result<Vec<_>, _>>()
            .map(Some)
            .map_err(|_| ApiError(StatusCode::BAD_REQUEST, Some("at: not heads".to_owned())))
    }
}

/// Reads `doc` with `read`, as it is or, with `at`, as it was at those heads
/// (`UnknownHeads` if it lacks one).
macro_rules! read_at {
    ($doc:expr, $at:expr, |$d:ident| $read:expr) => {
        match $at {
            None => {
                let $d = $doc;
                $read
            }
            Some(heads) => match crate::document::at::At::new($doc, heads) {
                Some(at) => {
                    let $d = &at;
                    $read
                }
                None => Err(StoreError::UnknownHeads),
            },
        }
    };
}

#[derive(Debug, Deserialize)]
pub(crate) struct PostQuery {
    /// Adds this to the counter at the path instead of inserting the body
    /// (`store::increment_path`).
    increment: Option<i64>,
}

#[derive(Debug, Deserialize)]
pub(crate) struct DeleteQuery {
    #[serde(default)]
    prune: bool,
}

fn conflicts_response(
    value: JsonValue,
    conflicts: Vec<store::ConflictReport>,
) -> Result<Response, ApiError> {
    let body = serde_json::json!({ "value": value, "conflicts": conflicts });
    response_for_value(&body, None) // always JSON
}

/// `GET /store`: the whole store (`docs/openapi.yaml`).
pub(crate) async fn get_store_root(
    State(state): State<AppState>,
    headers: HeaderMap,
    Query(query): Query<ReadQuery>,
) -> Result<Response, ApiError> {
    read(&state, Vec::new(), &headers, &query).await
}

/// `GET /store/{key}`.
pub(crate) async fn get_store_entry(
    State(state): State<AppState>,
    Path(key): Path<String>,
    headers: HeaderMap,
    Query(query): Query<ReadQuery>,
    RawQuery(raw_query): RawQuery,
) -> Result<Response, ApiError> {
    let path = parse_key_path(&key)?;
    let at = query.at()?;
    if let Some(redirect) =
        to_positive(&state, &path, false, at.as_deref(), raw_query.as_deref()).await?
    {
        return Ok(redirect);
    }
    read(&state, path, &headers, &query).await
}

/// A read of `path` as `query` asks (the empty path: the whole store). The root is read
/// while not ready too (`render_read`'s `unready_live`), and is never missing: an empty
/// store is an empty map.
async fn read(
    state: &AppState,
    path: Vec<String>,
    headers: &HeaderMap,
    query: &ReadQuery,
) -> Result<Response, ApiError> {
    let root = path.is_empty();
    if query.shallow {
        return shallow(state, path, root, query).await;
    }
    let accept = headers.get(ACCEPT).cloned();
    let want_conflicts = query.conflicts;
    let at = query.at()?;

    let (response, source, heads) = state
        .render_read(
            root,
            query.heads,
            move |doc| {
                read_at!(doc, at.as_deref(), |doc| if want_conflicts {
                    store::read_path_with_conflicts(doc, &path)
                        .map(|found| found.map(|(value, conflicts)| (value, Some(conflicts))))
                } else {
                    store::read_path(doc, &path).map(|found| found.map(|value| (value, None)))
                })
            },
            move |result| {
                let found = result.map_err(|e| store_status(e, "Failed to read store path"))?;
                let (value, conflicts) = match found {
                    Some(found) => found,
                    None if root => (JsonValue::Object(Default::default()), Some(Vec::new())),
                    None => return Err(StatusCode::NOT_FOUND.into()),
                };
                match conflicts {
                    Some(conflicts) => conflicts_response(value, conflicts),
                    None => response_for_value(&value, accept.as_ref()),
                }
            },
        )
        .await?;
    mark_source(response, source, heads)
}

/// `?shallow=true`: what's at `path`, one level deep, as JSON (`store::read_shallow`);
/// with conflicts asked for too, `400`. `unready_live` as `render_read` has it.
async fn shallow(
    state: &AppState,
    path: Vec<String>,
    unready_live: bool,
    query: &ReadQuery,
) -> Result<Response, ApiError> {
    if query.conflicts {
        return Err(StatusCode::BAD_REQUEST.into());
    }
    let at = query.at()?;
    let (response, source, heads) = state
        .render_read(
            unready_live,
            query.heads,
            move |doc| read_at!(doc, at.as_deref(), |doc| store::read_shallow(doc, &path)),
            move |result| {
                let value = result
                    .map_err(|e| store_status(e, "Failed to read store path"))?
                    .ok_or(StatusCode::NOT_FOUND)?;
                response_for_value(&value, None)
            },
        )
        .await?;
    mark_source(response, source, heads)
}

/// Header on reads served from the pre-rejoin document; its value is that epoch.
const STALE_HEADER: &str = "x-continuum-stale";

/// Header with the document's heads at the time of the read (sorted, comma-separated
/// hex), on request (`?heads=true`).
const HEADS_HEADER: &str = "x-continuum-heads";

fn mark_source(
    response: Result<Response, ApiError>,
    source: ReadSource,
    heads: Option<Vec<ChangeHash>>,
) -> Result<Response, ApiError> {
    let mut response = response?;
    if let Some(heads) = heads {
        let joined = heads
            .iter()
            .map(ToString::to_string)
            .collect::<Vec<_>>()
            .join(",");
        if let Ok(value) = HeaderValue::from_str(&joined) {
            response.headers_mut().insert(HEADS_HEADER, value);
        }
    }
    if let ReadSource::Stale(epoch) = source
        && let Ok(value) = HeaderValue::from_str(&epoch.to_string())
    {
        response.headers_mut().insert(STALE_HEADER, value);
    }
    Ok(response)
}

/// `PUT /store/{key}`.
pub(crate) async fn put_store_entry(
    State(state): State<AppState>,
    Extension(user): Extension<AuthenticatedUser>,
    Path(key): Path<String>,
    RawQuery(raw_query): RawQuery,
    headers: HeaderMap,
    body: Bytes,
) -> Result<Response, ApiError> {
    if !state.is_ready().await {
        return Err(StatusCode::SERVICE_UNAVAILABLE.into());
    }

    let path = parse_write_path(&state, &user, &key, "PUT")?;
    if let Some(redirect) = to_positive(&state, &path, false, None, raw_query.as_deref()).await? {
        return Ok(redirect);
    }
    let value = body_value(&headers, &body)?;

    let written = value.clone();
    state
        .mutate_doc_for(Some(user.0.clone()), move |doc| {
            store::write_path(doc, &path, &written)
        })
        .await
        .map_err(|e| store_status(e, "Failed to write store value"))?;

    response_for_value(&value, headers.get(ACCEPT))
}

fn etag_for(bytes: &[u8]) -> String {
    format!("\"{}\"", Hex(&Sha256::digest(bytes)))
}

fn parse_key_path(key: &str) -> Result<Vec<String>, StatusCode> {
    let segments: Vec<String> = key.split('/').map(str::to_owned).collect();

    if segments.is_empty() || segments.iter().any(|segment| segment.is_empty()) {
        return Err(StatusCode::BAD_REQUEST);
    }

    Ok(segments)
}

/// A path for a write. Protocol data (`_continuum`: remake records, members, forget
/// records, epoch metadata) is what every node acts on, so a write there acts on the
/// whole swarm: only users listed in `--protocol-writers` may, and each such write is
/// logged. Epoch metadata (and `_continuum` as a whole, which contains it) is never
/// writable: the protocol reads it at the snapshot change, so an edit only misleads
/// whoever reads it later. `403` otherwise.
fn parse_write_path(
    state: &AppState,
    user: &AuthenticatedUser,
    key: &str,
    method: &str,
) -> Result<Vec<String>, StatusCode> {
    let path = parse_key_path(key)?;
    if !meta::is_protocol(&path) {
        return Ok(path);
    }
    // `_continuum` itself, epoch metadata and the swarm ID: nobody.
    let fixed = path.len() < 2 || path[1] == meta::EPOCHS || path[1] == meta::SWARM;
    if fixed || !state.protocol_writers.contains(&user.0) {
        warn!(user = %user.0, method, path = %key, "Refused a write to protocol data");
        return Err(StatusCode::FORBIDDEN);
    }
    warn!(user = %user.0, method, path = %key, "Protocol data written through the store API; this acts on the whole swarm");
    Ok(path)
}

fn response_for_value(
    value: &JsonValue,
    accept: Option<&HeaderValue>,
) -> Result<Response, ApiError> {
    let wants_text = accept
        .and_then(|value| value.to_str().ok())
        .is_some_and(|value| {
            value
                .split(',')
                .any(|item| media_type(item) == "text/plain")
        });

    let (content_type, body) = if wants_text {
        let body = match value {
            JsonValue::String(value) => value.clone(),
            other => serde_json::to_string(other).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
        };

        ("text/plain; charset=utf-8", body)
    } else {
        (
            "application/json",
            serde_json::to_string(value).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
        )
    };

    let etag = etag_for(body.as_bytes());

    let mut response = body.into_response();

    response
        .headers_mut()
        .insert(CONTENT_TYPE, HeaderValue::from_static(content_type));

    response.headers_mut().insert(
        ETAG,
        HeaderValue::from_str(&etag).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
    );

    Ok(response)
}

/// A write's body as a value: JSON (the default), or text as a string.
fn body_value(headers: &HeaderMap, body: &Bytes) -> Result<JsonValue, StatusCode> {
    let content_type = headers
        .get(CONTENT_TYPE)
        .and_then(|value| value.to_str().ok())
        .unwrap_or("application/json");
    match media_type(content_type) {
        "application/json" => {
            serde_json::from_slice::<JsonValue>(body).map_err(|_| StatusCode::BAD_REQUEST)
        }
        "text/plain" => String::from_utf8(body.to_vec())
            .map(JsonValue::String)
            .map_err(|_| StatusCode::BAD_REQUEST),
        _ => Err(StatusCode::UNSUPPORTED_MEDIA_TYPE),
    }
}

/// `POST /store/{key}`: inserts the body into a list, `key` ending in the index it
/// goes to (`store::insert_path`: the items from there on move up; the length, or `-1`,
/// appends); `201` with where it went (`{"index": n}`). Not a list there, or an index
/// past its end: `409`, saying which. With `?increment=n`, adds `n` to the counter at
/// `key` instead (`store::increment_path`), answering with its value then; not a
/// counter there: `409`.
pub(crate) async fn post_store_entry(
    State(state): State<AppState>,
    Extension(user): Extension<AuthenticatedUser>,
    Path(key): Path<String>,
    Query(query): Query<PostQuery>,
    RawQuery(raw_query): RawQuery,
    headers: HeaderMap,
    body: Bytes,
) -> Result<Response, ApiError> {
    if !state.is_ready().await {
        return Err(StatusCode::SERVICE_UNAVAILABLE.into());
    }
    let path = parse_write_path(&state, &user, &key, "POST")?;
    // An increment names an item (`-1`: the last); an insert where it goes (`-1`: the end).
    let insertion = query.increment.is_none();
    if let Some(redirect) =
        to_positive(&state, &path, insertion, None, raw_query.as_deref()).await?
    {
        return Ok(redirect);
    }
    if let Some(by) = query.increment {
        let value = state
            .mutate_doc_for(Some(user.0.clone()), move |doc| {
                store::increment_path(doc, &path, by)
            })
            .await
            .map_err(|e| store_status(e, "Failed to increment a counter"))?;
        return response_for_value(&JsonValue::from(value), headers.get(ACCEPT));
    }
    let index = path.last().and_then(|leaf| store::list_index(leaf));
    let value = body_value(&headers, &body)?;
    let inserted = state
        .mutate_doc_for(Some(user.0.clone()), move |doc| {
            store::insert_path(doc, &path, &value)
        })
        .await;
    match inserted {
        Ok(()) => Ok((
            StatusCode::CREATED,
            axum::Json(serde_json::json!({ "index": index })),
        )
            .into_response()),
        Err(error) => Err(store_status(error, "Failed to insert into a list")),
    }
}

/// `DELETE /store/{key}`.
#[instrument(skip(state), fields(path, prune = query.prune))]
pub(crate) async fn delete_store_entry(
    State(state): State<AppState>,
    Extension(user): Extension<AuthenticatedUser>,
    Path(key): Path<String>,
    Query(query): Query<DeleteQuery>,
    RawQuery(raw_query): RawQuery,
) -> Result<Response, ApiError> {
    let path = parse_write_path(&state, &user, &key, "DELETE")?;

    tracing::Span::current().record("path", path.join("/").as_str());

    if !state.is_ready().await {
        return Err(StatusCode::SERVICE_UNAVAILABLE.into());
    }

    if path.is_empty() {
        return Err(StatusCode::BAD_REQUEST.into());
    }

    if let Some(redirect) = to_positive(&state, &path, false, None, raw_query.as_deref()).await? {
        return Ok(redirect);
    }

    let prune = query.prune;
    state
        .mutate_doc_for(Some(user.0.clone()), move |doc| {
            store::delete_path(doc, &path, prune)
        })
        .await
        .map_err(|e| store_status(e, "Failed to delete store value"))?;

    Ok(StatusCode::NO_CONTENT.into_response())
}

/// A path with negative list indexes (`-1` the last item; for an insert, the end): a
/// `307` to it with them made positive (`store::positive_path`), the query kept, so a
/// client reads, writes or deletes what the index names now. None without any.
async fn to_positive(
    state: &AppState,
    path: &[String],
    insertion: bool,
    at: Option<&[ChangeHash]>,
    query: Option<&str>,
) -> Result<Option<Response>, ApiError> {
    if !path.iter().any(|segment| segment.starts_with('-')) {
        return Ok(None);
    }
    let positive = {
        let replica = state.doc.read().await;
        // Counted from the end of the list as it was then, for a read at heads.
        read_at!(&replica.doc, at, |doc| store::positive_path(
            doc, path, insertion
        ))
    }
    .map_err(|e| store_status(e, "Failed to resolve a negative index"))?;
    let Some(positive) = positive else {
        return Ok(None);
    };
    let mut location = format!(
        "/store/{}",
        positive
            .iter()
            .map(|segment| encode_segment(segment))
            .collect::<Vec<_>>()
            .join("/")
    );
    if let Some(query) = query.filter(|query| !query.is_empty()) {
        location.push('?');
        location.push_str(query);
    }
    let location = HeaderValue::from_str(&location).map_err(|_| StatusCode::BAD_REQUEST)?;
    Ok(Some(
        (StatusCode::TEMPORARY_REDIRECT, [(LOCATION, location)]).into_response(),
    ))
}

/// `segment` as a URL path segment: unreserved characters as they are, the rest
/// percent-encoded (UTF-8).
fn encode_segment(segment: &str) -> String {
    let mut encoded = String::with_capacity(segment.len());
    for byte in segment.bytes() {
        if byte.is_ascii_alphanumeric() || b"-._~".contains(&byte) {
            encoded.push(char::from(byte));
        } else {
            encoded.push_str(&format!("%{byte:02X}"));
        }
    }
    encoded
}

// --- streams -----------------------------------------------------------------

async fn render_traced(
    state: &AppState,
    path: Vec<String>,
) -> Result<(Option<JsonValue>, bool), StoreError> {
    state
        .render(
            move |doc| store::read_path_traced(doc, &path),
            |result| result,
        )
        .await
}

/// `GET /stream`: follows the whole store, as `GET /store` reads it (events have an
/// empty `key`).
pub(crate) async fn stream_store_root(State(state): State<AppState>) -> Result<Response, ApiError> {
    stream(state, String::new(), Vec::new()).await
}

/// `GET /stream/{key}`: follows one value.
pub(crate) async fn stream_store_entry(
    State(state): State<AppState>,
    Path(key): Path<String>,
) -> Result<Response, ApiError> {
    let path = parse_key_path(&key)?;
    stream(state, key, path).await
}

/// Newline-delimited JSON events for the value at `path` (`key` as the client wrote it):
/// the current value, then one whenever the resolved value changes.
async fn stream(state: AppState, key: String, path: Vec<String>) -> Result<Response, ApiError> {
    if !state.is_ready().await {
        return Err(StatusCode::SERVICE_UNAVAILABLE.into());
    }

    // Subscribe first, so nothing between the initial read and the loop is missed.
    let mut changes = state.change_events.subscribe();

    let initial = render_traced(&state, path.clone())
        .await
        .map_err(|e| store_status(e, "Failed to read initial stream value"))?;

    let stream_state = state.clone();
    let stream_key = key.clone();

    let output = async_stream::stream! {
        let (mut last_value, mut conflicted) = initial;

        match stream_event_line(&stream_key, last_value.as_ref()) {
            Ok(line) => yield Ok::<Bytes, IoError>(line),
            Err(error) => {
                yield Err(error);
                return;
            }
        }

        loop {
            let relevant = match changes.recv().await {
                // Conflicted paths can change without a patch on their path; re-read always.
                Ok(set) => conflicted || set.paths.iter().any(|changed| paths_related(changed, &path)),
                Err(broadcast::error::RecvError::Lagged(skipped)) => {
                    debug!(skipped, key = %stream_key, "Stream lagged; re-reading");
                    true
                }
                Err(broadcast::error::RecvError::Closed) => break,
            };

            if !relevant {
                continue;
            }

            let (next_value, next_conflicted) = match render_traced(&stream_state, path.clone()).await {
                Ok(result) => result,
                Err(error) => {
                    error!(error = %format!("{error:#}"), key = %stream_key, "Failed to read updated stream value");
                    yield Err(IoError::other("Failed to read updated store value"));
                    return;
                }
            };
            conflicted = next_conflicted;

            // A related change can still leave the resolved value unchanged.
            if next_value == last_value {
                continue;
            }

            match stream_event_line(&stream_key, next_value.as_ref()) {
                Ok(line) => yield Ok(line),
                Err(error) => {
                    yield Err(error);
                    return;
                }
            }
            last_value = next_value;
        }
    };

    let mut response = Body::from_stream(output).into_response();
    response.headers_mut().insert(
        CONTENT_TYPE,
        HeaderValue::from_static("application/x-ndjson"),
    );
    response.headers_mut().insert(
        axum::http::header::CACHE_CONTROL,
        HeaderValue::from_static("no-cache, no-transform"),
    );
    Ok(response)
}

fn stream_event_line(key: &str, value: Option<&JsonValue>) -> std::io::Result<Bytes> {
    let mut event = serde_json::Map::new();

    event.insert("key".to_owned(), JsonValue::String(key.to_owned()));

    match value {
        Some(value) => {
            event.insert("value".to_owned(), value.clone());
        }
        None => {
            event.insert("deleted".to_owned(), JsonValue::Bool(true));
        }
    }

    let mut line = serde_json::to_vec(&JsonValue::Object(event)).map_err(IoError::other)?;

    line.push(b'\n');
    Ok(Bytes::from(line))
}