//! Print ASCII tables to the terminal.
//!
//! # Example
//!
//! ```
//! use ascii_table::AsciiTable;
//!
//! let ascii_table = AsciiTable::default();
//! let data = vec![&[1, 2, 3], &[4, 5, 6], &[7, 8, 9]];
//! ascii_table.print(data);
//! // ┌───┬───┬───┐
//! // │ 1 │ 2 │ 3 │
//! // │ 4 │ 5 │ 6 │
//! // │ 7 │ 8 │ 9 │
//! // └───┴───┴───┘
//! ```
//!
//! # Example
//!
//! ```
//! use std::fmt::Display;
//! use ascii_table::{AsciiTable, Align};
//!
//! let mut ascii_table = AsciiTable::default();
//! ascii_table.set_max_width(26);
//! ascii_table.column(0).set_header("H1").set_align(Align::Left);
//! ascii_table.column(1).set_header("H2").set_align(Align::Center);
//! ascii_table.column(2).set_header("H3").set_align(Align::Right);
//!
//! let data: Vec<Vec<&dyn Display>> = vec![
//!     vec![&'v', &'v', &'v'],
//!     vec![&123, &456, &789, &"abcdef"]
//! ];
//! ascii_table.print(data);
//! // ┌─────┬─────┬─────┬──────┐
//! // │ H1  │ H2  │ H3  │      │
//! // ├─────┼─────┼─────┼──────┤
//! // │ v   │  v  │   v │      │
//! // │ 123 │ 456 │ 789 │ abc+ │
//! // └─────┴─────┴─────┴──────┘
//! ```
//!
//! ## Features
//!
//! - `auto_table_width`: Sets the default max width of the ascii table to the width of the terminal.
//! - `color_codes`: Correctly calculates the width of a string when terminal color codes are present
//!   (like those from the `colorful` crate).
//! - `wide_characters`: Correctly calculates the width of a string when wide characters are present
//!   (like emoli's).

#[cfg(test)]
mod test;

#[cfg(feature = "color_codes")]
use lazy_static::lazy_static;
#[cfg(feature = "color_codes")]
use regex::Regex;
#[cfg(feature = "wide_characters")]
use unicode_width::UnicodeWidthStr;

use std::collections::BTreeMap;
use std::fmt::Display;

const SE: &str = "┌";
const NW: &str = "┘";
const SW: &str = "┐";
const NS: &str = "│";
const NE: &str = "└";
const EWS: &str = "┬";
const NES: &str = "├";
const NWS: &str = "┤";
const NEW: &str = "┴";
const NEWS: &str = "┼";
const EW: &str = "─";

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct AsciiTable {
    max_width: usize,
    columns: BTreeMap<usize, Column>,
}

impl AsciiTable {
    pub fn set_max_width(&mut self, max_width: usize) -> &mut Self {
        self.max_width = max_width;
        self
    }

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

    pub fn column(&mut self, index: usize) -> &mut Column {
        self.columns.entry(index).or_default()
    }
}

impl Default for AsciiTable {
    fn default() -> Self {
        Self {
            max_width: default_table_width(),
            columns: BTreeMap::new(),
        }
    }
}

#[cfg(feature = "auto_table_width")]
fn default_table_width() -> usize {
    termion::terminal_size()
        .map(|(width, _)| {
            let width: usize = width.into();
            width
        })
        .ok()
        .unwrap_or(100)
}

