use crate::error::Result;
use crate::font::FontRegistry;
use crate::model::TableWidth;
use crate::normalize::{NormalizedCellContent, NormalizedTable, NormalizedTableCell};
use crate::text::{TextMeasurer, wrap_inlines_with_fonts};
use super::{MeasuredTable, MeasuredTableCell, MeasuredTableRow};
use std::collections::HashSet;

pub(super) fn measure_table_block(table: &NormalizedTable, available_width: f32, measurer: &mut TextMeasurer, fonts: &FontRegistry) -> Result<MeasuredTable> {
    let (rows, column_count) = span_normalized_rows(&table.rows);
    let column_widths = resolve_column_widths(&rows, &table.widths, column_count, available_width, measurer, fonts)?;

    // 预计算 covered 位置（与 layout_table 一致的算法）
    let mut covered: HashSet<(usize, usize)> = HashSet::new();
    for (row_idx, row) in rows.iter().enumerate() {
        let mut col = 0usize;
        for cell in row {
            while col < column_count && covered.contains(&(row_idx, col)) {
                col += 1;
            }
            let row_span = cell.row_span.max(1);
            let col_span = cell.col_span.max(1);
            for r in 1..row_span {
                if row_idx + r < rows.len() {
                    for c in 0..col_span {
                        covered.insert((row_idx + r, col + c));
                    }
                }
            }
            col += col_span;
        }
    }

    let mut measured_rows = rows
        .iter()
        .enumerate()
        .map(|(row_idx, row)| measure_table_row(row, &column_widths, measurer, fonts, row_idx, &covered))
        .collect::<Result<Vec<_>>>()?;

    for row_idx in 0..measured_rows.len() {
        let row = &rows[row_idx];
        let mut col = 0usize;
        for (cell_idx, cell) in row.iter().enumerate() {
            while col < column_widths.len() && covered.contains(&(row_idx, col)) {
                col += 1;
            }

            let row_span = cell.row_span.max(1);
            let col_span = cell.col_span.max(1);

            if row_span > 1 && col < column_widths.len() {
                let cell_width = spanned_width(&column_widths, col, col_span);
                let content_width = (cell_width - cell.padding * 2.0).max(0.0);

                let (lines, items, min_content_height) = match &cell.content {
                    NormalizedCellContent::Inlines(inlines) => {
                        let l = wrap_inlines_with_fonts(inlines, content_width, measurer, fonts)?;
                        let h: f32 = l.iter().map(|li| li.line_height).sum();
                        (l, Vec::new(), h)
                    }
                    NormalizedCellContent::Blocks(blocks) => {
                        let mut measured = Vec::new();
                        let mut total_h = 0.0f32;
                        for block in blocks {
                            let m = super::measure_one_block(block, measurer, fonts, content_width)?;
                            total_h += super::measured_block_height(&m);
                            measured.push(m);
                        }
                        (Vec::new(), measured, total_h)
                    }
                };
                let cell_height = min_content_height + cell.padding * 2.0;

                let covered_height: f32 = (row_idx..(row_idx + row_span).min(measured_rows.len()))
                    .map(|i| measured_rows[i].height)
                    .sum();

                if cell_height > covered_height {
                    let extra = cell_height - covered_height;
                    let span_start = row_idx;
                    let span_end = (row_idx + row_span).min(measured_rows.len());
                    let span_count = span_end - span_start;
                    if span_count > 0 {
                        for i in span_start..span_end {
                            let proportion = if covered_height > 0.0 {
                                measured_rows[i].height / covered_height
                            } else {
                                1.0 / span_count as f32
                            };
                            measured_rows[i].height += extra * proportion;
                        }
                    }
                }

                measured_rows[row_idx].cells[cell_idx].lines = lines;
                measured_rows[row_idx].cells[cell_idx].items = items;
                measured_rows[row_idx].cells[cell_idx].height = cell_height;
            }
            col += col_span;
        }
    }

    if let Some(ref heights) = table.heights {
        for (i, row) in measured_rows.iter_mut().enumerate() {
            if let Some(&specified_height) = heights.get(i) {
                let new_row_height = specified_height.max(row.height);
                // Update cell heights to match the new row height
                for cell in &mut row.cells {
                    if cell.row_span <= 1 {
                        cell.height = new_row_height;
                    }
                }
                row.height = new_row_height;
            }
        }
    }

    let height = measured_rows.iter().map(|row| row.height).sum();
    Ok(MeasuredTable { column_widths, rows: measured_rows, style: table.style.clone(), height, header_rows: table.header_rows, dont_break_rows: table.dont_break_rows, page_break: table.page_break, heights: table.heights.clone(), keep_with_header_rows: table.keep_with_header_rows, table_style: table.table_style, layout_config: table.layout_config.clone(), h_line_when_broken: table.h_line_when_broken, origin_node_idx: table.origin_node_idx })
}

