use std::collections::{HashMap, 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::IntoStatic;
use jacquard_common::types::did::Did;
use jacquard_common::types::ident::AtIdentifier;
use jacquard_common::types::nsid::Nsid;
use jacquard_common::types::string::AtUri;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value, json};
use tower::ServiceExt;

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

pub const TYPE_COUNT: &str = "sh.tangled.query.enrichResponse#count";
pub const TYPE_DISTINCT_AUTHORS: &str = "sh.tangled.query.enrichResponse#distinctAuthors";
pub const TYPE_VIEWER: &str = "sh.tangled.query.enrichResponse#viewer";
/// vendored from upstream com.bad-example.identity.resolveMiniDoc's output
pub const TYPE_MINIDOC: &str = "com.bad-example.identity.miniDoc";

const KNOWN_TYPES: [&str; 4] = [TYPE_COUNT, TYPE_DISTINCT_AUTHORS, TYPE_VIEWER, TYPE_MINIDOC];

/// a payload type nsid, with an optional #fragment for lexicon defs. the raw
/// string is kept because it echoes into the data map as the payload key
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct PayloadType {
    raw: String,
    nsid: Nsid<DefaultStr>,
    #[allow(dead_code)]
    fragment: Option<String>,
}

impl PayloadType {
    pub fn parse(raw: &str) -> Result<Self, String> {
        let (nsid, fragment) = match raw.split_once('#') {
            Some((nsid, fragment)) => {
                let valid = !fragment.is_empty()
                    && fragment.starts_with(|c: char| c.is_ascii_alphabetic())
                    && fragment.chars().all(|c| c.is_ascii_alphanumeric());
                if !valid {
                    return Err(format!("invalid fragment #{fragment}"));
                }
                (nsid, Some(fragment.to_owned()))
            }
            None => (raw, None),
        };
        let nsid = Nsid::new_owned(nsid).map_err(|e| format!("invalid nsid: {e}"))?;
        Ok(Self {
            raw: raw.to_owned(),
            nsid,
            fragment,
        })
    }

    fn as_str(&self) -> &str {
        &self.raw
    }
}

impl<'de> Deserialize<'de> for PayloadType {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let raw = String::deserialize(deserializer)?;
        Self::parse(&raw).map_err(serde::de::Error::custom)
    }
}

impl Serialize for PayloadType {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_str(&self.raw)
    }
}

/// the paths the edge index can back: a record's reference field, or its
/// authoring repo
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum SourcePath {
    Subject,
    Repo,
}

/// a "collection:path" link source. raw echoes into the data map verbatim
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct LinkSource {
    raw: String,
    collection: Nsid<DefaultStr>,
    path: SourcePath,
}

impl LinkSource {
    pub fn parse(raw: &str) -> Result<Self, String> {
        let (collection, path) = match raw.split_once(':') {
            Some((collection, path)) if !collection.is_empty() && !path.is_empty() => {
                (collection, path)
            }
            _ => return Err("expected \"collection:path\"".to_owned()),
        };
        let collection =
            Nsid::new_owned(collection).map_err(|e| format!("invalid collection: {e}"))?;
        let path = match path {
            "subject" => SourcePath::Subject,
            ".repo" => SourcePath::Repo,
            _ if path.starts_with('.') => {
                return Err("envelope field not index-backed; try .repo".to_owned());
            }
            _ => {
                return Err(
                    "path not index-backed; only `subject` and `.repo` are supported".to_owned(),
                );
            }
        };
        Ok(Self {
            raw: raw.to_owned(),
            collection,
            path,
        })
    }

    fn as_str(&self) -> &str {
        &self.raw
    }
}

impl<'de> Deserialize<'de> for LinkSource {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let raw = String::deserialize(deserializer)?;
        Self::parse(&raw).map_err(serde::de::Error::custom)
    }
}

impl Serialize for LinkSource {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_str(&self.raw)
    }
}

#[derive(Debug, Deserialize)]
pub struct LinkDescriptor {
    /// constellation link source
    source: LinkSource,
    /// nsid of the sidecar payload to produce for this source
    #[serde(rename = "type")]
    ty: PayloadType,
    /// record paths into the inner response selecting targets for this payload
    #[serde(default)]
    targets: Option<Vec<String>>,
}

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

