use crate::buffer::BufferID;
use crate::state::EditorState;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt::Display;
use thiserror::Error;

#[derive(Debug, Eq, PartialEq, PartialOrd, Ord, Clone, Copy)]
pub enum Rank {
    Zero,
    One,
    Two,
    Three,
}

impl Display for Rank {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Rank::Zero => write!(f, "Rank 0"),
            Rank::One => write!(f, "Rank 1"),
            Rank::Two => write!(f, "Rank 2"),
            Rank::Three => write!(f, "Rank 3"),
        }
    }
}

#[derive(Error, Debug)]
pub enum SelectionError {
    #[error("Cannnot change selection of {from} to {to}")]
    InvalidPromotion { from: Rank, to: Rank },
}

#[derive(Debug, Clone)]
pub struct Interval {
    pub start: usize,
    pub end: usize,
    capture_groups: Vec<(usize, usize)>,
}

impl Interval {
    pub fn new(start: usize, end: usize) -> Self {
        Self {
            start,
            end,
            capture_groups: Vec::new(),
        }
    }

    pub fn is_disjoint(&self, other: &Interval) -> bool {
        self.start >= other.end || self.end <= other.start
    }
}

#[derive(Debug, Clone)]
pub enum Selection {
    Scalar {
        buffer_id: BufferID,
        pos: usize,
    },
    Vector {
        buffer_id: BufferID,
        interval: Interval,
    },
    Vectors {
        buffer_id: BufferID,
        ranges: Vec<Interval>,
    },
    MultiVectors {
        multi_ranges: HashMap<BufferID, Vec<Interval>>,
    },
}

/// A type that knows how to act on the rank-0 (scalar) and rank-1 (interval)
/// elements of a [`Selection`], and so can be vectorised over one.
pub trait Vectorisable<P, E> {
    /// Action on a rank-0 element.
    fn rank0(
        &self,
        param: &P,
        es: &EditorState,
        buffer_id: BufferID,
        pos: usize,
    ) -> Result<Selection, E>;

    /// Action on a rank-1 element.
    fn rank1(
        &self,
        param: &P,
        es: &EditorState,
        buffer_id: BufferID,
        interval: &Interval,
    ) -> Result<Selection, E>;

    /// Vectorise `sel`: apply `rank0`/`rank1` to every element and union
    /// the results.
    fn vectorise(&self, param: &P, es: &EditorState, sel: &Selection) -> Result<Selection, E> {
        // Single rank-0/rank-1 elements pass through as-is (no rank
        // promotion); multi-element selections are unioned.
        match sel {
            Selection::Scalar { buffer_id, pos } => Self::rank0(self, param, es, *buffer_id, *pos),
            Selection::Vector {
                buffer_id,
                interval,
            } => Self::rank1(self, param, es, *buffer_id, interval),
            Selection::Vectors { buffer_id, ranges } => {
                let mut selections = Vec::with_capacity(ranges.len());
                for interval in ranges {
                    selections.push(Self::rank1(self, param, es, *buffer_id, interval)?);
                }
                Ok(Selection::union(selections))
            }
            Selection::MultiVectors { multi_ranges } => {
                let mut selections = Vec::new();
                for (buffer_id, ranges) in multi_ranges {
                    for interval in ranges {
                        selections.push(Self::rank1(self, param, es, *buffer_id, interval)?);
                    }
                }
                Ok(Selection::union(selections))
            }
        }
    }
}

impl Selection {
    pub fn empty() -> Self {
        Self::MultiVectors {
            multi_ranges: HashMap::new(),
        }
    }

    fn normalise(&mut self) {
        if let Self::MultiVectors { multi_ranges } = self {
            multi_ranges.retain(|_, v| !v.is_empty())
        }
    }

    pub fn is_empty(&mut self) -> bool {
        self.normalise();
        match self {
            Self::Scalar { .. } => return false,
            Self::Vector { .. } => return false,
            Self::Vectors { ranges, .. } => ranges.is_empty(),
            Self::MultiVectors { multi_ranges } => multi_ranges.is_empty(),
        }
    }

    fn rank(&self) -> Rank {
        match self {
            Self::Scalar { .. } => Rank::Zero,
            Self::Vector { .. } => Rank::One,
            Self::Vectors { .. } => Rank::Two,
            Self::MultiVectors { .. } => Rank::Three,
        }
    }