#[cfg(not(feature = "auto_table_width"))]
fn default_table_width() -> usize {
    100
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Column {
    header: String,
    align: Align,
    max_width: usize,
}

impl Column {
    pub fn set_header<T>(&mut self, header: T) -> &mut Self
    where
        T: Into<String>,
    {
        self.header = header.into();
        self
    }

    pub fn header(&self) -> &str {
        &self.header
    }

    pub fn set_align(&mut self, align: Align) -> &mut Self {
        self.align = align;
        self
    }

    pub fn align(&self) -> Align {
        self.align
    }

    pub fn set_max_width(&mut self, max_width: usize) -> &mut Self {
        self.max_width = max_width;
        self
    }

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

impl Default for Column {
    fn default() -> Self {
        Column {
            header: String::new(),
            align: Default::default(),
            max_width: usize::max_value(),
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub enum Align {
    Left,
    Center,
    Right,
}

impl Default for Align {
    fn default() -> Self {
        Align::Left
    }
}

impl AsciiTable {
    pub fn print<L1, L2, T>(&self, data: L1)
    where
        L1: IntoIterator<Item = L2>,
        L2: IntoIterator<Item = T>,
        T: Display,
    {
        print!("{}", self.format(data))
    }

    pub fn format<L1, L2, T>(&self, data: L1) -> String
    where
        L1: IntoIterator<Item = L2>,
        L2: IntoIterator<Item = T>,
        T: Display,
    {
        self.format_inner(self.stringify(data))
    }

    fn format_inner(&self, data: Vec<Vec<SmartString>>) -> String {
        let num_cols = data.iter().map(|row| row.len()).max().unwrap_or(0);
        if !self.valid(&data, num_cols) {
            return self.format_empty();
        }

        let header = self.stringify_header(num_cols);
        let data = self.square_data(data, num_cols);
        let has_header = header.iter().any(|text| !text.is_empty());
        let widths = self.column_widths(&header, &data, num_cols);

        let mut result = String::new();
        result.push_str(&self.format_first(&widths));
        if has_header {
            result.push_str(&self.format_header_row(&header, &widths));
            result.push_str(&self.format_middle(&widths));
        }
        for row in data {
            result.push_str(&self.format_row(&row, &widths));
        }
        result.push_str(&self.format_last(&widths));
        result
    }

    fn valid(&self, data: &Vec<Vec<SmartString>>, num_cols: usize) -> bool {
        if data.len() == 0 {
            false
        } else if num_cols == 0 {
            false
        } else if self.max_width < Self::smallest_width(num_cols) {
            false
        } else {
            true
        }
    }

    fn smallest_width(num_cols: usize) -> usize {
        ((num_cols - 1) * 3) + 4
    }

    fn stringify<L1, L2, T>(&self, data: L1) -> Vec<Vec<SmartString>>
    where
        L1: IntoIterator<Item = L2>,
        L2: IntoIterator<Item = T>,
        T: Display,
    {
        data.into_iter()
            .map(|row| {
                row.into_iter()
                    .map(|cell| SmartString::from(cell.to_string()))
                    .collect()
            })
            .collect()
    }

    fn stringify_header(&self, num_cols: usize) -> Vec<SmartString> {
        let default_conf = &Default::default();
        (0..num_cols)
            .map(|a| SmartString::from(self.columns.get(&a).unwrap_or(default_conf).header.clone()))
            .collect()
    }

    fn square_data(
        &self,
        mut data: Vec<Vec<SmartString>>,
        num_cols: usize,
    ) -> Vec<Vec<SmartString>> {
        for row in data.iter_mut() {
            while row.len() < num_cols {
                row.push(SmartString::new())
            }
        }
        data
    }

    fn column_widths(
        &self,
        header: &[SmartString],
        data: &[Vec<SmartString>],
        num_cols: usize,
    ) -> Vec<usize> {
        let result: Vec<_> = (0..num_cols)
            .map(|a| {
                let default_conf = &Default::default();
                let conf = self.columns.get(&a).unwrap_or(default_conf);
                let column_width = data.iter().map(|row| row[a].width()).max().unwrap();
                let header_width = header[a].width();
                column_width.max(header_width).min(conf.max_width)
            })
            .collect();
        self.truncate_widths(result)
    }

    fn truncate_widths(&self, mut widths: Vec<usize>) -> Vec<usize> {
        let max_width = self.max_width;
        let table_padding = Self::smallest_width(widths.len());
        while widths.iter().sum::<usize>() + table_padding > max_width
            && *widths.iter().max().unwrap() > 0
        {
            let max = widths.iter().max().unwrap();
            let idx = widths.iter().rposition(|x| x == max).unwrap();
            widths[idx] -= 1;
        }
        widths
    }

    fn format_line(&self, row: &[SmartString], head: &str, delim: &str, tail: &str) -> String {
        let mut result = String::new();
        result.push_str(head);
        for cell in row {
            result.push_str(&format!("{}{}", cell, delim));
        }
        for _ in 0..delim.chars().count() {
            result.pop();
        }
        result.push_str(tail);
        result.push('\n');
        result
    }

    fn format_empty(&self) -> String {
        self.format_first(&vec![0])
            + &self.format_line(
                &[SmartString::new()],
                &format!("{}{}", NS, ' '),
                &format!("{}{}{}", ' ', NS, ' '),
                &format!("{}{}", ' ', NS),
            )
            + &self.format_last(&[0])
    }

    fn format_first(&self, widths: &[usize]) -> String {
        let row: Vec<_> = widths
            .iter()
            .map(|&x| SmartString::from_visible(EW.repeat(x)))
            .collect();
        self.format_line(
            &row,
            &format!("{}{}", SE, EW),
            &format!("{}{}{}", EW, EWS, EW),
            &format!("{}{}", EW, SW),
        )
    }

    fn format_middle(&self, widths: &[usize]) -> String {
        let row: Vec<_> = widths
            .iter()
            .map(|&x| SmartString::from_visible(EW.repeat(x)))
            .collect();
        self.format_line(
            &row,
            &format!("{}{}", NES, EW),
            &format!("{}{}{}", EW, NEWS, EW),
            &format!("{}{}", EW, NWS),
        )
    }

    fn format_row(&self, row: &[SmartString], widths: &[usize]) -> String {
        let row: Vec<_> = (0..widths.len())
            .map(|a| {
                let cell = &row[a];
                let width = widths[a];
                let default_conf = &Default::default();
                let conf = self.columns.get(&a).unwrap_or(default_conf);
                self.format_cell(cell, width, ' ', conf.align)
            })
            .collect();
        self.format_line(
            &row,
            &format!("{}{}", NS, ' '),
            &format!("{}{}{}", ' ', NS, ' '),
            &format!("{}{}", ' ', NS),
        )
    }

    fn format_header_row(&self, row: &[SmartString], widths: &[usize]) -> String {
        let row: Vec<_> = row
            .iter()
            .zip(widths.iter())
            .map(|(cell, &width)| self.format_cell(cell, width, ' ', Align::Left))
            .collect();
        self.format_line(
            &row,
            &format!("{}{}", NS, ' '),
            &format!("{}{}{}", ' ', NS, ' '),
            &format!("{}{}", ' ', NS),
        )
    }

    fn format_last(&self, widths: &[usize]) -> String {
        let row: Vec<_> = widths
            .iter()
            .map(|&x| SmartString::from_visible(EW.repeat(x)))
            .collect();
        self.format_line(
            &row,
            &format!("{}{}", NE, EW),
            &format!("{}{}{}", EW, NEW, EW),
            &format!("{}{}", EW, NW),
        )
    }

    fn format_cell(&self, text: &SmartString, len: usize, pad: char, align: Align) -> SmartString {
        if text.width() > len {
            let mut result = text.clone();
            while result.width() > len {
                result.pop();
            }
            if result.pop().is_some() {
                result.push_visible('+')
            }
            result
        } else {
            let mut result = text.clone();
            match align {
                Align::Left => {
                    while result.width() < len {
                        result.push_visible(pad)
                    }
                }
                Align::Right => {
                    while result.width() < len {
                        result.lpush_visible(pad)
                    }
                }
                Align::Center => {
                    while result.width() < len {
                        result.push_visible(pad);
                        if result.width() < len {
                            result.lpush_visible(pad)
                        }
                    }
                }
            }
            result
        }
    }
}

#[cfg(feature = "color_codes")]
lazy_static! {
    static ref COLOR_CODE_PARSER: Regex =
        Regex::new("\u{1b}\\[([0-9]+;)*[0-9]+m").expect("Regex compilation error");
}

#[derive(Clone, Debug)]
struct SmartString {
    fragments: Vec<SmartStringFragment>,
}

#[derive(Clone, Debug)]
struct SmartStringFragment {
    string: String,
    visible: bool,
}

impl SmartString {
    fn new() -> Self {
        Self {
            fragments: Vec::new(),
        }
    }

    #[cfg(feature = "color_codes")]
    fn from(string: String) -> Self {
        let mut fragments = Vec::new();
        let mut last = 0;
        for r#match in COLOR_CODE_PARSER.find_iter(&string) {
            let start = r#match.start();
            let end = r#match.end();

            if last < start {
                fragments.push(SmartStringFragment::new(&string[last..start], true));
            }
            fragments.push(SmartStringFragment::new(&string[start..end], false));

            last = end;
        }

        if last < string.len() {
            fragments.push(SmartStringFragment::new(&string[last..], true));
        }

        Self { fragments }
    }

    #[cfg(not(feature = "color_codes"))]
    fn from(string: String) -> Self {
        Self {
            fragments: vec![SmartStringFragment::new(string, true)],
        }
    }

    fn from_visible(string: String) -> Self {
        Self {
            fragments: vec![SmartStringFragment::new(string, true)],
        }
    }

    fn width(&self) -> usize {
        self.fragments
            .iter()
            .filter(|fragment| fragment.visible)
            .map(|fragment| fragment.width())
            .sum()
    }

    fn is_empty(&self) -> bool {
        self.fragments
            .iter()
            .filter(|fragment| fragment.visible)
            .all(|fragment| fragment.string.is_empty())
    }

    fn pop(&mut self) -> Option<char> {
        self.fragments
            .iter_mut()
            .filter(|fragment| fragment.visible && !fragment.string.is_empty())
            .last()
            .and_then(|fragment| fragment.string.pop())
    }

    fn push_visible(&mut self, ch: char) {
        let last_fragment = self
            .fragments
            .iter_mut()
            .filter(|fragment| fragment.visible)
            .map(|fragment| &mut fragment.string)
            .last();
        if let Some(fragment) = last_fragment {
            fragment.push(ch);
        } else {
            self.fragments.push(SmartStringFragment::new(ch, true));
        }
    }

    fn lpush_visible(&mut self, ch: char) {
        let first_fragment = self
            .fragments
            .iter_mut()
            .filter(|fragment| fragment.visible)
            .map(|fragment| &mut fragment.string)
            .next();
        if let Some(fragment) = first_fragment {
            fragment.insert(0, ch);
        } else {
            self.fragments.insert(0, SmartStringFragment::new(ch, true));
        }
    }
}

impl Display for SmartString {
    fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> Result<(), ::std::fmt::Error> {
        self.fragments
            .iter()
            .try_for_each(|fragment| fragment.string.fmt(fmt))
    }
}

impl SmartStringFragment {
    fn new<T>(string: T, visible: bool) -> Self
    where
        T: Into<String>,
    {
        Self {
            string: string.into(),
            visible,
        }
    }

    #[cfg(feature = "wide_characters")]
    fn width(&self) -> usize {
        UnicodeWidthStr::width(self.string.as_str())
    }

    #[cfg(not(feature = "wide_characters"))]
    fn width(&self) -> usize {
        self.string.chars().count()
    }
}
