//! Shared machinery for reading the current remote state of a package.
//!
//! Both `mlf status` and `mlf diff` (and later `mlf publish`) need to:
//!
//! 1. Load the workspace and generate Lexicon JSON for every local module.
//! 2. Compute the on-chain CID for each local record.
//! 3. Resolve the package's publishing DID via `_lexicon.<authority>` TXT.
//! 4. Fetch the current `com.atproto.lexicon.schema` records from the PDS.
//! 5. Diff local vs remote by (NSID, CID).
//!
//! This module exposes the first three as a single read-only operation
//! that produces a [`RemoteState`] the command-specific code can render.

use crate::check::collect_mlf_files_pub as collect_mlf_files;
use crate::check::extract_namespace_pub as extract_namespace;
use crate::config::{MlfConfig, PackageConfig, find_project_root};
use crate::workspace_ext::workspace_with_std_and_cache;
use miette::Diagnostic;
use mlf_atproto::identity::{self, IdentityError, RealDnsResolver};
use mlf_atproto::records::{self, RecordError};
use mlf_lang::Workspace;
use serde_json::Value;
use std::collections::{BTreeMap, HashSet};
use std::path::{Path, PathBuf};
use thiserror::Error;

/// Collection name the AT Protocol uses for lexicon schemas.
pub const LEXICON_COLLECTION: &str = "com.atproto.lexicon.schema";

#[derive(Error, Debug, Diagnostic)]
pub enum RemoteStateError {
    #[error("No mlf.toml found in current or parent directories")]
    #[diagnostic(
        code(mlf::remote_state::no_project),
        help("Run `mlf init` first to set up a project.")
    )]
    NoProject,

    #[error("Failed to load mlf.toml: {0}")]
    #[diagnostic(code(mlf::remote_state::config))]
    Config(String),

    #[error("Failed to read source files: {0}")]
    #[diagnostic(code(mlf::remote_state::io))]
    Io(String),

    #[error("Failed to parse {file}")]
    #[diagnostic(code(mlf::remote_state::parse))]
    Parse { file: String },

    #[error("Workspace resolution failed")]
    #[diagnostic(code(mlf::remote_state::resolve))]
    Resolve,

    #[error("Lexicon `{namespace}` (in {file}) is outside package scope `{package_name}.*`")]
    #[diagnostic(
        code(mlf::remote_state::scope_violation),
        help("Move the file or update `[package].name` in mlf.toml.")
    )]
    ScopeViolation {
        file: String,
        namespace: String,
        package_name: String,
    },

    #[error("Failed to generate Lexicon JSON for `{namespace}`: {reason}")]
    #[diagnostic(code(mlf::remote_state::codegen))]
    Codegen { namespace: String, reason: String },

    #[error("CID computation failed for `{namespace}`: {reason}")]
    #[diagnostic(code(mlf::remote_state::cid))]
    Cid { namespace: String, reason: String },

    #[error("Identity resolution failed")]
    #[diagnostic(code(mlf::remote_state::identity))]
    Identity(#[source] IdentityError),

    #[error("Record fetch failed")]
    #[diagnostic(code(mlf::remote_state::records))]
    Records(#[source] RecordError),

    #[error(
        "Package authorities resolve to multiple DIDs: {first_authority} → {first_did}, {second_authority} → {second_did}"
    )]
    #[diagnostic(
        code(mlf::remote_state::multiple_dids),
        help(
            "Every authority under `[package].name` must resolve to the same DID. This package spans two. Either fix the `_lexicon` TXT records, or split into separate workspaces."
        )
    )]
    MultipleDids {
        first_authority: String,
        first_did: String,
        second_authority: String,
        second_did: String,
    },
}

/// Fully-resolved view of the package's local state and the remote state
/// it would compare against on publish.
pub struct RemoteState {
    pub project_root: PathBuf,
    pub package: PackageConfig,

    /// The DID every authority under `[package].name` resolves to.
    /// `None` if we couldn't resolve any authority (e.g. no TXT records yet
    /// — a brand-new package).
    pub publishing_did: Option<String>,

    /// The PDS endpoint for [`publishing_did`]. `None` when the DID is.
    pub pds: Option<String>,

    /// Per-authority TXT lookup result. Keyed by authority (the reverse-
    /// DNS prefix of the NSID), so `[package.foo, package.bar]` under
    /// `com.example` both show up as `com.example` in the map — the
    /// authority covers every NSID that starts with it.
    pub authority_status: BTreeMap<String, AuthorityStatus>,

    /// Every lexicon defined locally, keyed by NSID.
    pub local: BTreeMap<String, LocalLexicon>,

    /// Every `com.atproto.lexicon.schema` record currently in the
    /// publishing DID's repo, keyed by NSID. Empty if we couldn't
    /// resolve a DID.
    pub remote: BTreeMap<String, RemoteLexicon>,
}

