import unittest
from dataclasses import dataclass, field

from ocr import drop_short_lines


@dataclass
class Line:
    baseline: list


@dataclass
class Bounds:
    lines: list = field(default_factory=list)


def row(x0, x1, y=100):
    return Line(baseline=[[x0, y], [x1, y]])


class DropShortLinesTests(unittest.TestCase):
    def test_drops_marginal_fragments(self):
        bounds = Bounds([row(100, 1100, y) for y in range(100, 1300, 100)] + [row(20, 60), row(1150, 1190)])
        kept, dropped = drop_short_lines(bounds, 0.2)
        self.assertEqual(dropped, 2)
        self.assertEqual(len(kept.lines), 12)

    def test_keeps_short_final_line_above_threshold(self):
        bounds = Bounds([row(100, 1100, y) for y in range(100, 900, 100)] + [row(100, 400, 1000)])
        kept, dropped = drop_short_lines(bounds, 0.2)
        self.assertEqual(dropped, 0)
        self.assertEqual(len(kept.lines), 9)

    def test_disabled_or_empty_is_unchanged(self):
        bounds = Bounds([row(0, 10), row(0, 1000)])
        self.assertEqual(drop_short_lines(bounds, 0)[1], 0)
        self.assertEqual(drop_short_lines(Bounds([]), 0.2)[1], 0)


if __name__ == "__main__":
    unittest.main()