import unittest
from src.investment.table_funcs import Table, get_table_info, recreate_page, remove_table_from_page, combine_continuous_tables, split_long_tables

class TestTableFunctions(unittest.TestCase):
    def setUp(self):
        self.sample_text_dict = {
            "1": "Some metadata\n-------Table Start--------\nTable 1\n{'A': ['1', '2'], 'B': ['3', '4']}\n-------Table End--------",
            "2": "More metadata\n-------Table Start--------\n\n{'A': ['5', '6'], 'B': ['7', '8']}\n-------Table End--------"
        }
        
        self.sample_table_dict = get_table_info(self.sample_text_dict)

    def test_get_table_info(self):
        self.assertIn("1", self.sample_table_dict)
        self.assertEqual(len(self.sample_table_dict["1"]), 1)
        self.assertEqual(self.sample_table_dict["1"][0].header, "Table 1")
        
    def test_recreate_page(self):
        page_tables = self.sample_table_dict["1"]
        recreated_text = recreate_page(page_tables)
        self.assertIn("Table 1", recreated_text)
        self.assertIn("-------Table Start--------", recreated_text)
        self.assertIn("-------Table End--------", recreated_text)
        
    def test_remove_table_from_page(self):
        table = self.sample_table_dict["1"][0]
        updated_text = remove_table_from_page(self.sample_text_dict["1"], table)
        self.assertNotIn("Table 1", updated_text)
        self.assertNotIn("-------Table Start--------", updated_text)
        
    def test_combine_continuous_tables(self):
        combined_text_dict = combine_continuous_tables(self.sample_text_dict, self.sample_table_dict)
        self.assertIn("1", combined_text_dict)
        self.assertNotIn("2", combined_text_dict)
        self.assertIn("5", combined_text_dict["1"])  # Ensuring second table is merged
        
    def test_split_long_tables(self):
        long_table_dict = {
            "1": "Some metadata\n-------Table Start--------\nTable 1\n{'A': ['1', '2', '3', '4', '5'], 'B': ['6', '7', '8', '9', '10']}\n-------Table End--------"
        }
        table_dict = get_table_info(long_table_dict)
        split_text_dict = split_long_tables(long_table_dict, table_dict, row_limit=2)
        self.assertGreater(len(split_text_dict), 1)
        for key, text in split_text_dict.items():
            self.assertIn("-------Table Start--------", text)
            self.assertIn("-------Table End--------", text)

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