use super::LexiconSource;
use jacquard_common::deps::smol_str::SmolStr;
use jacquard_common::types::ident::AtIdentifier;
use jacquard_common::types::string::Nsid;
use jacquard_common::xrpc::XrpcExt;
use jacquard_common::xrpc::atproto::{ListRecords, ListRecordsRecord};
use jacquard_common::{BosStr, IntoStatic};
use jacquard_identity::lexicon_resolver::LexiconSchemaResolver;
use jacquard_identity::resolver::{IdentityResolver, ResolverOptions};
use jacquard_identity::{JacquardResolver, PublicResolver};
use jacquard_lexicon::lexicon::LexiconDoc;
use miette::{Result, miette};
use serde::Serialize;
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub struct AtProtoSource {
    pub endpoint: String,
    pub slice: Option<String>,
}

impl AtProtoSource {
    /// Fetch a single lexicon schema by NSID using DNS + XRPC resolution
    async fn fetch_single_lexicon<S: BosStr + Sync>(
        &self,
        resolver: &PublicResolver,
        nsid: &Nsid<S>,
    ) -> Result<HashMap<String, LexiconDoc<'_>>> {
        let schema = resolver
            .resolve_lexicon_schema(nsid)
            .await
            .map_err(|e| miette!("Failed to resolve lexicon '{}': {}", nsid, e))?;

        let mut lexicons = HashMap::new();
        lexicons.insert(schema.nsid.to_string(), schema.doc);

        Ok(lexicons)
    }

    fn parse_lexicon_record<S: BosStr + Sync + Serialize>(
        record_data: &ListRecordsRecord<S>,
    ) -> Option<LexiconDoc<'static>> {
        // // Extract the 'value' field from the record
        // let value = match record_data {
        //     jacquard_common::types::value::Data::Object(map) => map.0.get("value")?,
        //     _ => {
        //         eprintln!("Warning: Record is not an object");
        //         return None;
        //     }
        // };

        match serde_json::to_string(&record_data.value) {
            Ok(json) => match serde_json::from_str::<LexiconDoc>(&json) {
                Ok(doc) => Some(doc.into_static()),
                Err(e) => {
                    eprintln!("Warning: Failed to parse lexicon from record value: {}", e);
                    None
                }
            },
            Err(e) => {
                eprintln!("Warning: Failed to serialize record value: {}", e);
                None
            }
        }
    }
}

impl LexiconSource for AtProtoSource {
    async fn fetch(&self) -> Result<HashMap<String, LexiconDoc<'_>>> {
        let http = reqwest::Client::new();
        let resolver = JacquardResolver::new_dns(http.clone(), ResolverOptions::default());

        // Try parsing as NSID first (for single lexicon fetch)
        if let Ok(nsid) = Nsid::new(self.endpoint.as_ref()) {
            return self.fetch_single_lexicon(&resolver, &nsid).await;
        }

        // Otherwise parse as at-identifier (handle or DID) for bulk fetch
        let identifier = AtIdentifier::new(self.endpoint.as_ref())
            .map_err(|e| miette!("Invalid endpoint '{}': {}", self.endpoint, e))?;

        // Resolve to get PDS endpoint
        let did = match &identifier {
            AtIdentifier::Did(d) => d.clone().into_static(),
            AtIdentifier::Handle(h) => resolver.resolve_handle(h).await?,
        };

        let did_doc_resp = resolver.resolve_did_doc(&did).await?;

        let did_doc = did_doc_resp.parse()?;

        let pds = did_doc
            .pds_endpoint()
            .ok_or_else(|| miette!("No PDS endpoint found for {}", did))?;

        // Determine repo - use slice if provided, otherwise use the resolved DID
        let repo = if let Some(ref slice) = self.slice {
            AtIdentifier::new(slice.as_ref())
                .map_err(|e| miette!("Invalid slice '{}': {}", slice, e))?
                .into_static()
        } else {
            AtIdentifier::Did(did.clone())
        };

        let collection = Nsid::new("com.atproto.lexicon.schema")
            .map_err(|e| miette!("Invalid collection NSID: {}", e))?;

        let mut lexicons = HashMap::new();

        // Try to fetch all records at once first
        let req = ListRecords {
            repo: repo.clone().into_static().into(),
            collection: collection.clone().into_static(),
            cursor: None,
            limit: None,
            reverse: None,
        };

        let resp = resolver.xrpc(pds.clone()).send(&req).await?;

        match resp.into_output() {
            Ok(output) => {
                // Batch fetch succeeded
                for record_data in output.records {
                    if let Some(doc) = Self::parse_lexicon_record(&record_data) {
                        let nsid = doc.id.to_string();
                        lexicons.insert(nsid, doc);
                    }
                }
            }
            Err(e) => {
                // Batch decode failed, try one-by-one with cursor
                eprintln!("Warning: Batch decode failed from {}: {}", self.endpoint, e);
                eprintln!("Retrying with limit=1 to skip invalid records...");

                let mut cursor: Option<SmolStr> = None;
                loop {
                    let req = ListRecords {
                        repo: repo.clone().into_static().into(),
                        collection: collection.clone().into_static(),
                        cursor: cursor.clone(),
                        limit: Some(1),
                        reverse: None,
                    };
                    let resp = resolver.xrpc(pds.clone()).send(&req).await?;

                    match resp.into_output() {
                        Ok(output) => {
                            for record_data in output.records {
                                if let Some(doc) = Self::parse_lexicon_record(&record_data) {
                                    let nsid = doc.id.to_string();
                                    lexicons.insert(nsid, doc);
                                }
                            }

                            if let Some(next_cursor) = output.cursor {
                                cursor = Some(next_cursor);
                            } else {
                                break;
                            }
                        }
                        Err(e) => {
                            eprintln!("Warning: Failed to decode record (skipping): {}", e);
                            // Try to continue with next record if possible
                            // This is a bit tricky since we don't have the cursor from failed decode
                            // For now, just break
                            break;
                        }
                    }
                }
            }
        }

        Ok(lexicons)
    }
}
