use std::collections::HashMap;

use anyhow::{Context, Result, anyhow, bail, ensure};
use log::debug;
use storage_backend::{
    s3_multi_file::{FileMeta, download_one_object_as_multi_files_with_retry},
    storage::{BackendStorage, ReadRange},
};
use tantivy::index::SegmentId;
use tici_grpc::{FragMeta, FragProperty, FragStorageRef};
use walkdir::WalkDir;

use crate::{
    utils::{self, FRAG_ROW_ENTRY_FILE, frag_multi_path},
    writer::{frag_batch_uploader::FragBatchUploader, tantivy_doc_buffer::HandleStat},
};

pub fn list_all_files_in_absolute(path: &str) -> Result<Vec<String>> {
    let mut files = Vec::new();
    for entry in WalkDir::new(path) {
        let entry = entry?;
        if entry.file_type().is_file() && !entry.file_name().to_string_lossy().starts_with('.') {
            files.push(entry.path().display().to_string());
        }
    }
    Ok(files)
}

fn strip_to_relative(parent_path: &str, path: &str) -> Result<String> {
    Ok(path
        .strip_prefix(parent_path)
        .ok_or_else(|| {
            anyhow!(
                "strip_prefix failed: parent_path={}, path={}",
                parent_path,
                path
            )
        })?
        .trim_start_matches('/')
        .trim_end_matches('/')
        .to_string())
}

pub fn list_all_files_in_relative(path: &str) -> Result<Vec<String>> {
    let mut files = Vec::new();
    for entry in WalkDir::new(path) {
        let entry = entry?;
        if entry.file_type().is_file() && !entry.file_name().to_string_lossy().starts_with('.') {
            files.push(strip_to_relative(
                path,
                &entry.path().display().to_string(),
            )?);
        }
    }
    Ok(files)
}

fn ensure_valid_s3_path(s3_path: &str) -> Result<()> {
    ensure!(
        !s3_path.is_empty() && !s3_path.starts_with('/') && !s3_path.ends_with('/'),
        "{s3_path} is empty or starts with '/' or ends with '/'"
    );
    Ok(())
}

pub fn to_s3_path(s3_prefix: &str, local_prefix_path: &str, local_path: &str) -> Result<String> {
    ensure_valid_s3_path(s3_prefix)?;
    let relative_path = strip_to_relative(local_prefix_path, local_path)?;
    // avoid s3 path starts or ends with '/'
    if relative_path.is_empty() {
        return Ok(s3_prefix.to_string());
    }
    Ok(format!("{}/{}", s3_prefix, relative_path))
}

fn to_local_path(local_parent_path: &str, s3_path: &str) -> String {
    format!("{}/{}", local_parent_path, s3_path)
}

fn strip_storage_ref_key_frag_path(frag_path: &str, key: &str) -> Result<String> {
    ensure_valid_s3_path(frag_path)?;
    let sub_key = key.strip_prefix(frag_path).ok_or_else(|| {
        anyhow!(
            "storage ref key does not start with frag_path: key={}, frag_path={}",
            key,
            frag_path
        )
    })?;
    ensure!(
        sub_key.starts_with('/'),
        "storage ref key does not have frag_path boundary: key={}, frag_path={}",
        key,
        frag_path
    );
    Ok(sub_key.trim_start_matches('/').to_string())
}

fn storage_ref_download_key(frag_path: &str, storage_ref_key: &str) -> Result<String> {
    ensure_valid_s3_path(frag_path)?;
    ensure_valid_s3_path(storage_ref_key)?;
    Ok(format!("{frag_path}/{storage_ref_key}"))
}

pub async fn prepare_local_frag_path(local_frag_path: &str) -> Result<()> {
    let local_frag_index_path = utils::frag_index_path(local_frag_path);
    tokio::fs::create_dir_all(&local_frag_index_path)
        .await
        .with_context(|| {
            format!(
                "Failed to create local frag index path={}",
                local_frag_index_path
            )
        })?;
    Ok(())
}

pub fn temporary_download_path(local_data_path: &str) -> String {
    format!("{}/downloading", local_data_path)
}