pub struct LocalLexicon {
    pub namespace: String,
    pub file: String,
    /// The JSON as it would be written to the PDS (with `$type`).
    pub record_json: Value,
    /// CID of [`record_json`] — comparable to whatever listRecords returns.
    pub cid: String,
}

pub struct RemoteLexicon {
    pub nsid: String,
    pub cid: String,
    pub record_json: Value,
}

#[derive(Debug, Clone)]
pub enum AuthorityStatus {
    /// `_lexicon.<auth>` TXT resolved to a DID.
    Resolved { did: String },
    /// `_lexicon.<auth>` TXT was missing or empty.
    Missing,
    /// Lookup failed for some other reason (network, parse, etc.).
    Error(String),
}

impl RemoteState {
    /// Build a [`RemoteState`] by walking the workspace, generating JSON,
    /// and performing the necessary network lookups. Uses CWD to find
    /// the project root. For loading a specific package in a multi-
    /// package workspace, use [`RemoteState::load_at`].
    pub async fn load() -> Result<Self, RemoteStateError> {
        let current_dir =
            std::env::current_dir().map_err(|e| RemoteStateError::Io(e.to_string()))?;
        let project_root =
            find_project_root(&current_dir).map_err(|_| RemoteStateError::NoProject)?;
        let config_path = project_root.join("mlf.toml");
        let config =
            MlfConfig::load(&config_path).map_err(|e| RemoteStateError::Config(e.to_string()))?;
        RemoteState::load_at(&project_root, &config).await
    }

    /// Build a [`RemoteState`] for a specific package rooted at
    /// `project_root`. The caller has already loaded `config`.
    pub async fn load_at(
        project_root: &Path,
        config: &MlfConfig,
    ) -> Result<Self, RemoteStateError> {
        let project_root = project_root.to_path_buf();
        let package = config.package.clone();
        let source_dir = project_root.join(&config.source.directory);

        // 1. Load every .mlf file, parse it, add to a resolver workspace.
        let local = collect_local(&source_dir, &project_root, &package)?;

        // 2. Group NSIDs into their reverse-DNS authorities.
        let authorities = authorities_for_nsids(local.keys().map(|s| s.as_str()))?;

        // 3. Resolve each authority via DNS; enforce single-DID gate.
        let resolver = RealDnsResolver::new().map_err(RemoteStateError::Identity)?;
        let (publishing_did, authority_status) =
            resolve_authorities(&resolver, &authorities).await?;

        // 4. If we have a DID, resolve its PDS and fetch the lexicon records.
        let (pds, remote) = if let Some(did) = publishing_did.as_deref() {
            let client = reqwest::Client::new();
            let pds = identity::resolve_did_to_pds(&client, did)
                .await
                .map_err(RemoteStateError::Identity)?;
            let records = records::list_all_records(&client, &pds, did, LEXICON_COLLECTION)
                .await
                .map_err(RemoteStateError::Records)?;
            let mut remote_map: BTreeMap<String, RemoteLexicon> = BTreeMap::new();
            for r in records {
                let nsid = record_nsid(&r).unwrap_or_else(|| rkey_from_uri(&r.uri));
                let cid = r.cid.clone().unwrap_or_default();
                remote_map.insert(
                    nsid.clone(),
                    RemoteLexicon {
                        nsid,
                        cid,
                        record_json: r.value,
                    },
                );
            }
            (Some(pds), remote_map)
        } else {
            (None, BTreeMap::new())
        };

        Ok(Self {
            project_root,
            package,
            publishing_did,
            pds,
            authority_status,
            local,
            remote,
        })
    }
}