pub async fn enrich(
    State(state): State<AppState>,
    Json(input): Json<EnrichInput>,
) -> Result<Json<Value>, XrpcError> {
    let mut seen_descriptors = HashSet::new();
    let mut descriptors = Vec::new();
    for (descriptor_index, descriptor) in input.enrich.iter().enumerate() {
        if !KNOWN_TYPES.contains(&descriptor.ty.as_str()) {
            return Err(descriptor_error(&descriptor.source, "unknown enrich type"));
        }
        validate_source(&descriptor.source, &descriptor.ty)?;
        if descriptor.ty.as_str() == TYPE_VIEWER && input.viewer.is_none() {
            return Err(descriptor_error(
                &descriptor.source,
                "viewer payloads require a viewer param",
            ));
        }
        let targets = descriptor
            .targets
            .as_ref()
            .map(|targets| {
                targets
                    .iter()
                    .enumerate()
                    .map(|(target_index, path)| {
                        RecordPath::parse(path).map_err(|e| {
                            XrpcError::InvalidParams(format!(
                                "enrich[{descriptor_index}].targets[{target_index}] {path:?}: {e}"
                            ))
                        })
                    })
                    .collect::<Result<Vec<_>, _>>()
            })
            .transpose()?;
        if seen_descriptors.insert((
            &descriptor.source,
            &descriptor.ty,
            descriptor.targets.as_deref(),
        )) {
            descriptors.push((descriptor, targets));
        }
    }

    let inner = run_inner(&state, &input.xrpc, input.params.unwrap_or_default()).await?;
    if descriptors.is_empty() {
        return Ok(Json(json!({ "output": inner, "data": {} })));
    }

    let mut data = Map::new();
    let mut refs_by_targets: HashMap<Option<&[String]>, Vec<SubjectRef>> = HashMap::new();
    let mut minidoc_targets: HashMap<Did<DefaultStr>, Vec<LinkSource>> = HashMap::new();
    for (descriptor, targets) in descriptors {
        let refs = refs_by_targets
            .entry(descriptor.targets.as_deref())
            .or_insert_with(|| {
                let mut refs: Vec<SubjectRef> = Vec::new();
                let mut seen: HashSet<SubjectRef> = HashSet::new();
                match &targets {
                    Some(targets) => {
                        for path in targets {
                            for node in walk_path(path, [&inner]) {
                                collect_ref(node, &mut refs, &mut seen);
                            }
                        }
                    }
                    None => discover_refs(&inner, &mut refs, &mut seen),
                }
                // the authority of every at-uri is a reference too
                // this lets each enrich address record authors
                let authorities: Vec<SubjectRef> = refs
                    .iter()
                    .filter_map(|reference| repo_did(reference).map(SubjectRef::Did))
                    .collect();
                for authority in authorities {
                    if seen.insert(authority.clone()) {
                        refs.push(authority);
                    }
                }
                refs
            });
        for reference in refs.iter() {
            let Some(subject) = applicable_subject(descriptor, reference)? else {
                continue;
            };
            match descriptor.ty.as_str() {
                TYPE_COUNT | TYPE_DISTINCT_AUTHORS | TYPE_VIEWER => {
                    let kind = edge_kind(&descriptor.source)?;
                    let key = EdgeKey::new(kind, subject.clone());
                    let result = match descriptor.ty.as_str() {
                        TYPE_COUNT => Value::from(state.edges.count(&key)),
                        TYPE_DISTINCT_AUTHORS => {
                            Value::from(state.edges.count_distinct_authors(&key))
                        }
                        _ => {
                            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,
                            }
                        }
                    };
                    put(
                        &mut data,
                        subject.as_str(),
                        descriptor.source.as_str(),
                        descriptor.ty.as_str(),
                        result,
                    );
                }
                TYPE_MINIDOC => {
                    if let Some(target) = repo_did(&subject) {
                        let sources = minidoc_targets.entry(target).or_default();
                        if !sources.contains(&descriptor.source) {
                            sources.push(descriptor.source.clone());
                        }
                    }
                }
                _ => unreachable!("validated above!"),
            }
        }
    }

    let docs = cached_minidocs(&state, minidoc_targets);
    for (target, sources, doc) in docs {
        for source in sources {
            put(
                &mut data,
                target.as_str(),
                source.as_str(),
                TYPE_MINIDOC,
                doc.clone(),
            );
        }
    }

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

fn put(data: &mut Map<String, Value>, reference: &str, source: &str, ty: &str, payload: Value) {
    let per_ref = data
        .entry(reference.to_owned())
        .or_insert_with(|| Value::Object(Map::new()));
    let Value::Object(per_ref) = per_ref else {
        return;
    };
    let per_source = per_ref
        .entry(source.to_owned())
        .or_insert_with(|| Value::Object(Map::new()));
    if let Value::Object(per_source) = per_source {
        per_source.insert(ty.to_owned(), payload);
    }
}

