use std::{
    collections::{BTreeMap, VecDeque},
    fs::File,
    path::Path,
};

use anyhow::{Context, Result, anyhow, bail};
use parquet::file::reader::{FileReader, RowGroupReader, SerializedFileReader};

use crate::profile::{ColumnProfile, profile_column_iter};

const INITIAL_CHUNK_SIZE: usize = 1000;
const CSV_LOAD_CHUNK_SIZE: usize = 10_000;
const PARQUET_LOAD_CHUNK_SIZE: usize = 20_000;
const MAX_CACHED_SEGMENTS: usize = 4;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FileFormat {
    Csv,
    Parquet,
}

#[derive(Debug, Clone)]
pub struct DatasetPreview {
    pub path: String,
    pub format: FileFormat,
    pub headers: Vec<String>,
    pub limit: Option<usize>,
    pub exhausted: bool,
    pub total_rows_hint: Option<usize>,
    loaded_rows: usize,
    segments: BTreeMap<usize, Vec<Vec<String>>>,
    segment_order: VecDeque<usize>,
    chunk_size: usize,
    cache_version: u64,
}

impl DatasetPreview {
    pub fn column_count(&self) -> usize {
        self.headers.len()
    }

    pub fn row_count(&self) -> usize {
        self.loaded_rows
    }

    pub fn total_rows_hint(&self) -> Option<usize> {
        self.total_rows_hint
    }

    pub fn cache_version(&self) -> u64 {
        self.cache_version
    }

    pub fn load_more(&mut self) -> Result<usize> {
        if self.exhausted {
            return Ok(0);
        }

        let start = self.loaded_rows;
        let Some(count) = read_count(self.limit, start, self.chunk_size) else {
            self.exhausted = true;
            return Ok(0);
        };

        let loaded = load_rows(Path::new(&self.path), &self.format, start, count)?;
        let added = loaded.len();
        self.loaded_rows = self.loaded_rows.max(start + added);
        self.insert_segment(start, loaded);

        if added < count || self.limit.is_some_and(|limit| self.loaded_rows >= limit) {
            self.exhausted = true;
        }

        Ok(added)
    }

    pub fn ensure_row_loaded(&mut self, row: usize) -> Result<()> {
        while row >= self.loaded_rows && !self.exhausted {
            self.load_more()?;
        }

        if self.has_row_cached(row) || row >= self.loaded_rows {
            return Ok(());
        }

        let start = (row / self.chunk_size) * self.chunk_size;
        let Some(count) = read_count(self.limit, start, self.chunk_size) else {
            return Ok(());
        };
        let loaded = load_rows(Path::new(&self.path), &self.format, start, count)?;
        self.insert_segment(start, loaded);
        Ok(())
    }

    pub fn value(&self, row: usize, col: usize) -> &str {
        self.segment_for_row(row)
            .and_then(|(start, rows)| rows.get(row - start))
            .and_then(|row| row.get(col))
            .map(String::as_str)
            .unwrap_or("")
    }

    pub fn column_profile(&self, col: usize) -> Option<ColumnProfile> {
        let name = self.headers.get(col)?;
        Some(profile_column_iter(
            name,
            self.segments.values().flat_map(|rows| rows.iter()),
            col,
        ))
    }

    fn insert_segment(&mut self, start: usize, rows: Vec<Vec<String>>) {
        if rows.is_empty() {
            return;
        }

        if !self.segments.contains_key(&start) {
            self.segment_order.push_back(start);
        }
        self.segments.insert(start, rows);
        self.cache_version = self.cache_version.saturating_add(1);

        while self.segment_order.len() > MAX_CACHED_SEGMENTS {
            if let Some(old_start) = self.segment_order.pop_front()
                && old_start != start
            {
                self.segments.remove(&old_start);
                self.cache_version = self.cache_version.saturating_add(1);
            }
        }
    }

    fn has_row_cached(&self, row: usize) -> bool {
        self.segment_for_row(row).is_some()
    }

    fn segment_for_row(&self, row: usize) -> Option<(usize, &Vec<Vec<String>>)> {
        self.segments
            .range(..=row)
            .next_back()
            .filter(|(start, rows)| row < **start + rows.len())
            .map(|(start, rows)| (*start, rows))
    }
}

pub fn detect_format(path: &Path) -> Result<FileFormat> {
    match path
        .extension()
        .and_then(|ext| ext.to_str())
        .map(str::to_ascii_lowercase)
    {
        Some(ext) if ext == "csv" => Ok(FileFormat::Csv),
        Some(ext) if ext == "parquet" || ext == "pq" => Ok(FileFormat::Parquet),
        Some(ext) => bail!("unsupported file extension: .{ext}; expected .csv or .parquet"),
        None => bail!("file has no extension; expected .csv or .parquet"),
    }
}