/// Compute what a publish run would do given a fully-loaded [`RemoteState`].
pub fn classify(state: &RemoteState) -> BTreeMap<String, Classification> {
    let mut out = BTreeMap::new();
    let mut seen: HashSet<&String> = HashSet::new();

    for (nsid, local) in &state.local {
        seen.insert(nsid);
        let cls = match state.remote.get(nsid) {
            None => Classification::New,
            Some(remote) if remote.cid == local.cid => Classification::Unchanged,
            Some(remote) => Classification::Changed {
                local_cid: local.cid.clone(),
                remote_cid: remote.cid.clone(),
            },
        };
        out.insert(nsid.clone(), cls);
    }

    for nsid in state.remote.keys() {
        if state.package.namespace_is_in_scope(nsid) && !seen.contains(nsid) {
            // We only flag remote-only records that fall under this package's
            // scope. Anything else in the repo belongs to a different package
            // hosted on the same account and isn't ours to touch.
            out.insert(nsid.clone(), Classification::Removed);
        }
    }

    out
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Classification {
    Unchanged,
    New,
    Changed {
        local_cid: String,
        remote_cid: String,
    },
    Removed,
}

// ---------------------------------------------------------------------------
// helpers
// ---------------------------------------------------------------------------

fn collect_local(
    source_dir: &Path,
    project_root: &Path,
    package: &PackageConfig,
) -> Result<BTreeMap<String, LocalLexicon>, RemoteStateError> {
    let cache_dir = crate::config::get_mlf_cache_dir(project_root);
    let mut workspace = workspace_with_std_and_cache(Some(&cache_dir))
        .map_err(|e| RemoteStateError::Io(format!("loading workspace: {e}")))?;

    let files = collect_mlf_files(source_dir)
        .map_err(|e| RemoteStateError::Io(format!("collecting .mlf files: {e}")))?;

    let mut source_order: Vec<(String, String)> = Vec::new();

    for file in &files {
        let source = std::fs::read_to_string(file)
            .map_err(|e| RemoteStateError::Io(format!("reading {}: {e}", file.display())))?;
        let namespace = extract_namespace(file, source_dir)
            .map_err(|e| RemoteStateError::Io(format!("namespace for {}: {e}", file.display())))?;
        if !package.namespace_is_in_scope(&namespace) {
            return Err(RemoteStateError::ScopeViolation {
                file: file.display().to_string(),
                namespace,
                package_name: package.name.clone(),
            });
        }
        let lexicon = mlf_lang::parse_lexicon(&source).map_err(|_| RemoteStateError::Parse {
            file: file.display().to_string(),
        })?;
        workspace
            .add_module(namespace.clone(), lexicon)
            .map_err(|_| RemoteStateError::Resolve)?;
        source_order.push((namespace, file.display().to_string()));
    }

    workspace.resolve().map_err(|_| RemoteStateError::Resolve)?;

    // Produce JSON + CID for each module in the order we loaded them.
    let mut out = BTreeMap::new();
    for (namespace, file) in source_order {
        let lexicon = workspace_module(&workspace, &namespace)?;
        let codegen_out = mlf_codegen::generate_lexicon(&namespace, lexicon, &workspace);
        let record_json = wrap_as_schema_record(codegen_out.json);
        let cid =
            mlf_atproto::cid::cid_for_json(&record_json).map_err(|e| RemoteStateError::Cid {
                namespace: namespace.clone(),
                reason: e.to_string(),
            })?;
        out.insert(
            namespace.clone(),
            LocalLexicon {
                namespace: namespace.clone(),
                file,
                record_json,
                cid,
            },
        );
    }

    Ok(out)
}

fn workspace_module<'a>(
    workspace: &'a Workspace,
    namespace: &str,
) -> Result<&'a mlf_lang::Lexicon, RemoteStateError> {
    workspace
        .get_lexicon(namespace)
        .ok_or_else(|| RemoteStateError::Codegen {
            namespace: namespace.to_string(),
            reason: "module not present in workspace after resolve".to_string(),
        })
}

/// Wrap a generated Lexicon JSON with the `$type` required on a PDS record.
fn wrap_as_schema_record(mut lexicon_json: Value) -> Value {
    if let Value::Object(ref mut map) = lexicon_json {
        map.insert(
            "$type".to_string(),
            Value::String(LEXICON_COLLECTION.to_string()),
        );
    }
    lexicon_json
}

/// Take every NSID and return the unique set of authorities covering
/// them per the ATProto lexicon spec: drop the final name segment and
/// use the remainder as the authority.
///
/// NSIDs `com.example.forum.post` and `com.example.forum.thread` share
/// authority `com.example.forum`. `com.example.directory.listing` adds
/// a second authority `com.example.directory`.
fn authorities_for_nsids<'a>(
    nsids: impl Iterator<Item = &'a str>,
) -> Result<HashSet<String>, RemoteStateError> {
    nsids
        .map(|nsid| identity::authority_of_nsid(nsid).map_err(RemoteStateError::Identity))
        .collect()
}

async fn resolve_authorities(
    resolver: &RealDnsResolver,
    authorities: &HashSet<String>,
) -> Result<(Option<String>, BTreeMap<String, AuthorityStatus>), RemoteStateError> {
    let mut statuses: BTreeMap<String, AuthorityStatus> = BTreeMap::new();
    let mut dids: Vec<(String, String)> = Vec::new();

    for authority in authorities {
        let label = identity::authority_dns_name(authority);
        let res = resolver.resolve_authority_did(authority).await;
        match res {
            Ok(did) => {
                statuses.insert(
                    label.clone(),
                    AuthorityStatus::Resolved { did: did.clone() },
                );
                dids.push((label, did));
            }
            Err(IdentityError::NoDidInTxt(_)) => {
                statuses.insert(label, AuthorityStatus::Missing);
            }
            Err(IdentityError::DnsLookupFailed { error, .. }) => {
                statuses.insert(label, AuthorityStatus::Error(error));
            }
            Err(e) => return Err(RemoteStateError::Identity(e)),
        }
    }

    // Enforce single-DID gate: every resolved authority must point at the
    // same DID. Multiple DIDs is a v1 error per the plan.
    let mut publishing_did: Option<String> = None;
    if let Some((first_auth, first_did)) = dids.first().cloned() {
        for (auth, did) in dids.iter().skip(1) {
            if did != &first_did {
                return Err(RemoteStateError::MultipleDids {
                    first_authority: first_auth,
                    first_did,
                    second_authority: auth.clone(),
                    second_did: did.clone(),
                });
            }
        }
        publishing_did = Some(first_did);
    }

    Ok((publishing_did, statuses))
}