fn cached_minidocs(
    state: &AppState,
    targets: HashMap<Did<DefaultStr>, Vec<LinkSource>>,
) -> Vec<(Did<DefaultStr>, Vec<LinkSource>, Value)> {
    targets
        .into_iter()
        .filter_map(|(did, sources)| {
            state
                .identity
                .get_by_did(&did)
                .ok()
                .and_then(|doc| serde_json::to_value(doc).ok())
                .map(|doc| (did, sources, doc))
        })
        .collect()
}

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

fn validate_source(source: &LinkSource, ty: &PayloadType) -> Result<(), XrpcError> {
    match source.path {
        SourcePath::Subject => subject_shape(source.collection.as_ref())
            .map(|_| ())
            .ok_or_else(|| descriptor_error(source, "unknown collection")),
        // minidocs read the uri's authority straight out of the response, only
        // edge payloads need the author index
        SourcePath::Repo if ty.as_str() == TYPE_MINIDOC => Ok(()),
        SourcePath::Repo => mirror_kind(source.collection.as_ref())
            .map(|_| ())
            .ok_or_else(|| descriptor_error(source, "collection has no author index")),
    }
}

fn repo_did(reference: &SubjectRef) -> Option<Did<DefaultStr>> {
    match reference {
        SubjectRef::Did(did) => Some(did.clone()),
        SubjectRef::Uri(uri) => match uri.authority() {
            AtIdentifier::Did(did) => Some(did.clone().into_static()),
            AtIdentifier::Handle(_) => None,
        },
        SubjectRef::Global => None,
    }
}

/// the subject an edge lookup or payload should be keyed by, or None if the
/// descriptor's shape doesn't apply to this ref, eg. a did-only source asked
/// about an at-uri. `.repo` sources normalize to the repo did
fn applicable_subject(
    descriptor: &LinkDescriptor,
    reference: &SubjectRef,
) -> Result<Option<SubjectRef>, XrpcError> {
    match descriptor.source.path {
        SourcePath::Subject => {
            let (_, shape) = subject_shape(descriptor.source.collection.as_ref())
                .ok_or_else(|| descriptor_error(&descriptor.source, "unknown collection"))?;
            Ok(shape_accepts(shape, reference).then(|| reference.clone()))
        }
        SourcePath::Repo => Ok(repo_did(reference).map(SubjectRef::Did)),
    }
}

/// the edge kind behind a source for stats payloads
fn edge_kind(source: &LinkSource) -> Result<Nsid<DefaultStr>, XrpcError> {
    let kind = match source.path {
        SourcePath::Subject => subject_shape(source.collection.as_ref())
            .map(|(nsid, _)| nsid)
            .ok_or_else(|| descriptor_error(source, "unknown collection"))?,
        SourcePath::Repo => mirror_kind(source.collection.as_ref())
            .ok_or_else(|| descriptor_error(source, "collection has no author index"))?,
    };
    Ok(nsid_static(kind))
}

/// 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,
    }
}

/// 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);
    }

    #[test]
    fn repo_did_normalizes_uri_authorities() {
        let uri = SubjectRef::Uri(AtUri::new_owned("at://did:plc:abc/sh.tangled.repo/x").unwrap());
        let bare = SubjectRef::Did(Did::new_owned("did:plc:abc").unwrap());
        assert_eq!(
            repo_did(&uri).as_ref().map(|d| d.as_str()),
            Some("did:plc:abc")
        );
        assert_eq!(
            repo_did(&bare).as_ref().map(|d| d.as_str()),
            Some("did:plc:abc")
        );
        assert_eq!(repo_did(&SubjectRef::Global), None);
    }

    #[test]
    fn parses_link_sources_and_payload_types() {
        let source = LinkSource::parse("sh.tangled.feed.star:subject").unwrap();
        assert_eq!(source.collection.as_ref(), "sh.tangled.feed.star");
        assert_eq!(source.path, SourcePath::Subject);
        let source = LinkSource::parse("sh.tangled.graph.follow:.repo").unwrap();
        assert_eq!(source.path, SourcePath::Repo);
        assert!(LinkSource::parse("sh.tangled.feed.star").is_err());
        assert!(LinkSource::parse("sh.tangled.feed.star:.rkey").is_err());
        assert!(LinkSource::parse("not an nsid:subject").is_err());

        let ty = PayloadType::parse(TYPE_MINIDOC).unwrap();
        assert_eq!(ty.nsid.as_ref(), "com.bad-example.identity.miniDoc");
        assert_eq!(ty.fragment.as_deref(), None);
        let ty = PayloadType::parse(TYPE_COUNT).unwrap();
        assert_eq!(ty.nsid.as_ref(), "sh.tangled.query.enrichResponse");
        assert_eq!(ty.fragment.as_deref(), Some("count"));
        assert!(PayloadType::parse("sh.tangled.query.enrichResponse#").is_err());
        assert!(PayloadType::parse("nope#count").is_err());
    }
}