use std::collections::BTreeMap;

use crate::storage::{
    PackTime, RawByHashEntry, RepositoryViewEntry, RepositoryViewState, RocksStore,
    RrdpSourceMemberRecord, RrdpSourceRecord, RrdpSourceSyncState, RrdpUriOwnerRecord,
    RrdpUriOwnerState,
};
use sha2::Digest;

#[allow(dead_code)]
pub struct PreparedRawByHashBatch {
    pub uri_to_hash: BTreeMap<String, String>,
    pub entries_to_write: Vec<RawByHashEntry>,
}

pub struct PreparedRepoBytesBatch {
    pub uri_to_hash: BTreeMap<String, String>,
    pub blobs_to_write: Vec<(String, Vec<u8>)>,
}

pub fn prepare_repo_bytes_batch(
    objects: &[(String, Vec<u8>)],
) -> Result<PreparedRepoBytesBatch, String> {
    prepare_repo_bytes_batch_owned(objects.to_vec())
}

pub fn prepare_repo_bytes_batch_owned(
    objects: Vec<(String, Vec<u8>)>,
) -> Result<PreparedRepoBytesBatch, String> {
    let mut uri_to_hash: BTreeMap<String, String> = BTreeMap::new();
    let mut pending: BTreeMap<String, Vec<u8>> = BTreeMap::new();

    for (uri, bytes) in objects {
        if bytes.is_empty() {
            return Err(format!("repo bytes for {uri} must not be empty"));
        }
        let sha256_hex = compute_sha256_hex(&bytes);
        uri_to_hash.insert(uri.clone(), sha256_hex.clone());
        match pending.entry(sha256_hex) {
            std::collections::btree_map::Entry::Vacant(slot) => {
                slot.insert(bytes);
            }
            std::collections::btree_map::Entry::Occupied(existing) => {
                if existing.get() != &bytes {
                    return Err(format!(
                        "repo bytes collision for {uri}: same sha256 maps to different bytes"
                    ));
                }
            }
        }
    }

    Ok(PreparedRepoBytesBatch {
        uri_to_hash,
        blobs_to_write: pending.into_iter().collect(),
    })
}

#[allow(dead_code)]
pub fn prepare_raw_by_hash_evidence_batch(
    store: &RocksStore,
    objects: &[(String, Vec<u8>)],
) -> Result<PreparedRawByHashBatch, String> {
    let mut pending: BTreeMap<String, RawByHashEntry> = BTreeMap::new();
    let mut uri_to_hash: BTreeMap<String, String> = BTreeMap::new();

    for (uri, bytes) in objects {
        let sha256_hex = compute_sha256_hex(bytes);
        uri_to_hash.insert(uri.clone(), sha256_hex.clone());
        let entry = pending
            .entry(sha256_hex.clone())
            .or_insert_with(|| RawByHashEntry::from_bytes(sha256_hex.clone(), bytes.clone()));

        if entry.bytes != *bytes {
            return Err(format!(
                "raw_by_hash collision for {uri}: same sha256 maps to different bytes"
            ));
        }
        if !entry.origin_uris.iter().any(|existing| existing == uri) {
            entry.origin_uris.push(uri.clone());
        }
        if entry.object_type.is_none() {
            entry.object_type = infer_object_type_from_uri(uri);
        }
    }

    let hashes: Vec<String> = pending.keys().cloned().collect();
    let existing_entries = store
        .get_raw_by_hash_entries_batch(&hashes)
        .map_err(|e| e.to_string())?;

    let mut entries_to_write = Vec::new();
    for (hash, existing_opt) in hashes.into_iter().zip(existing_entries.into_iter()) {
        let mut pending_entry = pending.remove(&hash).expect("pending raw_by_hash entry");
        match existing_opt {
            Some(mut existing) => {
                if existing.bytes != pending_entry.bytes {
                    return Err(format!(
                        "raw_by_hash collision for hash {hash}: same sha256 maps to different bytes"
                    ));
                }
                let mut changed = false;
                for uri in pending_entry.origin_uris.drain(..) {
                    if !existing
                        .origin_uris
                        .iter()
                        .any(|existing_uri| existing_uri == &uri)
                    {
                        existing.origin_uris.push(uri);
                        changed = true;
                    }
                }
                if existing.object_type.is_none() && pending_entry.object_type.is_some() {
                    existing.object_type = pending_entry.object_type;
                    changed = true;
                }
                if changed {
                    entries_to_write.push(existing);
                }
            }
            None => entries_to_write.push(pending_entry),
        }
    }

    Ok(PreparedRawByHashBatch {
        uri_to_hash,
        entries_to_write,
    })
}

