import unittest
from tabulate import tabulate
from base_test import BaseAllocationTest


class AccrualStateTest(BaseAllocationTest):
    """Test accrual state calculations"""

    def test_accrual_state_calculation(self):
        """Test accrual state calculations after some blocks"""
        print("\nACCRUAL STATE CALCULATIONS")

        # Get latest block height for testing
        test_block_height = self._get_latest_block_height_from_api()
        if not test_block_height:
            self.skipTest("Could not retrieve latest block height")

        # Test first star and galaxy from each lock period (limit to avoid too many API calls)
        test_points = []
        for lock_period in sorted(self.points_by_duration.keys()):
            duration_data = self.points_by_duration[lock_period]

            if duration_data['star']:
                star_data = duration_data['star']
                zenith_addr = star_data['zenith_address']
                if zenith_addr:
                    test_points.append((star_data['point'], zenith_addr, lock_period))

            if duration_data['galaxy']:
                galaxy_data = duration_data['galaxy']
                zenith_addr = galaxy_data['zenith_address']
                if zenith_addr:
                    test_points.append((galaxy_data['point'], zenith_addr, lock_period))

        # Collect data for table
        accrual_data = []

        for point, zenith_addr, lock_period in test_points:
            with self.subTest(point=point, block_height=test_block_height):
                # Get allocation data first
                allocation_data = self._get_point_allocation_from_api(zenith_addr, point)
                self.assertIsNotNone(allocation_data, f"No allocation data for {point}")

                unlock_schedule = allocation_data.get('unlock_schedule')
                self.assertIsNotNone(unlock_schedule, f"No unlock_schedule for {point}")

                # Get accrual state at test block height
                accrual_state = self._get_point_accrual_state_from_api(zenith_addr, point, test_block_height)
                self.assertIsNotNone(accrual_state, f"No accrual state for {point} at block {test_block_height}")

                # Extract values for calculation
                total_allocation = int(allocation_data['allocated_amount']['amount'])
                initial_unlock_amount = int(unlock_schedule['initial_unlock_amount']['amount'])
                unlock_blocks = int(unlock_schedule['unlock_blocks'])

                api_total_unlocked = int(accrual_state['total_unlocked']['amount'])
                last_unlock_block = int(accrual_state['last_unlock_block'])

                # Calculate expected last unlock block using unlock_frequency
                expected_last_unlock_block = (test_block_height // self.unlock_frequency_blocks) * self.unlock_frequency_blocks

                # Assert on last unlock block
                self.assertEqual(expected_last_unlock_block, last_unlock_block,
                    f"Last unlock block mismatch for {point} at block {test_block_height}: "
                    f"Expected={expected_last_unlock_block}, "
                    f"zenithd={last_unlock_block}, "
                    f"Diff={last_unlock_block - expected_last_unlock_block}")

                # Calculate expected total_unlocked using expected last unlock block
                # Formula: initial_unlock_amount + (expected_last_unlock_block * remaining_amount / unlock_blocks)
                remaining_amount = total_allocation - initial_unlock_amount
                expected_total_unlocked = initial_unlock_amount + (expected_last_unlock_block * remaining_amount // unlock_blocks)

                difference = api_total_unlocked - expected_total_unlocked

                accrual_data.append([
                    point,
                    f"{lock_period} years",
                    f"Block {test_block_height}",
                    f"Block {last_unlock_block}",
                    f"{expected_total_unlocked:,}",
                    f"{api_total_unlocked:,}",
                    f"{difference:+,}" if difference != 0 else "0"
                ])

                self.assertEqual(expected_total_unlocked, api_total_unlocked,
                    f"Total unlocked mismatch for {point} at block {test_block_height}: "
                    f"Expected={expected_total_unlocked:,} $sZ, "
                    f"zenithd={api_total_unlocked:,} $sZ, "
                    f"Diff={difference:+,} $sZ")

        # Print table
        print(f"\nTotal Unlocked at Block {test_block_height}:")
        headers = ["Point", "Lock Period", "Block Height", "Last Unlocked At", "Expected ($sZ)", "zenithd ($sZ)", "Difference"]
        print(tabulate(accrual_data, headers=headers, tablefmt="grid"))


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