import unittest
from datetime import datetime
from tabulate import tabulate
from base_test import BaseAllocationTest, SECONDS_PER_YEAR, BLOCK_DURATION_SECONDS


class UnlockScheduleTest(BaseAllocationTest):
    """Test unlock schedule calculations"""

    def test_unlock_schedule_calculation(self):
        """Test unlock schedule calculations for all lock periods"""
        print("\nUNLOCK SCHEDULE CALCULATIONS")

        genesis_timestamp = self._get_genesis_time_from_api()
        if not genesis_timestamp:
            self.skipTest("Could not retrieve genesis time")

        # 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, star_data['block_timestamp']))

            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, galaxy_data['block_timestamp']))

        # Collect data for tables
        unlock_blocks_data = []
        initial_unlock_data = []

        for point, zenith_addr, lock_period, start_timestamp in test_points:
            with self.subTest(point=point):
                # Get allocation with unlock schedule from zenithd
                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}")

                # Use start timestamp from watcher events
                # Calculate time difference and pregenesis blocks
                time_diff_seconds = genesis_timestamp - start_timestamp
                pregenesis_blocks = time_diff_seconds // BLOCK_DURATION_SECONDS

                # Calculate total lock duration in blocks (years * seconds_per_year / block_time)
                total_lock_blocks = lock_period * SECONDS_PER_YEAR // BLOCK_DURATION_SECONDS

                # Calculate expected remaining blocks
                expected_unlock_blocks = total_lock_blocks - pregenesis_blocks
                api_unlock_blocks = int(unlock_schedule['unlock_blocks'])
                blocks_diff = api_unlock_blocks - expected_unlock_blocks

                # Convert start timestamp to ISO format
                start_time_iso = datetime.fromtimestamp(start_timestamp).isoformat()

                unlock_blocks_data.append([
                    point,
                    f"{lock_period} years",
                    start_time_iso,
                    f"{expected_unlock_blocks:,}",
                    f"{api_unlock_blocks:,}",
                    f"{blocks_diff:+,}" if blocks_diff != 0 else "0"
                ])

                self.assertEqual(expected_unlock_blocks, api_unlock_blocks,
                    f"Unlock blocks mismatch for {point}: "
                    f"Expected={expected_unlock_blocks}, "
                    f"zenithd={api_unlock_blocks}, "
                    f"Diff={blocks_diff}")

                # Calculate expected initial unlock amount
                total_allocation = int(allocation_data['allocated_amount']['amount'])
                expected_initial_unlock = (pregenesis_blocks * total_allocation) // total_lock_blocks
                api_initial_unlock = int(unlock_schedule['initial_unlock_amount']['amount'])
                unlock_diff = api_initial_unlock - expected_initial_unlock

                initial_unlock_data.append([
                    point,
                    f"{lock_period} years",
                    start_time_iso,
                    f"{expected_initial_unlock:,}",
                    f"{api_initial_unlock:,}",
                    f"{unlock_diff:+,}" if unlock_diff != 0 else "0"
                ])

                self.assertEqual(expected_initial_unlock, api_initial_unlock,
                    f"Initial unlock amount mismatch for {point}: "
                    f"Expected={expected_initial_unlock:,} $sZ, "
                    f"zenithd={api_initial_unlock:,} $sZ, "
                    f"Diff={unlock_diff:+,} $sZ")

        # Print tables
        print("\nUnlock Blocks Comparison:")
        unlock_blocks_headers = ["Point", "Lock Period", "Start Time", "Expected Blocks", "zenithd Blocks", "Difference"]
        print(tabulate(unlock_blocks_data, headers=unlock_blocks_headers, tablefmt="grid"))

        print("\nInitial Unlock Amounts Comparison:")
        initial_unlock_headers = ["Point", "Lock Period", "Start Time", "Expected ($sZ)", "zenithd ($sZ)", "Difference"]
        print(tabulate(initial_unlock_data, headers=initial_unlock_headers, tablefmt="grid"))


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