// Need a way to avoid concurrently download the same frag.
pub async fn download_frag(
    backend_storage: &dyn BackendStorage,
    bucket: &str,
    frag_meta: &FragMeta,
    local_data_path: &str,
) -> Result<String> {
    frag_meta.validate()?;
    let tmp_local_data_path = temporary_download_path(local_data_path);
    let tmp_local_frag_path = to_local_path(&tmp_local_data_path, &frag_meta.frag_path);
    prepare_local_frag_path(&tmp_local_frag_path).await?;
    for storage_ref in &frag_meta.storage_refs {
        let range = ReadRange::new(storage_ref.offset, storage_ref.size).with_context(|| {
            format!(
                "invalid frag storage ref: sub_key={}, offset={}, size={}",
                storage_ref.sub_key, storage_ref.offset, storage_ref.size
            )
        })?;
        download_one_object_as_multi_files_with_retry(
            backend_storage,
            bucket,
            &storage_ref_download_key(&frag_meta.frag_path, &storage_ref.sub_key)?,
            Some(range),
            &tmp_local_frag_path,
        )
        .await?;
    }

    let local_frag_path = to_local_path(local_data_path, &frag_meta.frag_path);
    // Create parent path of local_frag_path
    {
        let local_frag_path = std::path::Path::new(&local_frag_path);
        if let Some(parent) = local_frag_path.parent() {
            std::fs::create_dir_all(parent)
                .with_context(|| format!("Create parent directory for {:?}", local_frag_path))?;
        }
    }
    std::fs::rename(&tmp_local_frag_path, &local_frag_path)
        .with_context(|| format!("Rename {} to {}", tmp_local_frag_path, local_frag_path))?;
    Ok(local_frag_path)
}

// const FRAG_INDEX_FILES_TO_FILTER_OUT: &[&str] = &["index/.managed.json",
// "index/.tantivy-meta.lock", "index/.tantivy-writer.lock"];
pub const FRAG_INDEX_SEG_FILE_SUFFIXES: &[&str] =
    &[".fast", ".fieldnorm", ".idx", ".pos", ".store", ".term"];
pub const FRAG_INDEX_META_FILE: &str = "index/meta.json";

fn extract_seg_id_and_suffix(fname: &str) -> Result<(String, String)> {
    let s = fname
        .split('/')
        .next_back()
        .ok_or_else(|| anyhow::anyhow!("split '/' failed, file={}", fname))?;
    if let Some((seg_id, suffix)) = s.split_once('.') {
        Ok((seg_id.to_string(), format!(".{}", suffix)))
    } else {
        Err(anyhow!("split '.' failed, str={}, fname={}", s, fname))
    }
}

fn encode_frag_meta(
    local_frag_path: &str,
    s3_frag_path: &str,
    relative_fnames: &[String],
    data_size: u64,
    handle_stat: HandleStat,
    storage_refs: Vec<FragStorageRef>,
) -> Result<FragMeta> {
    let mut seg_files: HashMap<String, Vec<String>> = HashMap::new();
    let mut has_index_meta_file = false;
    for fname in relative_fnames {
        if fname == FRAG_ROW_ENTRY_FILE {
            continue;
        }
        if fname == FRAG_INDEX_META_FILE {
            has_index_meta_file = true;
            continue;
        }
        if !fname.starts_with("index/") {
            bail!(
                "Unknown file={}, local_frag_path={}, relative_fnames={:?}",
                fname,
                local_frag_path,
                relative_fnames
            );
        }

        let (seg_id, suffix) = extract_seg_id_and_suffix(fname)?;
        // The suffix of delmark file is ".{opstamp}.del", opstamp is a incremental
        // number.
        if FRAG_INDEX_SEG_FILE_SUFFIXES.contains(&suffix.as_str()) || suffix.ends_with(".del") {
            seg_files
                .entry(seg_id.to_string())
                .or_default()
                .push(suffix.to_string());
        } else {
            bail!(
                "Unknown file={}, seg_id={}, suffix={}, local_frag_path={}, relative_fnames={:?}",
                fname,
                seg_id,
                suffix,
                local_frag_path,
                relative_fnames
            );
        }
    }

    if !has_index_meta_file {
        bail!(
            "Missing index/meta.json, local_frag_path={}, relative_fnames={:?}",
            local_frag_path,
            relative_fnames
        );
    }

    for (seg_id, suffixes) in &mut seg_files {
        suffixes.sort();
        // The delmark will always be the first file after sorting.
        if suffixes[0].ends_with(".del") {
            suffixes.remove(0);
        };

        if suffixes != FRAG_INDEX_SEG_FILE_SUFFIXES {
            bail!(
                "Missing seg files for seg_id={}, local_frag_path={}, relative_filenames={:?}",
                seg_id,
                local_frag_path,
                relative_fnames
            );
        }
    }

    FragMeta::new(
        1,
        s3_frag_path,
        FragProperty {
            size: data_size,
            count: handle_stat.handle_count as u64,
            min_handle: handle_stat.min_handle,
            split_handles: vec![handle_stat.mid_handle],
            max_handle: handle_stat.max_handle,
            delete_count: handle_stat.delete_count as u64,
        },
        storage_refs,
        u32::try_from(seg_files.len()).context("too many index segments")?,
    )
}

