use std::collections::HashSet;

use axum::{
    Json,
    body::{Body, to_bytes},
    extract::State,
    http::{Request, StatusCode},
    response::Response,
};
use bobbin_types::ids::{EdgeKey, SubjectRef, nsid_static};
use jacquard_common::DefaultStr;
use jacquard_common::types::did::Did;
use jacquard_common::types::ident::AtIdentifier;
use jacquard_common::types::string::AtUri;
use serde::Deserialize;
use serde_json::{Map, Value, json};
use tower::ServiceExt;

use crate::recordpath::{parse_record_path, walk_path};
use crate::{AppState, SubjectShape, XrpcError, mirror_kind, subject_shape};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum Aggregation {
    Count,
    DistinctAuthors,
    Viewer,
}

impl Aggregation {
    fn key(self) -> &'static str {
        match self {
            Aggregation::Count => "count",
            Aggregation::DistinctAuthors => "distinctAuthors",
            Aggregation::Viewer => "viewer",
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct LinkDescriptor {
    /// constellation link source
    source: String,
    #[serde(rename = "type")]
    aggregation: Aggregation,
}

impl LinkDescriptor {
    fn parts(&self) -> Result<(&str, &str, Aggregation), XrpcError> {
        match self.source.split_once(':') {
            Some((collection, path)) if !collection.is_empty() && !path.is_empty() => {
                Ok((collection, path, self.aggregation))
            }
            _ => Err(XrpcError::InvalidParams(format!(
                "enrich {:?}: expected \"collection:path\"",
                self.source
            ))),
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct EnrichInput {
    xrpc: String,
    #[serde(default)]
    params: Option<Map<String, Value>>,
    enrich: Vec<LinkDescriptor>,
    #[serde(default)]
    sources: Option<Vec<String>>,
    #[serde(default)]
    viewer: Option<Did>,
}

pub async fn enrich(
    State(state): State<AppState>,
    Json(input): Json<EnrichInput>,
) -> Result<Json<Value>, XrpcError> {
    for descriptor in &input.enrich {
        validate_descriptor(descriptor)?;
        if descriptor.aggregation == Aggregation::Viewer && input.viewer.is_none() {
            return Err(descriptor_error(
                descriptor,
                "viewer aggregation requires a viewer param",
            ));
        }
    }

    let inner = run_inner(&state, &input.xrpc, input.params.unwrap_or_default()).await?;

    let mut refs: Vec<SubjectRef> = Vec::new();
    let mut seen: HashSet<SubjectRef> = HashSet::new();
    match &input.sources {
        Some(sources) => {
            for (i, path) in sources.iter().enumerate() {
                let segs = parse_record_path(path)
                    .map_err(|e| XrpcError::InvalidParams(format!("sources[{i}] {path:?}: {e}")))?;
                for node in walk_path(&segs, [&inner]) {
                    collect_ref(node, &mut refs, &mut seen);
                }
            }
        }
        None => discover_refs(&inner, &mut refs, &mut seen),
    }

    let mut stats = Map::new();
    for reference in &refs {
        let mut per_ref = Map::new();
        for descriptor in &input.enrich {
            let (.., aggregation) = descriptor.parts()?;
            let Some(kind) = resolve_descriptor(descriptor, reference)? else {
                continue;
            };
            let key = EdgeKey::new(kind, reference.clone());
            let result = match aggregation {
                Aggregation::Count => Value::from(state.edges.count(&key)),
                Aggregation::DistinctAuthors => {
                    Value::from(state.edges.count_distinct_authors(&key))
                }
                Aggregation::Viewer => {
                    // viewer presence is validated up front
                    let viewer = input.viewer.as_ref().expect("viewer param present");
                    match state.edges.viewer_source(&key, viewer.as_str()) {
                        Some(uri) => Value::String(uri.to_string()),
                        None => Value::Null,
                    }
                }
            };
            let entry = per_ref
                .entry(descriptor.source.clone())
                .or_insert_with(|| Value::Object(Map::new()));
            if let Value::Object(counts) = entry {
                counts.insert(aggregation.key().to_owned(), result);
            }
        }
        if !per_ref.is_empty() {
            stats.insert(subject_string(reference), Value::Object(per_ref));
        }
    }

    Ok(Json(json!({ "output": inner, "stats": stats })))
}

fn descriptor_error(descriptor: &LinkDescriptor, msg: &str) -> XrpcError {
    XrpcError::InvalidParams(format!("enrich {}: {msg}", descriptor.source))
}

fn validate_descriptor(descriptor: &LinkDescriptor) -> Result<(), XrpcError> {
    let (collection, path, _) = descriptor.parts()?;
    match path {
        "subject" => subject_shape(collection)
            .map(|_| ())
            .ok_or_else(|| descriptor_error(descriptor, "unknown collection")),
        ".repo" => mirror_kind(collection)
            .map(|_| ())
            .ok_or_else(|| descriptor_error(descriptor, "collection has no author index")),
        _ if path.starts_with('.') => Err(descriptor_error(
            descriptor,
            "envelope field not index-backed; try .repo",
        )),
        _ => Err(descriptor_error(
            descriptor,
            "path not index-backed; only `subject` and `.repo` are supported",
        )),
    }
}

/// the edge kind to look up for one reference, or None if the descriptor's shape
/// doesn't apply to this ref kind, eg. a did-only descriptor asked about an at-uri
fn resolve_descriptor(
    descriptor: &LinkDescriptor,
    reference: &SubjectRef,
) -> Result<Option<jacquard_common::types::nsid::Nsid<DefaultStr>>, XrpcError> {
    let (collection, path, _) = descriptor.parts()?;
    match path {
        "subject" => {
            let Some((nsid, shape)) = subject_shape(collection) else {
                return Err(descriptor_error(descriptor, "unknown collection"));
            };
            Ok(shape_accepts(shape, reference).then(|| nsid_static(nsid)))
        }
        ".repo" => {
            let Some(kind) = mirror_kind(collection) else {
                return Err(descriptor_error(
                    descriptor,
                    "collection has no author index",
                ));
            };
            Ok(matches!(reference, SubjectRef::Did(_)).then(|| nsid_static(kind)))
        }
        _ => unreachable!("validated up front"),
    }
}

/// mismatch means skip, not reject
fn shape_accepts(shape: SubjectShape, reference: &SubjectRef) -> bool {
    match (shape, reference) {
        (SubjectShape::BareDid, SubjectRef::Did(_)) => true,
        (SubjectShape::Collection(expected), SubjectRef::Uri(uri)) => {
            uri.collection().is_some_and(|c| c.as_ref() == expected)
        }
        (SubjectShape::OneOfCollections(allowed), SubjectRef::Uri(uri))
        | (SubjectShape::BareDidOrOneOfCollections(allowed), SubjectRef::Uri(uri)) => uri
            .collection()
            .is_some_and(|c| allowed.contains(&c.as_ref())),
        (SubjectShape::BareDidOrOneOfCollections(_), SubjectRef::Did(_)) => true,
        (SubjectShape::AnyAtUri, SubjectRef::Uri(_)) => true,
        _ => false,
    }
}

fn subject_string(reference: &SubjectRef) -> String {
    match reference {
        SubjectRef::Did(d) => d.as_ref().to_owned(),
        SubjectRef::Uri(u) => u.as_ref().to_owned(),
    }
}

/// a value counts as a reference if it is a did string or an at-uri string
/// with collection and rkey, or a {uri, cid} strong ref
fn collect_ref(value: &Value, refs: &mut Vec<SubjectRef>, seen: &mut HashSet<SubjectRef>) {
    let candidate = match value {
        Value::String(s) => Some(s.as_str()),
        Value::Object(o) => o
            .get("uri")
            .and_then(Value::as_str)
            .filter(|_| o.contains_key("cid")),
        _ => None,
    };
    let Some(candidate) = candidate else { return };
    let reference = if candidate.starts_with("did:") {
        Did::<DefaultStr>::new_owned(candidate)
            .ok()
            .map(SubjectRef::Did)
    } else if candidate.starts_with("at://") {
        AtUri::<DefaultStr>::new_owned(candidate)
            .ok()
            .filter(|u| {
                matches!(u.authority(), AtIdentifier::Did(_))
                    && u.collection().is_some()
                    && u.rkey().is_some()
            })
            .map(SubjectRef::Uri)
    } else {
        None
    };
    if let Some(reference) = reference
        && seen.insert(reference.clone())
    {
        refs.push(reference);
    }
}

fn discover_refs(value: &Value, refs: &mut Vec<SubjectRef>, seen: &mut HashSet<SubjectRef>) {
    collect_ref(value, refs, seen);
    match value {
        Value::Array(items) => {
            for item in items {
                discover_refs(item, refs, seen);
            }
        }
        Value::Object(map) => {
            for v in map.values() {
                discover_refs(v, refs, seen);
            }
        }
        _ => {}
    }
}

async fn run_inner(
    state: &AppState,
    nsid: &str,
    params: Map<String, Value>,
) -> Result<Value, XrpcError> {
    let qs = encode_params(&params);
    let uri = format!("/xrpc/{nsid}?{qs}");
    let request = Request::builder()
        .method("GET")
        .uri(&uri)
        .body(Body::empty())
        .map_err(|e| XrpcError::Internal(format!("inner request: {e}")))?;
    let response = state
        .self_router()
        .oneshot(request)
        .await
        .map_err(|e| XrpcError::Internal(format!("inner dispatch: {e}")))?;
    if response.status() == StatusCode::NOT_FOUND {
        let bytes = to_bytes(response.into_body(), usize::MAX)
            .await
            .map_err(|e| XrpcError::Internal(format!("inner response: {e}")))?;
        return Err(bytes
            .is_empty()
            .then(|| XrpcError::InvalidParams(format!("unknown or unenrichable query: {nsid}")))
            .unwrap_or(XrpcError::NotFound));
    }
    finish(response).await
}

fn encode_params(params: &Map<String, Value>) -> String {
    let mut out = url::form_urlencoded::Serializer::new(String::new());
    for (key, value) in params {
        match value {
            Value::Array(items) => {
                for item in items {
                    if let Some(scalar) = scalar_str(item) {
                        out.append_pair(key, &scalar);
                    }
                }
            }
            _ => {
                if let Some(scalar) = scalar_str(value) {
                    out.append_pair(key, &scalar);
                }
            }
        }
    }
    out.finish()
}

fn scalar_str(value: &Value) -> Option<String> {
    match value {
        Value::String(s) => Some(s.clone()),
        Value::Number(n) => Some(n.to_string()),
        Value::Bool(b) => Some(b.to_string()),
        _ => None,
    }
}

async fn finish(resp: Response) -> Result<Value, XrpcError> {
    let status = resp.status();
    let bytes = to_bytes(resp.into_body(), usize::MAX)
        .await
        .map_err(|e| XrpcError::Internal(format!("inner response: {e}")))?;
    if !status.is_success() {
        let msg = String::from_utf8_lossy(&bytes).into_owned();
        return Err(match status {
            StatusCode::BAD_REQUEST => XrpcError::InvalidParams(msg),
            StatusCode::NOT_FOUND => XrpcError::NotFound,
            StatusCode::SERVICE_UNAVAILABLE => XrpcError::Overloaded,
            _ => XrpcError::UpstreamUnavailable(format!("inner query ({status}): {msg}")),
        });
    }
    serde_json::from_slice(&bytes)
        .map_err(|e| XrpcError::Internal(format!("inner response decode: {e}")))
}

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

    #[test]
    fn collects_dids_uris_and_strong_refs() {
        let mut refs = Vec::new();
        let mut seen = HashSet::new();
        collect_ref(&json!("did:plc:abc"), &mut refs, &mut seen);
        collect_ref(
            &json!("at://did:plc:abc/sh.tangled.repo/x"),
            &mut refs,
            &mut seen,
        );
        collect_ref(
            &json!({"uri": "at://did:plc:abc/sh.tangled.repo/y", "cid": "bafy"}),
            &mut refs,
            &mut seen,
        );
        collect_ref(&json!("at://did:plc:abc"), &mut refs, &mut seen);
        collect_ref(&json!("oppi.li"), &mut refs, &mut seen);
        collect_ref(&json!("did:plc:abc"), &mut refs, &mut seen);
        assert_eq!(refs.len(), 3);
    }
}