fn record_nsid(record: &records::Record) -> Option<String> {
    record
        .value
        .get("id")
        .and_then(|v| v.as_str())
        .map(|s| s.to_string())
}

fn rkey_from_uri(uri: &str) -> String {
    uri.rsplit('/').next().unwrap_or(uri).to_string()
}

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

    fn mk_local(nsid: &str, cid: &str) -> LocalLexicon {
        LocalLexicon {
            namespace: nsid.to_string(),
            file: format!("{nsid}.mlf"),
            record_json: json!({"id": nsid}),
            cid: cid.to_string(),
        }
    }

    fn mk_remote(nsid: &str, cid: &str) -> RemoteLexicon {
        RemoteLexicon {
            nsid: nsid.to_string(),
            cid: cid.to_string(),
            record_json: json!({"id": nsid}),
        }
    }

    fn state_with(
        package: &str,
        local: Vec<LocalLexicon>,
        remote: Vec<RemoteLexicon>,
    ) -> RemoteState {
        let mut local_map = BTreeMap::new();
        for l in local {
            local_map.insert(l.namespace.clone(), l);
        }
        let mut remote_map = BTreeMap::new();
        for r in remote {
            remote_map.insert(r.nsid.clone(), r);
        }
        RemoteState {
            project_root: PathBuf::from("/tmp"),
            package: PackageConfig {
                name: package.to_string(),
            },
            publishing_did: Some("did:plc:x".to_string()),
            pds: Some("https://pds".to_string()),
            authority_status: BTreeMap::new(),
            local: local_map,
            remote: remote_map,
        }
    }

    #[test]
    fn classify_detects_new_changed_unchanged_removed() {
        let state = state_with(
            "com.example.forum",
            vec![
                mk_local("com.example.forum.post", "bafyA"), // unchanged
                mk_local("com.example.forum.thread", "bafyNEW"), // changed
                mk_local("com.example.forum.reply", "bafyR"), // new
            ],
            vec![
                mk_remote("com.example.forum.post", "bafyA"),
                mk_remote("com.example.forum.thread", "bafyOLD"),
                mk_remote("com.example.forum.orphan", "bafyO"), // removed
            ],
        );

        let classes = classify(&state);

        assert_eq!(
            classes.get("com.example.forum.post"),
            Some(&Classification::Unchanged)
        );
        assert_eq!(
            classes.get("com.example.forum.thread"),
            Some(&Classification::Changed {
                local_cid: "bafyNEW".into(),
                remote_cid: "bafyOLD".into(),
            }),
        );
        assert_eq!(
            classes.get("com.example.forum.reply"),
            Some(&Classification::New)
        );
        assert_eq!(
            classes.get("com.example.forum.orphan"),
            Some(&Classification::Removed)
        );
    }

    #[test]
    fn classify_ignores_remote_records_outside_package_scope() {
        // The PDS may hold records for other packages under the same
        // account. We must never touch records whose NSID isn't under
        // `[package].name`.
        let state = state_with(
            "com.example.forum",
            vec![mk_local("com.example.forum.post", "bafyA")],
            vec![
                mk_remote("com.example.forum.post", "bafyA"),
                // Under the same account but a different package — leave it alone.
                mk_remote("com.example.directory.listing", "bafyD"),
            ],
        );

        let classes = classify(&state);
        assert_eq!(classes.len(), 1);
        assert!(classes.contains_key("com.example.forum.post"));
        assert!(!classes.contains_key("com.example.directory.listing"));
    }

    #[test]
    fn authorities_collapse_by_full_prefix() {
        let auths = authorities_for_nsids(
            ["com.example.forum.post", "com.example.forum.thread"].into_iter(),
        )
        .unwrap();
        assert_eq!(auths.len(), 1);
        assert!(auths.contains("com.example.forum"));
    }

    #[test]
    fn authorities_split_on_different_sub_authority() {
        let auths = authorities_for_nsids(
            ["com.example.forum.post", "com.example.directory.listing"].into_iter(),
        )
        .unwrap();
        assert_eq!(auths.len(), 2);
        assert!(auths.contains("com.example.forum"));
        assert!(auths.contains("com.example.directory"));
    }
}