import unittest
from tabulate import tabulate
from base_test import BaseAllocationTest


class AllocationTest(BaseAllocationTest):
    """Test allocation comparisons between notebook and zenithd"""

    def test_0_star_allocations(self):
        """Test star allocations for all lock periods"""
        print("\nSTAR ALLOCATIONS COMPARISON")

        table_data = []
        headers = ["Lock Period", "Point", "Zenith Address", "Notebook ($sZ)", "zenithd ($sZ)", "Difference"]

        for lock_period in sorted(self.points_by_duration.keys()):
            with self.subTest(lock_period=lock_period):
                duration_data = self.points_by_duration[lock_period]

                star_data = duration_data['star']
                if star_data == None:
                    continue

                zenith_addr = star_data['zenith_address']
                self.assertIsNotNone(zenith_addr,
                    f"Could not find zenith address for star {star_data['point']}")

                allocation_data = self._get_point_allocation_from_api(zenith_addr, star_data['point'])
                self.assertIsNotNone(allocation_data,
                    f"zenithd returned no allocation for star {star_data['point']}")

                api_allocation = int(allocation_data['allocated_amount']['amount'])

                notebook_key = f'{lock_period}_years'
                notebook_allocation = self.notebook_allocations['stars'].get(notebook_key, 0)

                difference = api_allocation - notebook_allocation

                table_data.append([
                    f"{lock_period} years",
                    star_data['point'],
                    zenith_addr,
                    f"{notebook_allocation:,}",
                    f"{api_allocation:,}",
                    f"{difference:+,}" if difference != 0 else "0"
                ])

                self.assertEqual(notebook_allocation, api_allocation,
                    f"Star {star_data['point']} ({lock_period}Y): "
                    f"Notebook={notebook_allocation:,} $sZ, "
                    f"zenithd={api_allocation:,} $sZ, "
                    f"Diff={difference:+,} $sZ")

        print(tabulate(table_data, headers=headers, tablefmt="grid"))

    def test_1_galaxy_allocations(self):
        """Test galaxy allocations for all lock periods"""
        print("\nGALAXY ALLOCATIONS COMPARISON")

        table_data = []
        headers = ["Lock Period", "Point", "Zenith Address", "Notebook ($sZ)", "zenithd ($sZ)", "Difference"]

        for lock_period in sorted(self.points_by_duration.keys()):
            with self.subTest(lock_period=lock_period):
                duration_data = self.points_by_duration[lock_period]

                galaxy_data = duration_data['galaxy']
                if galaxy_data == None:
                    continue

                zenith_addr = galaxy_data['zenith_address']
                self.assertIsNotNone(zenith_addr,
                    f"Could not find zenith address for galaxy {galaxy_data['point']}")

                allocation_data = self._get_point_allocation_from_api(zenith_addr, galaxy_data['point'])
                self.assertIsNotNone(allocation_data,
                    f"zenithd returned no allocation for galaxy {galaxy_data['point']}")

                api_allocation = int(allocation_data['allocated_amount']['amount'])

                notebook_key = f'{lock_period}_years'
                notebook_allocation = self.notebook_allocations['galaxies'].get(notebook_key, 0)

                difference = api_allocation - notebook_allocation

                table_data.append([
                    f"{lock_period} years",
                    galaxy_data['point'],
                    zenith_addr,
                    f"{notebook_allocation:,}",
                    f"{api_allocation:,}",
                    f"{difference:+,}" if difference != 0 else "0"
                ])

                self.assertEqual(notebook_allocation, api_allocation,
                    f"Galaxy {galaxy_data['point']} ({lock_period}Y): "
                    f"Notebook={notebook_allocation:,} $sZ, "
                    f"zenithd={api_allocation:,} $sZ, "
                    f"Diff={difference:+,} $sZ")

        print(tabulate(table_data, headers=headers, tablefmt="grid"))

    def test_2_total_allocations(self):
        """Test total allocations for all participants"""
        print("\nTOTAL ALLOCATIONS COMPARISON")

        notebook_total = self.notebook_allocations.get('total', 0)

        if notebook_total == 0:
            self.skipTest("No total_allocation found in notebook data")

        api_total = 0

        for i, participant in enumerate(self.participants):
            zenith_address = participant['attestation']['payload']['payload']['address']

            with self.subTest(participant=i, address=zenith_address):
                participant_total = self._get_total_address_allocation_amount_from_api(zenith_address)
                api_total += participant_total

        # Summary table
        difference = api_total - notebook_total
        table_data = [
            ["Notebook", f"{notebook_total:,}"],
            ["zenithd", f"{api_total:,}"],
            ["Difference", f"{difference:+,}" if difference != 0 else "0"]
        ]
        headers = ["Source", "Total Allocation ($sZ)"]

        print(tabulate(table_data, headers=headers, tablefmt="grid"))

        self.assertEqual(notebook_total, api_total,
            f"Total allocation mismatch: "
            f"Notebook={notebook_total:,} $sZ, "
            f"zenithd={api_total:,} $sZ, "
            f"Diff={difference:+,} $sZ")


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