fn filter_files_by_seg_ids(relative_fnames: &mut Vec<String>, seg_ids: &[SegmentId]) {
    let seg_path_prefixes = seg_ids
        .iter()
        .map(|seg_id| format!("index/{}", seg_id.uuid_string()))
        .collect::<Vec<_>>();

    debug!(
        "Input relative_fnames={:?}, seg_ids={:?}",
        relative_fnames, seg_ids
    );
    relative_fnames.retain(|fname| {
        if fname == FRAG_ROW_ENTRY_FILE || fname == FRAG_INDEX_META_FILE {
            true
        } else {
            seg_path_prefixes
                .iter()
                .any(|prefix| fname.starts_with(prefix))
        }
    });
    debug!("Filtered relative_fnames={:?}", relative_fnames);
}

pub fn parse_frag_id(frag_path: &str) -> Result<u64> {
    let frag_name = frag_path
        .split('/')
        .next_back()
        .ok_or_else(|| anyhow!("Invalid frag_path={}", frag_path))?;
    if let Some(frag_id_str) = frag_name.strip_prefix("f_") {
        let frag_id = frag_id_str
            .parse::<u64>()
            .with_context(|| format!("Failed to parse frag_id from frag_name={}", frag_name))?;
        Ok(frag_id)
    } else {
        Err(anyhow!("Invalid frag_name={}", frag_name))
    }
}

pub async fn upload_frag(
    frag_batch_uploader: &FragBatchUploader,
    bucket: &str,
    s3_prefix: &str,
    local_data_path: &str,
    local_frag_path: &str,
    handle_stat: HandleStat,
    seg_ids: &[SegmentId],
) -> Result<FragMeta> {
    let s3_frag_path = to_s3_path(s3_prefix, local_data_path, local_frag_path)?;
    let mut relative_fnames = list_all_files_in_relative(local_frag_path)?;
    filter_files_by_seg_ids(&mut relative_fnames, seg_ids);
    if relative_fnames.is_empty() {
        bail!(
            "Missing files. local_frag_path={}, seg_ids={:?}",
            local_frag_path,
            seg_ids
        );
    }

    let mut file_metas = Vec::with_capacity(relative_fnames.len());
    for fname in &relative_fnames {
        let full_path = std::path::Path::new(local_frag_path).join(fname);
        let meta = std::fs::metadata(&full_path)
            .with_context(|| format!("Get metadata of {:?}", full_path))?;
        file_metas.push(FileMeta {
            relative_path: fname.clone(),
            size: meta.len(),
        });
    }

    let receivers = frag_batch_uploader
        .upload(bucket, &s3_frag_path, local_frag_path, file_metas)
        .await?;

    let mut data_size = 0u64;
    let mut storage_refs = Vec::with_capacity(receivers.len());
    for (i, rx) in receivers.into_iter().enumerate() {
        let resp = rx
            .await
            .map_err(|e| anyhow!("Upload channel closed: {:?}", e))??;
        let expected_key = frag_multi_path(&s3_frag_path, i as u32);
        ensure!(
            resp.key == expected_key,
            "Uploaded file key mismatch, expected={}, actual={}",
            expected_key,
            resp.key
        );
        // current contract: each storage_ref is its own object
        ensure!(
            resp.offset == 0,
            "Uploaded file offset mismatch, expected=0, actual={}",
            resp.offset
        );
        ensure!(
            resp.size > 0,
            "Uploaded file size mismatch, expected>0, actual={}",
            resp.size
        );
        data_size += resp.size;
        storage_refs.push(FragStorageRef {
            sub_key: strip_storage_ref_key_frag_path(&s3_frag_path, &resp.key)?,
            offset: resp.offset,
            size: resp.size,
        });
    }

    encode_frag_meta(
        local_frag_path,
        &s3_frag_path,
        &relative_fnames,
        data_size,
        handle_stat,
        storage_refs,
    )
}

