sxd_html_table 0.2.0

use std::collections::HashMap;

use sxd_xpath::{nodeset::Node, Context, Factory, Value};

use crate::{element_utils, table::Table, Error};

pub(crate) const MAX_TABLE_COLUMNS: usize = 1000;

impl Table<Node<'_>> {
    pub fn to_string_table(&self) -> Table<String> {
        self.map(|_, _, node| node.string_value())
    }

    pub fn to_string_table_with_header(&self) -> Table<(String, bool)> {
        self.map(|_, _, node| {
            let Some(element) = node.element() else {
                return (node.string_value(), false);
            };
            let is_header = element.name() == "th".into();
            (node.string_value(), is_header)
        })
    }
}

struct TableSupport<'a>(Node<'a>);

impl<'a> TableSupport<'a> {
    fn tr_nodes(&self) -> Result<Vec<Node<'a>>, Error> {
        let tr_nodes = match evaluate_xpath_node(self.0, "./thead/tr|./tbody/tr|./tfoot/tr") {
            Ok(Value::Nodeset(tr_nodes)) => tr_nodes,
            _ => {
                return Err(Error::InvalidDocument(
                    "XPath ./tbody/tr did not return a nodeset",
                ))
            }
        };
        Ok(tr_nodes.document_order())
    }
    fn td_nodes(&self, tr: Node<'a>) -> Result<Vec<Node<'a>>, Error> {
        let td_nodes = match evaluate_xpath_node(tr, "./td|./th") {
            Ok(Value::Nodeset(td_nodes)) => td_nodes,
            _ => {
                return Err(Error::InvalidDocument(
                    "XPath ./td|./th did not return a nodeset",
                ))
            }
        };
        Ok(td_nodes.document_order())
    }
}

pub fn evaluate_xpath_node<'a>(
    node: impl Into<Node<'a>>,
    expr: &str,
) -> Result<Value<'a>, sxd_xpath::Error> {
    let factory = Factory::new();
    let expression = factory.build(expr)?;
    let expression = expression.ok_or(sxd_xpath::Error::NoXPath)?;
    let context = Context::new();
    expression
        .evaluate(&context, node.into())
        .map_err(Into::into)
}

fn extract_table_nodes<'a>(node: impl Into<Node<'a>>) -> Result<Vec<Node<'a>>, Error> {
    let val = evaluate_xpath_node(node, "//table[not(ancestor::table)]").map_err(Error::from)?;
    let Value::Nodeset(table_nodes) = val else {
        unreachable!("//table XPath always returns a Nodeset");
    };
    let nodes = table_nodes.document_order();
    if nodes.is_empty() {
        return Err(Error::TableNotFound);
    }
    Ok(nodes)
}

pub fn extract_table_nodes_to_table<'a>(
    node: impl Into<Node<'a>>,
) -> Result<Vec<Table<Node<'a>>>, Error> {
    let mut tables = vec![];
    for node in extract_table_nodes(node)? {
        tables.push(node_to_table(node)?);
    }
    Ok(tables)
}

fn node_to_table<'a>(node: impl Into<Node<'a>>) -> Result<Table<Node<'a>>, Error> {
    let mut map: HashMap<(usize, usize), Node> = HashMap::new();
    let t = TableSupport(node.into());
    let tr_nodes = t.tr_nodes()?;
    for (row_index, tr_node) in tr_nodes.iter().enumerate() {
        let mut rowspan_zero_seen = false;
        for td_node in t.td_nodes(*tr_node)? {
            let mut col_index = 0;
            let Some(element) = td_node.element() else {
                return Err(Error::InvalidDocument(
                    "td/th node could not be cast to an element",
                ));
            };
            let (mut row_size, col_size) = element_utils::extract_rowspan_and_colspan(element);
            let remaining_rows = tr_nodes.len() - row_index;
            if row_size == 0 {
                if rowspan_zero_seen {
                    row_size = 1;
                } else {
                    row_size = remaining_rows;
                    rowspan_zero_seen = true;
                }
            } else if row_size > remaining_rows {
                row_size = remaining_rows;
            }
            while col_index < MAX_TABLE_COLUMNS && map.contains_key(&(row_index, col_index)) {
                col_index += 1;
            }
            if col_index >= MAX_TABLE_COLUMNS {
                return Err(Error::InvalidDocument(
                    "table row exceeds the maximum number of columns",
                ));
            }
            if col_size > MAX_TABLE_COLUMNS {
                return Err(Error::InvalidDocument(
                    "colspan exceeds the maximum number of columns",
                ));
            }
            let Some(end_col) = col_index.checked_add(col_size) else {
                return Err(Error::InvalidDocument(
                    "column index overflowed while placing a cell",
                ));
            };
            if end_col > MAX_TABLE_COLUMNS {
                return Err(Error::InvalidDocument(
                    "cell placement exceeds the maximum number of columns",
                ));
            }
            for k in 0..row_size {
                for l in 0..col_size {
                    let row = row_index + k;
                    let col = col_index + l;
                    map.insert((row, col), td_node);
                }
            }
        }
    }
    let rows = map
        .keys()
        .map(|(i, _)| i)
        .max()
        .map(|&i| i + 1)
        .unwrap_or(0);
    let cols = map
        .keys()
        .map(|(_, j)| j)
        .max()
        .map(|&j| j + 1)
        .unwrap_or(0);
    let mut table = Table::new((rows, cols));
    for ((i, j), item) in map {
        table.set(i, j, item);
    }
    Ok(table)
}