pub fn upsert_repo_blob_bytes(store: &RocksStore, bytes: &[u8]) -> Result<String, String> {
    if bytes.is_empty() {
        return Err("repo bytes must not be empty".to_string());
    }
    let sha256_hex = compute_sha256_hex(bytes);
    store
        .put_blob_bytes_batch(&[(sha256_hex.clone(), bytes.to_vec())])
        .map_err(|e| e.to_string())?;
    Ok(sha256_hex)
}

pub fn infer_object_type_from_uri(uri: &str) -> Option<String> {
    let ext = uri.rsplit('.').next()?;
    let ext = ext.to_ascii_lowercase();
    match ext.as_str() {
        "cer" | "crl" | "mft" | "roa" | "asa" | "gbr" | "tal" | "xml" => Some(ext),
        _ => None,
    }
}

pub fn build_repository_view_present_entry(
    repository_source: &str,
    rsync_uri: &str,
    current_hash: &str,
) -> RepositoryViewEntry {
    RepositoryViewEntry {
        rsync_uri: rsync_uri.to_string(),
        current_hash: Some(current_hash.to_string()),
        repository_source: Some(repository_source.to_string()),
        object_type: infer_object_type_from_uri(rsync_uri),
        state: RepositoryViewState::Present,
    }
}

pub fn build_repository_view_withdrawn_entry(
    repository_source: &str,
    rsync_uri: &str,
    current_hash: Option<String>,
) -> RepositoryViewEntry {
    RepositoryViewEntry {
        rsync_uri: rsync_uri.to_string(),
        current_hash,
        repository_source: Some(repository_source.to_string()),
        object_type: infer_object_type_from_uri(rsync_uri),
        state: RepositoryViewState::Withdrawn,
    }
}

pub fn build_rrdp_source_member_present_record(
    notification_uri: &str,
    session_id: &str,
    serial: u64,
    rsync_uri: &str,
    current_hash: &str,
) -> RrdpSourceMemberRecord {
    RrdpSourceMemberRecord {
        notify_uri: notification_uri.to_string(),
        rsync_uri: rsync_uri.to_string(),
        current_hash: Some(current_hash.to_string()),
        object_type: infer_object_type_from_uri(rsync_uri),
        present: true,
        last_confirmed_session_id: session_id.to_string(),
        last_confirmed_serial: serial,
        last_changed_at: now_pack_time(),
    }
}

pub fn build_rrdp_source_member_withdrawn_record(
    notification_uri: &str,
    session_id: &str,
    serial: u64,
    rsync_uri: &str,
    current_hash: Option<String>,
) -> RrdpSourceMemberRecord {
    RrdpSourceMemberRecord {
        notify_uri: notification_uri.to_string(),
        rsync_uri: rsync_uri.to_string(),
        current_hash,
        object_type: infer_object_type_from_uri(rsync_uri),
        present: false,
        last_confirmed_session_id: session_id.to_string(),
        last_confirmed_serial: serial,
        last_changed_at: now_pack_time(),
    }
}

pub fn build_rrdp_uri_owner_active_record(
    notification_uri: &str,
    session_id: &str,
    serial: u64,
    rsync_uri: &str,
    current_hash: &str,
) -> RrdpUriOwnerRecord {
    RrdpUriOwnerRecord {
        rsync_uri: rsync_uri.to_string(),
        notify_uri: notification_uri.to_string(),
        current_hash: Some(current_hash.to_string()),
        last_confirmed_session_id: session_id.to_string(),
        last_confirmed_serial: serial,
        last_changed_at: now_pack_time(),
        owner_state: RrdpUriOwnerState::Active,
    }
}

