use std::collections::{HashMap, HashSet, VecDeque};

use crate::types::Transform;

pub const DEFAULT_MAX_HISTORY: usize = 1000;

#[derive(Clone, Debug, PartialEq)]
pub enum TfError {
    FramePairNotFound {
        frame_from: String,
        frame_to: String,
    },
    NoDataInBuffer {
        frame_from: String,
        frame_to: String,
    },
    ExtrapolationError {
        frame_from: String,
        frame_to: String,
        requested_ns: u64,
        min_ns: u64,
        max_ns: u64,
    },
    NoPathFound {
        frame_from: String,
        frame_to: String,
    },
    /// e.g. fewer than 2 frames passed to `lookup_transform_chain`.
    InvalidArgument(String),
}

impl std::fmt::Display for TfError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TfError::FramePairNotFound {
                frame_from,
                frame_to,
            } => write!(f, "frame pair '{frame_from}' -> '{frame_to}' not found"),
            TfError::NoDataInBuffer {
                frame_from,
                frame_to,
            } => write!(f, "no data for '{frame_from}' -> '{frame_to}'"),
            TfError::ExtrapolationError {
                frame_from,
                frame_to,
                requested_ns,
                min_ns,
                max_ns,
            } => write!(
                f,
                "extrapolation for '{frame_from}'->''{frame_to}': \
                           {requested_ns}ns not in [{min_ns}, {max_ns}]"
            ),
            TfError::NoPathFound {
                frame_from,
                frame_to,
            } => write!(
                f,
                "no path between '{frame_from}' and '{frame_to}' in the TF graph"
            ),
            TfError::InvalidArgument(msg) => write!(f, "invalid argument: {msg}"),
        }
    }
}

type Entry = (u64, Transform);

/// Time-indexed coordinate transform buffer with interpolation.
///
/// Maintains a bounded history per `(frame_from, frame_to)` pair.
/// Supports static transforms, inverse-pair fallback, and multi-hop chain lookup.
pub struct TransformBuffer {
    max_history: usize,
    data: HashMap<(String, String), VecDeque<Entry>>,
    /// Static transforms never expire and provide a fallback for `lookup_transform`.
    static_tfs: HashMap<(String, String), Transform>,
}

impl TransformBuffer {
    pub fn new() -> Self {
        Self::with_max_history(DEFAULT_MAX_HISTORY)
    }

    pub fn with_max_history(max_history: usize) -> Self {
        Self {
            max_history: max_history.max(1),
            data: HashMap::new(),
            static_tfs: HashMap::new(),
        }
    }

    pub fn register_static(
        &mut self,
        frame_from: impl Into<String>,
        frame_to: impl Into<String>,
        transform: Transform,
    ) {
        self.static_tfs
            .insert((frame_from.into(), frame_to.into()), transform);
    }

    pub fn broadcast(
        &mut self,
        frame_from: impl Into<String>,
        frame_to: impl Into<String>,
        transform: Transform,
        timestamp_ns: u64,
    ) {
        let key = (frame_from.into(), frame_to.into());
        let entries = self.data.entry(key).or_default();
        let pos = entries.partition_point(|(ts, _)| *ts < timestamp_ns);
        if pos < entries.len() && entries[pos].0 == timestamp_ns {
            entries[pos].1 = transform;
        } else {
            entries.insert(pos, (timestamp_ns, transform));
            while entries.len() > self.max_history {
                entries.pop_front();
            }
        }
    }

    /// Broadcast via a `StampedTransform`.
    pub fn broadcast_stamped(&mut self, stamped: crate::types::StampedTransform) {
        self.broadcast(
            stamped.frame_from,
            stamped.frame_to,
            stamped.transform,
            stamped.timestamp_ns,
        );
    }

    /// Lookup with interpolation. Falls back to static transforms if no dynamic pair found.
    pub fn lookup_transform(
        &self,
        frame_from: &str,
        frame_to: &str,
        timestamp_ns: u64,
    ) -> Result<Transform, TfError> {
        let key = (frame_from.to_string(), frame_to.to_string());
        if let Some(entries) = self.data.get(&key) {
            return self.interpolate(frame_from, frame_to, entries, timestamp_ns);
        }
        // Static fallback.
        if let Some(tf) = self.static_tfs.get(&key) {
            return Ok(tf.clone());
        }
        Err(TfError::FramePairNotFound {
            frame_from: frame_from.to_string(),
            frame_to: frame_to.to_string(),
        })
    }

    /// Like `lookup_transform`, but also tries the inverted pair if direct lookup fails.
    pub fn lookup_transform_or_invert(
        &self,
        frame_from: &str,
        frame_to: &str,
        timestamp_ns: u64,
    ) -> Result<Transform, TfError> {
        match self.lookup_transform(frame_from, frame_to, timestamp_ns) {
            Ok(tf) => Ok(tf),
            Err(TfError::FramePairNotFound { .. }) => {
                // Try reversed pair and invert the result.
                self.lookup_transform(frame_to, frame_from, timestamp_ns)
                    .map(|tf| tf.inverse())
            }
            Err(e) => Err(e),
        }
    }