pub fn load_preview(path: &Path, limit: Option<usize>) -> Result<DatasetPreview> {
    let format = detect_format(path)?;
    let headers = load_headers(path, &format)?;
    let total_rows_hint = total_rows_hint(path, &format)?;
    let count = read_count(limit, 0, INITIAL_CHUNK_SIZE).unwrap_or(0);
    let rows = load_rows(path, &format, 0, count)?;
    let exhausted = rows.len() < count || limit.is_some_and(|limit| rows.len() >= limit);
    let chunk_size = load_chunk_size(&format);
    let loaded_rows = rows.len();
    let mut segments = BTreeMap::new();
    if !rows.is_empty() {
        segments.insert(0, rows);
    }
    let segment_order = segments.keys().copied().collect();

    Ok(DatasetPreview {
        path: path.display().to_string(),
        format,
        headers,
        limit,
        exhausted,
        total_rows_hint,
        loaded_rows,
        segments,
        segment_order,
        chunk_size,
        cache_version: 0,
    })
}

fn load_chunk_size(format: &FileFormat) -> usize {
    match format {
        FileFormat::Csv => CSV_LOAD_CHUNK_SIZE,
        FileFormat::Parquet => PARQUET_LOAD_CHUNK_SIZE,
    }
}

fn read_count(limit: Option<usize>, loaded: usize, chunk_size: usize) -> Option<usize> {
    match limit {
        Some(limit) if loaded >= limit => None,
        Some(limit) => Some((limit - loaded).min(chunk_size)),
        None => Some(chunk_size),
    }
}

fn load_headers(path: &Path, format: &FileFormat) -> Result<Vec<String>> {
    match format {
        FileFormat::Csv => load_csv_headers(path),
        FileFormat::Parquet => load_parquet_headers(path),
    }
}

fn total_rows_hint(path: &Path, format: &FileFormat) -> Result<Option<usize>> {
    match format {
        FileFormat::Csv => Ok(None),
        FileFormat::Parquet => {
            let file = File::open(path)
                .with_context(|| format!("failed to open Parquet: {}", path.display()))?;
            let reader = SerializedFileReader::new(file)
                .with_context(|| format!("failed to read Parquet metadata: {}", path.display()))?;
            Ok(usize::try_from(reader.metadata().file_metadata().num_rows()).ok())
        }
    }
}

fn load_rows(
    path: &Path,
    format: &FileFormat,
    offset: usize,
    count: usize,
) -> Result<Vec<Vec<String>>> {
    match format {
        FileFormat::Csv => load_csv_rows(path, offset, count),
        FileFormat::Parquet => load_parquet_rows(path, offset, count),
    }
}

fn load_csv_headers(path: &Path) -> Result<Vec<String>> {
    let mut reader = csv::Reader::from_path(path)
        .with_context(|| format!("failed to open CSV: {}", path.display()))?;
    let headers = reader
        .headers()
        .with_context(|| format!("failed to read CSV header: {}", path.display()))?
        .iter()
        .map(ToOwned::to_owned)
        .collect::<Vec<_>>();

    Ok(headers)
}

fn load_csv_rows(path: &Path, offset: usize, count: usize) -> Result<Vec<Vec<String>>> {
    let mut reader = csv::Reader::from_path(path)
        .with_context(|| format!("failed to open CSV: {}", path.display()))?;

    let mut rows = Vec::new();
    for (idx, record) in reader.records().enumerate() {
        if idx < offset {
            continue;
        }

        if rows.len() >= count {
            break;
        }

        let record =
            record.with_context(|| format!("failed to read CSV row: {}", path.display()))?;
        rows.push(record.iter().map(ToOwned::to_owned).collect());
    }

    Ok(rows)
}

fn load_parquet_headers(path: &Path) -> Result<Vec<String>> {
    let file =
        File::open(path).with_context(|| format!("failed to open Parquet: {}", path.display()))?;
    let reader = SerializedFileReader::new(file)
        .with_context(|| format!("failed to read Parquet metadata: {}", path.display()))?;

    let headers = reader
        .metadata()
        .file_metadata()
        .schema_descr()
        .columns()
        .iter()
        .map(|col| col.name().to_string())
        .collect::<Vec<_>>();

    if headers.is_empty() {
        return Err(anyhow!(
            "Parquet file has no visible columns: {}",
            path.display()
        ));
    }

    Ok(headers)
}

fn load_parquet_rows(path: &Path, offset: usize, count: usize) -> Result<Vec<Vec<String>>> {
    let file =
        File::open(path).with_context(|| format!("failed to open Parquet: {}", path.display()))?;
    let reader = SerializedFileReader::new(file)
        .with_context(|| format!("failed to read Parquet metadata: {}", path.display()))?;
    let mut rows = Vec::new();
    let end = offset.saturating_add(count);
    let mut row_group_start = 0usize;

    for row_group_idx in 0..reader.num_row_groups() {
        let row_group = reader.get_row_group(row_group_idx).with_context(|| {
            format!(
                "failed to read Parquet row group {row_group_idx}: {}",
                path.display()
            )
        })?;
        let row_group_rows = usize::try_from(row_group.metadata().num_rows()).unwrap_or(usize::MAX);
        let row_group_end = row_group_start.saturating_add(row_group_rows);

        if row_group_end <= offset {
            row_group_start = row_group_end;
            continue;
        }

        if row_group_start >= end {
            break;
        }

        read_parquet_row_group_rows(
            path,
            row_group.as_ref(),
            row_group_start,
            offset,
            end,
            &mut rows,
        )?;

        if rows.len() >= count {
            break;
        }

        row_group_start = row_group_end;
    }

    Ok(rows)
}