fn span_normalized_rows(rows: &[Vec<NormalizedTableCell>]) -> (Vec<Vec<NormalizedTableCell>>, usize) {
    let mut normalized = Vec::with_capacity(rows.len());
    let mut covered: HashSet<(usize, usize)> = HashSet::new();
    let mut column_count = 0usize;

    for (row_idx, row) in rows.iter().enumerate() {
        let mut normalized_row = Vec::new();
        let mut cell_idx = 0usize;
        let mut col = 0usize;

        while cell_idx < row.len() {
            while covered.contains(&(row_idx, col)) {
                if cell_idx < row.len() && is_span_placeholder(&row[cell_idx]) {
                    cell_idx += 1;
                }
                col += 1;
            }
            if cell_idx >= row.len() {
                break;
            }

            let cell = row[cell_idx].clone();
            let row_span = cell.row_span.max(1);
            let col_span = cell.col_span.max(1);
            for r in 1..row_span {
                if row_idx + r < rows.len() {
                    for c in 0..col_span {
                        covered.insert((row_idx + r, col + c));
                    }
                }
            }

            normalized_row.push(cell);
            col += col_span;
            cell_idx += 1;

            for _ in 1..col_span {
                if cell_idx < row.len() && is_span_placeholder(&row[cell_idx]) {
                    cell_idx += 1;
                }
            }
        }

        column_count = column_count.max(col);
        normalized.push(normalized_row);
    }

    (normalized, column_count)
}

fn is_span_placeholder(cell: &NormalizedTableCell) -> bool {
    cell.col_span.max(1) == 1
        && cell.row_span.max(1) == 1
        && matches!(&cell.content, NormalizedCellContent::Inlines(inlines) if inlines.iter().all(|inline| inline.text.is_empty()))
}

fn intrinsic_cell_width(cell: &NormalizedTableCell, measurer: &mut TextMeasurer, fonts: &FontRegistry) -> Result<f32> {
    let mut width = 0.0;
    match &cell.content {
        NormalizedCellContent::Inlines(inlines) => {
            for inline in inlines {
                width += measurer.measure_with_font(
                    &inline.text,
                    inline.style.font_family.as_deref(),
                    inline.style.font_size,
                    fonts,
                )?;
            }
        }
        NormalizedCellContent::Blocks(_) => {}
    }
    Ok(width + cell.padding * 2.0)
}

fn spanned_width(column_widths: &[f32], start_column: usize, col_span: usize) -> f32 {
    let end = (start_column + col_span.max(1)).min(column_widths.len());
    column_widths[start_column..end].iter().sum()
}

fn measure_table_row(row: &[NormalizedTableCell], column_widths: &[f32], measurer: &mut TextMeasurer, fonts: &FontRegistry, row_idx: usize, covered: &HashSet<(usize, usize)>) -> Result<MeasuredTableRow> {
    let mut column_index = 0usize;
    let mut cells = Vec::new();
    let mut cell_iter = row.iter();
    loop {
        while column_index < column_widths.len() && covered.contains(&(row_idx, column_index)) {
            column_index += 1;
        }
        let cell = match cell_iter.next() {
            Some(c) => c,
            None => break,
        };
        let col_span = cell.col_span.max(1);
        let width = spanned_width(column_widths, column_index, col_span);
        let content_width = (width - cell.padding * 2.0).max(0.0);
        let (lines, items, text_height) = match &cell.content {
            NormalizedCellContent::Inlines(inlines) => {
                let l = wrap_inlines_with_fonts(inlines, content_width, measurer, fonts)?;
                let h: f32 = l.iter().map(|line| line.line_height).sum();
                (l, Vec::new(), h)
            }
            NormalizedCellContent::Blocks(blocks) => {
                let mut measured = Vec::new();
                let mut total_h = 0.0f32;
                for block in blocks {
                    let m = super::measure_one_block(block, measurer, fonts, content_width)?;
                    total_h += super::measured_block_height(&m);
                    measured.push(m);
                }
                (Vec::new(), measured, total_h)
            }
        };
        let height = text_height + cell.padding * 2.0;
        cells.push(MeasuredTableCell { lines, items, width, height, padding: cell.padding, style: cell.style.clone(), col_span, row_span: cell.row_span.max(1), vertical_alignment: cell.vertical_alignment, border: cell.border.clone() });
        column_index += col_span;
    }
    let height = cells.iter().map(|cell| cell.height).fold(0.0f32, f32::max);
    Ok(MeasuredTableRow { cells, height })
}

fn resolve_column_widths(
    rows: &[Vec<NormalizedTableCell>],
    widths: &[TableWidth],
    column_count: usize,
    available_width: f32,
    measurer: &mut TextMeasurer,
    fonts: &FontRegistry,
) -> Result<Vec<f32>> {
    if column_count == 0 {
        return Ok(Vec::new());
    }

    let mut resolved = vec![0.0; column_count];
    let mut fixed_total = 0.0;
    let mut star_count = 0;

    for index in 0..column_count {
        match widths.get(index).copied().unwrap_or(TableWidth::Star) {
            TableWidth::Fixed(value) => {
                resolved[index] = value.max(0.0);
                fixed_total += resolved[index];
            }
            TableWidth::Auto => {
                let mut auto_width = 0.0;
                for row in rows {
                    if let Some(cell) = row.get(index) {
                        auto_width = f32::max(auto_width, intrinsic_cell_width(cell, measurer, fonts)?);
                    }
                }
                resolved[index] = auto_width;
                fixed_total += auto_width;
            }
            TableWidth::Star => star_count += 1,
        }
    }

    let remaining = (available_width - fixed_total).max(0.0);
    let star_width = if star_count == 0 { 0.0 } else { remaining / star_count as f32 };

    for index in 0..column_count {
        if matches!(widths.get(index).copied().unwrap_or(TableWidth::Star), TableWidth::Star) {
            resolved[index] = star_width;
        }
    }

    Ok(resolved)
}