pub fn build_rrdp_uri_owner_withdrawn_record(
    notification_uri: &str,
    session_id: &str,
    serial: u64,
    rsync_uri: &str,
    current_hash: Option<String>,
) -> RrdpUriOwnerRecord {
    RrdpUriOwnerRecord {
        rsync_uri: rsync_uri.to_string(),
        notify_uri: notification_uri.to_string(),
        current_hash,
        last_confirmed_session_id: session_id.to_string(),
        last_confirmed_serial: serial,
        last_changed_at: now_pack_time(),
        owner_state: RrdpUriOwnerState::Withdrawn,
    }
}

#[allow(dead_code)]
pub fn upsert_raw_by_hash_evidence(
    store: &RocksStore,
    rsync_uri: &str,
    bytes: &[u8],
) -> Result<String, String> {
    let prepared =
        prepare_raw_by_hash_evidence_batch(store, &[(rsync_uri.to_string(), bytes.to_vec())])?;
    store
        .put_raw_by_hash_entries_batch_unchecked(&prepared.entries_to_write)
        .map_err(|e| e.to_string())?;
    let sha256_hex = prepared
        .uri_to_hash
        .get(rsync_uri)
        .cloned()
        .expect("raw_by_hash mapping for upsert input");
    Ok(sha256_hex)
}

pub fn put_repository_view_present(
    store: &RocksStore,
    repository_source: &str,
    rsync_uri: &str,
    current_hash: &str,
) -> Result<(), String> {
    let entry = build_repository_view_present_entry(repository_source, rsync_uri, current_hash);
    store
        .put_repository_view_entry(&entry)
        .map_err(|e| e.to_string())
}

pub fn put_repository_view_withdrawn(
    store: &RocksStore,
    repository_source: &str,
    rsync_uri: &str,
    current_hash: Option<String>,
) -> Result<(), String> {
    let entry = build_repository_view_withdrawn_entry(repository_source, rsync_uri, current_hash);
    store
        .put_repository_view_entry(&entry)
        .map_err(|e| e.to_string())
}

pub fn ensure_rrdp_uri_can_be_owned_by(
    store: &RocksStore,
    notification_uri: &str,
    rsync_uri: &str,
) -> Result<(), String> {
    let Some(owner) = store
        .get_rrdp_uri_owner_record(rsync_uri)
        .map_err(|e| e.to_string())?
    else {
        return Ok(());
    };

    if owner.notify_uri != notification_uri && owner.owner_state == RrdpUriOwnerState::Active {
        return Err(format!(
            "RRDP source owner conflict for {rsync_uri}: current owner {} but incoming source {}",
            owner.notify_uri, notification_uri
        ));
    }
    Ok(())
}

pub fn current_rrdp_owner_is(
    store: &RocksStore,
    notification_uri: &str,
    rsync_uri: &str,
) -> Result<bool, String> {
    Ok(matches!(
        store
            .get_rrdp_uri_owner_record(rsync_uri)
            .map_err(|e| e.to_string())?,
        Some(owner)
            if owner.notify_uri == notification_uri && owner.owner_state == RrdpUriOwnerState::Active
    ))
}

pub fn put_rrdp_source_member_present(
    store: &RocksStore,
    notification_uri: &str,
    session_id: &str,
    serial: u64,
    rsync_uri: &str,
    current_hash: &str,
) -> Result<(), String> {
    let record = build_rrdp_source_member_present_record(
        notification_uri,
        session_id,
        serial,
        rsync_uri,
        current_hash,
    );
    store
        .put_rrdp_source_member_record(&record)
        .map_err(|e| e.to_string())
}

pub fn put_rrdp_source_member_withdrawn(
    store: &RocksStore,
    notification_uri: &str,
    session_id: &str,
    serial: u64,
    rsync_uri: &str,
    current_hash: Option<String>,
) -> Result<(), String> {
    let record = build_rrdp_source_member_withdrawn_record(
        notification_uri,
        session_id,
        serial,
        rsync_uri,
        current_hash,
    );
    store
        .put_rrdp_source_member_record(&record)
        .map_err(|e| e.to_string())
}