#[cfg(test)]
mod tests {
    use bytes::Bytes;
    use futures::StreamExt;
    use storage_backend::s3_multi_file::{
        FileBundle, FileMeta, MultiFileStream, upload_multi_files_as_one_object_with_retry,
    };
    use tici_grpc::{FragMeta, FragProperty, FragStorageRef};
    use tici_test_utils::test_utils::init_logger_for_ut;
    use tokio::io::AsyncWriteExt;

    use crate::{
        test_utils::{
            MockBackendStorage, TEST_S3_BUCKET, test_create_backend_storage,
            test_get_frag_batch_uploader,
        },
        writer::tantivy_doc_buffer::HandleStat,
    };

    async fn create_multi_file_object_data() -> Vec<u8> {
        let temp_dir = tempfile::tempdir().unwrap();
        let local_frag_path = temp_dir.path().join("frag");
        std::fs::create_dir_all(local_frag_path.join("index")).unwrap();
        let meta_path = local_frag_path.join("index/meta.json");
        std::fs::write(&meta_path, b"{}").unwrap();

        let file_bundle = FileBundle::new(
            local_frag_path.display().to_string(),
            1,
            vec![FileMeta {
                relative_path: "index/meta.json".to_string(),
                size: std::fs::metadata(&meta_path).unwrap().len(),
            }],
        );
        let stream = MultiFileStream::new(&file_bundle).unwrap().into_stream();
        let chunks = stream.map(|chunk| chunk.unwrap()).collect::<Vec<_>>().await;
        let total_size = chunks.iter().map(Bytes::len).sum();
        let mut data = Vec::with_capacity(total_size);
        for chunk in chunks {
            data.extend_from_slice(&chunk);
        }
        data
    }

    #[test]
    fn test_strip_storage_ref_key_frag_path() {
        assert_eq!(
            super::strip_storage_ref_key_frag_path(
                "upload-prefix/t_1/i_2/f_3",
                "upload-prefix/t_1/i_2/f_3/multi.0"
            )
            .unwrap(),
            "multi.0"
        );
        assert_eq!(
            super::strip_storage_ref_key_frag_path(
                "upload-prefix/t_1/i_2/f_3",
                "upload-prefix/t_1/i_2/f_3/index/multi.0"
            )
            .unwrap(),
            "index/multi.0"
        );
        assert!(
            super::strip_storage_ref_key_frag_path(
                "upload-prefix/t_1/i_2/f_3",
                "upload-prefix/t_1/i_2/f_30/multi.0"
            )
            .is_err()
        );
    }

    #[test]
    fn test_storage_ref_download_key() {
        assert_eq!(
            super::storage_ref_download_key("upload-prefix/t_1/i_2/f_3", "multi.0").unwrap(),
            "upload-prefix/t_1/i_2/f_3/multi.0"
        );
        assert_eq!(
            super::storage_ref_download_key("upload-prefix/t_1/i_2/f_3", "index/multi.0").unwrap(),
            "upload-prefix/t_1/i_2/f_3/index/multi.0"
        );
        assert!(super::storage_ref_download_key("upload-prefix/", "multi.0").is_err());
        assert!(super::storage_ref_download_key("upload-prefix", "/multi.0").is_err());
    }