    fn promote(&self, to_rank: Rank) -> Result<Self, SelectionError> {
        use Rank::*;

        match (self, to_rank) {
            (Self::Scalar { .. }, Zero) => Ok(self.clone()),
            (Self::Vector { .. }, One) => Ok(self.clone()),
            (Self::Vectors { .. }, Two) => Ok(self.clone()),
            (Self::MultiVectors { .. }, Three) => Ok(self.clone()),

            // Scalar
            (Self::Scalar { pos, buffer_id }, One) => Ok(Self::Vector {
                interval: Interval {
                    start: pos + 0,
                    end: pos + 1,
                    capture_groups: Vec::new(),
                },
                buffer_id: *buffer_id,
            }),
            (Self::Scalar { pos, buffer_id }, Two) => Ok(Self::Vectors {
                ranges: vec![Interval {
                    start: pos + 0,
                    end: pos + 1,
                    capture_groups: Vec::new(),
                }],
                buffer_id: *buffer_id,
            }),
            (Self::Scalar { pos, buffer_id }, Three) => {
                let mut map = HashMap::new();
                map.insert(
                    *buffer_id,
                    vec![Interval {
                        start: pos + 0,
                        end: pos + 1,
                        capture_groups: Vec::new(),
                    }],
                );
                Ok(Self::MultiVectors { multi_ranges: map })
            }

            // Vector promotions
            (
                Self::Vector {
                    interval,
                    buffer_id,
                },
                Two,
            ) => Ok(Self::Vectors {
                ranges: vec![interval.clone()],
                buffer_id: *buffer_id,
            }),
            (
                Self::Vector {
                    interval,
                    buffer_id,
                },
                Three,
            ) => {
                let mut map = HashMap::new();
                map.insert(*buffer_id, vec![interval.clone()]);
                Ok(Self::MultiVectors { multi_ranges: map })
            }

            // Vectors promotion
            (Self::Vectors { ranges, buffer_id }, Three) => {
                let mut map = HashMap::new();
                map.insert(*buffer_id, ranges.clone());
                Ok(Self::MultiVectors { multi_ranges: map })
            }

            // all other cases
            (sel, target) => Err(SelectionError::InvalidPromotion {
                from: sel.rank(),
                to: target,
            }),
        }
    }

    fn buffers(&self) -> Vec<BufferID> {
        match self {
            Self::Scalar { buffer_id, .. } => vec![*buffer_id],
            Self::Vector { buffer_id, .. } => vec![*buffer_id],
            Self::Vectors { buffer_id, .. } => vec![*buffer_id],
            Self::MultiVectors { multi_ranges } => multi_ranges.keys().copied().collect(),
        }
    }

    /// Union the given selections into a single selection.
    ///
    /// The target rank is the maximum rank needed to hold all inputs (rank 3
    /// if they span multiple buffers), to which every input can be promoted,
    /// so this operation cannot fail.
    pub fn union(selections: Vec<Self>) -> Self {
        if selections.is_empty() {
            return Self::empty();
        }

        let all_buffers: HashSet<BufferID> = selections.iter().flat_map(|s| s.buffers()).collect();

        if all_buffers.is_empty() {
            return Self::empty();
        }

        let max_rank = if all_buffers.len() > 1 {
            Rank::Three
        } else {
            selections
                .iter()
                .map(|s| s.rank())
                .max()
                .unwrap_or(Rank::Two)
                .max(Rank::Two) // At least rank 2
        };

        let promoted: Vec<Self> = selections
            .iter()
            .map(|s| {
                s.promote(max_rank)
                    .expect("union promotes to the max rank present, which cannot fail")
            })
            .collect();

        match max_rank {
            Rank::Two => {
                let mut all_intervals = Vec::new();
                let buffer_id = *all_buffers.iter().next().unwrap();

                for sel in &promoted {
                    match sel {
                        Self::Vectors { ranges, .. } => all_intervals.extend(ranges.clone()),
                        _ => unreachable!("promote() to Rank::Two should always produce Vectors"),
                    }
                }

                Self::Vectors {
                    ranges: all_intervals,
                    buffer_id,
                }
            }
            Rank::Three => {
                let mut buffer_intervals: HashMap<BufferID, Vec<Interval>> =
                    all_buffers.iter().map(|&b| (b, Vec::new())).collect();

                for sel in &promoted {
                    match sel {
                        Self::MultiVectors { multi_ranges } => {
                            for (&buffer, intervals) in multi_ranges {
                                buffer_intervals
                                    .entry(buffer)
                                    .or_insert_with(Vec::new)
                                    .extend(intervals.clone());
                            }
                        }
                        _ => unreachable!(
                            "promote() to Rank::Three should always produce MultiVectors"
                        ),
                    }
                }

                Self::MultiVectors {
                    multi_ranges: buffer_intervals,
                }
            }
            _ => unreachable!(),
        }
    }
}