pub fn put_rrdp_uri_owner_active(
    store: &RocksStore,
    notification_uri: &str,
    session_id: &str,
    serial: u64,
    rsync_uri: &str,
    current_hash: &str,
) -> Result<(), String> {
    let record = build_rrdp_uri_owner_active_record(
        notification_uri,
        session_id,
        serial,
        rsync_uri,
        current_hash,
    );
    store
        .put_rrdp_uri_owner_record(&record)
        .map_err(|e| e.to_string())
}

pub fn put_rrdp_uri_owner_withdrawn(
    store: &RocksStore,
    notification_uri: &str,
    session_id: &str,
    serial: u64,
    rsync_uri: &str,
    current_hash: Option<String>,
) -> Result<(), String> {
    let record = build_rrdp_uri_owner_withdrawn_record(
        notification_uri,
        session_id,
        serial,
        rsync_uri,
        current_hash,
    );
    store
        .put_rrdp_uri_owner_record(&record)
        .map_err(|e| e.to_string())
}

pub fn update_rrdp_source_record_on_success(
    store: &RocksStore,
    notification_uri: &str,
    session_id: &str,
    serial: u64,
    sync_state: RrdpSourceSyncState,
    last_snapshot_uri: Option<&str>,
    last_snapshot_hash_hex: Option<&str>,
) -> Result<(), String> {
    let now = now_pack_time();
    let mut record = match store
        .get_rrdp_source_record(notification_uri)
        .map_err(|e| e.to_string())?
    {
        Some(existing) => existing,
        None => RrdpSourceRecord {
            notify_uri: notification_uri.to_string(),
            last_session_id: None,
            last_serial: None,
            first_seen_at: now.clone(),
            last_seen_at: now.clone(),
            last_sync_at: None,
            sync_state,
            last_snapshot_uri: None,
            last_snapshot_hash: None,
            last_error: None,
        },
    };

    record.last_session_id = Some(session_id.to_string());
    record.last_serial = Some(serial);
    record.last_seen_at = now.clone();
    record.last_sync_at = Some(now);
    record.sync_state = sync_state;
    record.last_snapshot_uri = last_snapshot_uri.map(str::to_string);
    record.last_snapshot_hash = last_snapshot_hash_hex.map(str::to_string);
    record.last_error = None;

    store
        .put_rrdp_source_record(&record)
        .map_err(|e| e.to_string())
}

pub fn compute_sha256_hex(bytes: &[u8]) -> String {
    hex::encode(sha2::Sha256::digest(bytes))
}

pub fn now_pack_time() -> PackTime {
    PackTime::from_utc_offset_datetime(time::OffsetDateTime::now_utc())
}

#[cfg(test)]
mod tests {
    use super::{prepare_repo_bytes_batch, prepare_repo_bytes_batch_owned};
    use std::collections::BTreeSet;

    #[test]
    fn prepare_repo_bytes_batch_deduplicates_by_hash() {
        let objects = vec![
            (
                "rsync://example.test/repo/a.roa".to_string(),
                b"same".to_vec(),
            ),
            (
                "rsync://example.test/repo/b.roa".to_string(),
                b"same".to_vec(),
            ),
            (
                "rsync://example.test/repo/c.roa".to_string(),
                b"other".to_vec(),
            ),
        ];
        let prepared = prepare_repo_bytes_batch(&objects).expect("prepare repo bytes");
        assert_eq!(prepared.uri_to_hash.len(), 3);
        assert_eq!(prepared.blobs_to_write.len(), 2);
        let unique_hashes = prepared
            .blobs_to_write
            .iter()
            .map(|(hash, _)| hash.clone())
            .collect::<BTreeSet<_>>();
        assert_eq!(unique_hashes.len(), 2);
    }

    #[test]
    fn prepare_repo_bytes_batch_owned_deduplicates_without_borrowed_input() {
        let objects = vec![
            (
                "rsync://example.test/repo/a.roa".to_string(),
                b"same".to_vec(),
            ),
            (
                "rsync://example.test/repo/b.roa".to_string(),
                b"same".to_vec(),
            ),
        ];
        let prepared = prepare_repo_bytes_batch_owned(objects).expect("prepare repo bytes");
        assert_eq!(prepared.uri_to_hash.len(), 2);
        assert_eq!(prepared.blobs_to_write.len(), 1);
    }
}