fn read_parquet_row_group_rows(
    path: &Path,
    row_group: &dyn RowGroupReader,
    row_group_start: usize,
    offset: usize,
    end: usize,
    rows: &mut Vec<Vec<String>>,
) -> Result<()> {
    let mut row_iter = row_group
        .get_row_iter(None)
        .with_context(|| format!("failed to read Parquet row group rows: {}", path.display()))?;

    for (local_idx, row) in (&mut row_iter).enumerate() {
        let global_idx = row_group_start + local_idx;
        if global_idx < offset {
            continue;
        }

        if global_idx >= end {
            break;
        }

        let row =
            row.with_context(|| format!("failed to decode Parquet row: {}", path.display()))?;
        let columns = row
            .get_column_iter()
            .map(|(_, field)| field.to_string())
            .collect::<Vec<_>>();

        rows.push(columns);
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::{fs, time::SystemTime};

    #[test]
    fn detects_csv_extension() {
        assert_eq!(
            detect_format(Path::new("data.csv")).unwrap(),
            FileFormat::Csv
        );
    }

    #[test]
    fn detects_parquet_extension() {
        assert_eq!(
            detect_format(Path::new("data.parquet")).unwrap(),
            FileFormat::Parquet
        );
        assert_eq!(
            detect_format(Path::new("data.pq")).unwrap(),
            FileFormat::Parquet
        );
    }

    #[test]
    fn rejects_unknown_extension() {
        assert!(detect_format(Path::new("data.jsonl")).is_err());
    }

    #[test]
    fn loads_csv_preview_with_limit() {
        let stamp = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("dap-test-{stamp}.csv"));
        fs::write(&path, "id,name\n1,Ada\n2,Grace\n3,Katherine\n").unwrap();

        let preview = load_preview(&path, Some(2)).unwrap();
        fs::remove_file(&path).unwrap();

        assert_eq!(preview.headers, vec!["id", "name"]);
        assert_eq!(preview.row_count(), 2);
        assert_eq!(preview.value(1, 1), "Grace");
    }

    #[test]
    fn loads_entire_small_csv_when_limit_is_omitted() {
        let stamp = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("dap-test-full-{stamp}.csv"));
        fs::write(&path, "id,name\n1,Ada\n2,Grace\n3,Katherine\n").unwrap();

        let preview = load_preview(&path, None).unwrap();
        fs::remove_file(&path).unwrap();

        assert_eq!(preview.row_count(), 3);
        assert_eq!(preview.value(2, 1), "Katherine");
    }

    #[test]
    fn lazily_loads_csv_past_initial_chunk() {
        let stamp = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("dap-test-lazy-{stamp}.csv"));
        let mut contents = String::from("id,name\n");
        for idx in 0..1005 {
            contents.push_str(&format!("{idx},name-{idx}\n"));
        }
        fs::write(&path, contents).unwrap();

        let mut preview = load_preview(&path, None).unwrap();
        assert_eq!(preview.row_count(), 1000);
        assert!(!preview.exhausted);

        let added = preview.load_more().unwrap();
        fs::remove_file(&path).unwrap();

        assert_eq!(added, 5);
        assert_eq!(preview.row_count(), 1005);
        assert_eq!(preview.value(1004, 1), "name-1004");
        assert!(preview.exhausted);
    }

    #[test]
    fn lazy_load_respects_limit() {
        let stamp = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("dap-test-limit-{stamp}.csv"));
        let mut contents = String::from("id,name\n");
        for idx in 0..2000 {
            contents.push_str(&format!("{idx},name-{idx}\n"));
        }
        fs::write(&path, contents).unwrap();

        let mut preview = load_preview(&path, Some(1500)).unwrap();
        assert_eq!(preview.row_count(), 1000);

        let added = preview.load_more().unwrap();
        fs::remove_file(&path).unwrap();

        assert_eq!(added, 500);
        assert_eq!(preview.row_count(), 1500);
        assert!(preview.exhausted);
    }

    #[test]
    fn evicted_rows_can_be_reloaded_from_disk() {
        let stamp = SystemTime::now()
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("dap-test-cache-{stamp}.csv"));
        let mut contents = String::from("id,name\n");
        for idx in 0..45_000 {
            contents.push_str(&format!("{idx},name-{idx}\n"));
        }
        fs::write(&path, contents).unwrap();

        let mut preview = load_preview(&path, None).unwrap();
        for _ in 0..5 {
            preview.load_more().unwrap();
        }

        preview.ensure_row_loaded(10).unwrap();
        fs::remove_file(&path).unwrap();

        assert_eq!(preview.value(10, 1), "name-10");
    }
}