    #[tokio::test]
    async fn test_download_frag_resolves_relative_storage_ref_key() {
        let object_key = "upload-prefix/t_1/i_2/f_3/multi.0";
        let data = create_multi_file_object_data().await;
        let backend_storage = MockBackendStorage::new();
        backend_storage.push_object_response(TEST_S3_BUCKET, object_key, Ok(data.clone()));
        let temp_dir = tempfile::tempdir().unwrap();
        let frag_meta = FragMeta::new(
            1,
            "upload-prefix/t_1/i_2/f_3",
            FragProperty {
                size: data.len() as u64,
                split_handles: vec![vec![]],
                ..Default::default()
            },
            vec![FragStorageRef {
                sub_key: "multi.0".to_string(),
                offset: 0,
                size: data.len() as u64,
            }],
            0,
        )
        .unwrap();

        let local_frag_path = super::download_frag(
            &backend_storage,
            TEST_S3_BUCKET,
            &frag_meta,
            temp_dir.path().to_str().unwrap(),
        )
        .await
        .unwrap();

        assert!(
            std::path::Path::new(&local_frag_path)
                .join("index/meta.json")
                .exists()
        );
    }

    #[tokio::test]
    async fn test_upload_frag_strips_frag_path_from_storage_ref_key() {
        let frag_batch_uploader = test_get_frag_batch_uploader().await.unwrap();
        let temp_dir = tempfile::tempdir().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let local_frag_path = temp_dir.path().join("t_1/i_2/f_3");
        std::fs::create_dir_all(local_frag_path.join("index")).unwrap();
        std::fs::write(local_frag_path.join("index/meta.json"), b"{}").unwrap();
        let s3_prefix = format!("test-upload-strip-{}", uuid::Uuid::new_v4());

        let frag_meta = super::upload_frag(
            &frag_batch_uploader,
            TEST_S3_BUCKET,
            &s3_prefix,
            local_data_path,
            local_frag_path.to_str().unwrap(),
            HandleStat::default(),
            &[],
        )
        .await
        .unwrap();

        assert_eq!(frag_meta.frag_path, format!("{s3_prefix}/t_1/i_2/f_3"));
        assert_eq!(frag_meta.storage_refs.len(), 1);
        assert_eq!(frag_meta.storage_refs[0].sub_key, "multi.0");
    }

    #[tokio::test]
    async fn test_upload_large_file_chunk_limit() {
        init_logger_for_ut();
        let tmpdir = tempfile::tempdir().unwrap();
        let big_file_path = tmpdir.path().join("big_file.bin");
        let mut f = tokio::fs::File::create(&big_file_path).await.unwrap();
        let buf = vec![0u8; 20 * 1024 * 1024];
        f.write_all(&buf).await.unwrap();
        f.flush().await.unwrap();

        let backend_storage = test_create_backend_storage().await.unwrap();

        let files = vec![FileMeta {
            relative_path: "big_file.bin".into(),
            size: 20 * 1024 * 1024,
        }];

        let bundle = FileBundle::new(tmpdir.path().to_str().unwrap(), 123, files);
        upload_multi_files_as_one_object_with_retry(
            backend_storage.as_ref(),
            TEST_S3_BUCKET,
            "big_file.bin",
            &bundle,
        )
        .await
        .unwrap();
    }

    #[tokio::test]
    async fn test_download_frag_rejects_empty_storage_refs() {
        let backend_storage = test_create_backend_storage().await.unwrap();
        let temp_dir = tempfile::tempdir().unwrap();
        let frag_meta = FragMeta {
            frag_path: "prefix/f_1".to_string(),
            property: Some(FragProperty {
                split_handles: vec![vec![]],
                ..Default::default()
            }),
            ..Default::default()
        };

        let err = super::download_frag(
            backend_storage.as_ref(),
            TEST_S3_BUCKET,
            &frag_meta,
            temp_dir.path().to_str().unwrap(),
        )
        .await
        .unwrap_err();

        assert!(err.to_string().contains("storage_refs must not be empty"));
    }
}