use std::path::{Path, PathBuf};

use bt_kernel::{BtError, Result, SymbolId, TimeRange, Timestamp, NANOS_PER_SECOND};
use chrono::{Datelike, TimeZone, Utc};

pub trait PartitionScheme: Send + Sync {
    fn partition_path(&self, dataset: &str, symbol: SymbolId, ts: Timestamp) -> PathBuf;

    fn list_partitions(
        &self,
        dataset: &str,
        symbol: SymbolId,
        range: TimeRange,
    ) -> Result<Vec<PathBuf>>;
}

#[derive(Clone, Debug)]
pub struct DefaultPartitionScheme {
    root: PathBuf,
}

impl DefaultPartitionScheme {
    pub fn new(root: impl AsRef<Path>) -> Self {
        Self {
            root: root.as_ref().to_path_buf(),
        }
    }

    pub fn root(&self) -> &Path {
        &self.root
    }
}

impl PartitionScheme for DefaultPartitionScheme {
    fn partition_path(&self, dataset: &str, symbol: SymbolId, ts: Timestamp) -> PathBuf {
        let dt = match timestamp_to_datetime(ts) {
            Ok(value) => value,
            Err(_) => Utc.timestamp_opt(0, 0).single().expect("unix epoch exists"),
        };

        self.root
            .join(dataset)
            .join(symbol.0.to_string())
            .join(format!("{:04}", dt.year()))
            .join(format!("{:02}", dt.month()))
            .join(format!("{:02}.parquet", dt.day()))
    }

    fn list_partitions(
        &self,
        dataset: &str,
        symbol: SymbolId,
        range: TimeRange,
    ) -> Result<Vec<PathBuf>> {
        if range.end <= range.start {
            return Ok(Vec::new());
        }

        let base = self.root.join(dataset).join(symbol.0.to_string());

        let start_dt = timestamp_to_datetime(range.start)?;
        let last_nanos = range.end.as_nanos().checked_sub(1).ok_or_else(|| {
            BtError::TimeRange("range end underflow while listing partitions".to_string())
        })?;
        let end_dt = timestamp_to_datetime(Timestamp(last_nanos))?;

        let start_ym = (start_dt.year(), start_dt.month());
        let end_ym = (end_dt.year(), end_dt.month());

        // Iterate month-by-month. Prefer monthly file over daily files.
        let mut paths = Vec::new();
        let mut y = start_ym.0;
        let mut m = start_ym.1;

        while (y, m) <= end_ym {
            let month_dir = base.join(format!("{y:04}")).join(format!("{m:02}"));

            // Prefer Arrow IPC (.arrow) > Parquet (.parquet) > daily files
            let arrow_path = base.join(format!("{y:04}")).join(format!("{m:02}.arrow"));
            let parquet_path = base.join(format!("{y:04}")).join(format!("{m:02}.parquet"));
            if arrow_path.exists() {
                paths.push(arrow_path);
            } else if parquet_path.exists() {
                paths.push(parquet_path);
            } else {
                // Fallback: enumerate daily files within this month.
                let first_day = if (y, m) == start_ym {
                    start_dt.day()
                } else {
                    1
                };
                let last_day = if (y, m) == end_ym {
                    end_dt.day()
                } else {
                    days_in_month(y, m)
                };
                for d in first_day..=last_day {
                    let path = month_dir.join(format!("{d:02}.parquet"));
                    if path.exists() {
                        paths.push(path);
                    }
                }
            }

            // Advance to next month.
            m += 1;
            if m > 12 {
                m = 1;
                y += 1;
            }
        }

        Ok(paths)
    }
}

pub(crate) fn timestamp_to_datetime(ts: Timestamp) -> Result<chrono::DateTime<Utc>> {
    let secs = ts.as_nanos().div_euclid(NANOS_PER_SECOND);
    let nanos = ts.as_nanos().rem_euclid(NANOS_PER_SECOND) as u32;

    Utc.timestamp_opt(secs, nanos)
        .single()
        .ok_or_else(|| BtError::Data(format!("invalid timestamp nanos: {}", ts.as_nanos())))
}

fn days_in_month(year: i32, month: u32) -> u32 {
    match month {
        1 | 3 | 5 | 7 | 8 | 10 | 12 => 31,
        4 | 6 | 9 | 11 => 30,
        2 => {
            if (year % 4 == 0 && year % 100 != 0) || year % 400 == 0 {
                29
            } else {
                28
            }
        }
        _ => 31,
    }
}

pub(crate) fn datetime_to_timestamp(dt: chrono::DateTime<Utc>) -> Result<Timestamp> {
    let nanos = dt
        .timestamp()
        .checked_mul(NANOS_PER_SECOND)
        .and_then(|v| v.checked_add(i64::from(dt.timestamp_subsec_nanos())))
        .ok_or_else(|| BtError::Data("timestamp overflow".to_string()))?;

    Ok(Timestamp(nanos))
}