    /// Compose transforms along a frame chain.
    /// `frames = ["a", "b", "c"]` computes `T_{a→b} · T_{b→c}`.
    /// Requires at least 2 frames.
    pub fn lookup_transform_chain(
        &self,
        frames: &[&str],
        timestamp_ns: u64,
    ) -> Result<Transform, TfError> {
        if frames.len() < 2 {
            return Err(TfError::InvalidArgument(format!(
                "chain needs ≥2 frames, got {}",
                frames.len()
            )));
        }
        let mut result = Transform::identity();
        for pair in frames.windows(2) {
            let tf = self.lookup_transform(pair[0], pair[1], timestamp_ns)?;
            result = result.compose(&tf);
        }
        Ok(result)
    }

    pub fn lookup_transform_any(
        &self,
        frame_from: &str,
        frame_to: &str,
        timestamp_ns: u64,
    ) -> Result<Transform, TfError> {
        match self.lookup_transform_or_invert(frame_from, frame_to, timestamp_ns) {
            Err(TfError::FramePairNotFound { .. }) => {}
            result => return result,
        }
        self.bfs_lookup(frame_from, frame_to, timestamp_ns)
    }

    pub fn known_frames(&self) -> Vec<String> {
        let mut frames: HashSet<String> = HashSet::new();
        for (a, b) in self.data.keys().chain(self.static_tfs.keys()) {
            frames.insert(a.clone());
            frames.insert(b.clone());
        }
        frames.into_iter().collect()
    }

    fn build_adjacency(&self) -> HashMap<String, Vec<String>> {
        let mut adj: HashMap<String, Vec<String>> = HashMap::new();
        for (from, to) in self.data.keys().chain(self.static_tfs.keys()) {
            adj.entry(from.clone()).or_default().push(to.clone());
            adj.entry(to.clone()).or_default().push(from.clone());
        }
        adj
    }

    fn bfs_lookup(
        &self,
        frame_from: &str,
        frame_to: &str,
        timestamp_ns: u64,
    ) -> Result<Transform, TfError> {
        let adj = self.build_adjacency();
        let start = frame_from.to_string();
        let goal = frame_to.to_string();

        let mut parent: HashMap<String, String> = HashMap::new();
        let mut queue: VecDeque<String> = VecDeque::new();
        parent.insert(start.clone(), start.clone());
        queue.push_back(start);

        while let Some(current) = queue.pop_front() {
            if current == goal {
                let mut path = Vec::new();
                let mut node = goal.clone();
                loop {
                    let p = parent[&node].clone();
                    path.push(node.clone());
                    if p == node {
                        break;
                    }
                    node = p;
                }
                path.reverse();
                let mut result = Transform::identity();
                for pair in path.windows(2) {
                    let tf = self.lookup_transform_or_invert(&pair[0], &pair[1], timestamp_ns)?;
                    result = result.compose(&tf);
                }
                return Ok(result);
            }
            for neighbor in adj.get(&current).into_iter().flatten() {
                if !parent.contains_key(neighbor.as_str()) {
                    parent.insert(neighbor.clone(), current.clone());
                    queue.push_back(neighbor.clone());
                }
            }
        }

        Err(TfError::NoPathFound {
            frame_from: frame_from.to_string(),
            frame_to: frame_to.to_string(),
        })
    }

    /// Delete all dynamic entries older than `timestamp_ns` for a specific pair.
    pub fn clear_before(&mut self, frame_from: &str, frame_to: &str, timestamp_ns: u64) {
        if let Some(entries) = self
            .data
            .get_mut(&(frame_from.to_string(), frame_to.to_string()))
        {
            entries.retain(|(ts, _)| *ts >= timestamp_ns);
        }
    }

    /// Delete all dynamic entries older than `timestamp_ns` across all pairs.
    pub fn gc_before(&mut self, timestamp_ns: u64) {
        for entries in self.data.values_mut() {
            entries.retain(|(ts, _)| *ts >= timestamp_ns);
        }
        self.data.retain(|_, v| !v.is_empty());
    }

    pub fn frame_pair_count(&self) -> usize {
        self.data.len()
    }
    pub fn entry_count(&self, frame_from: &str, frame_to: &str) -> usize {
        self.data
            .get(&(frame_from.to_string(), frame_to.to_string()))
            .map_or(0, |v| v.len())
    }

    fn interpolate(
        &self,
        frame_from: &str,
        frame_to: &str,
        entries: &VecDeque<Entry>,
        timestamp_ns: u64,
    ) -> Result<Transform, TfError> {
        if entries.is_empty() {
            return Err(TfError::NoDataInBuffer {
                frame_from: frame_from.to_string(),
                frame_to: frame_to.to_string(),
            });
        }
        let min_ns = entries.front().unwrap().0;
        let max_ns = entries.back().unwrap().0;
        if timestamp_ns < min_ns || timestamp_ns > max_ns {
            return Err(TfError::ExtrapolationError {
                frame_from: frame_from.to_string(),
                frame_to: frame_to.to_string(),
                requested_ns: timestamp_ns,
                min_ns,
                max_ns,
            });
        }
        let hi = entries.partition_point(|(ts, _)| *ts <= timestamp_ns);
        if entries[hi - 1].0 == timestamp_ns {
            return Ok(entries[hi - 1].1.clone());
        }
        let (t0, tf0) = &entries[hi - 1];
        let (t1, tf1) = &entries[hi];
        let t = (timestamp_ns - t0) as f64 / (t1 - t0) as f64;
        Ok(tf0.interpolate(tf1, t))
    }
}

impl Default for TransformBuffer {
    fn default() -> Self {
        Self::new()
    }
}
