import org.testng.annotations.Test;
import static org.testng.Assert.assertThrows;
import static org.testng.AssertJUnit.assertEquals;

import org.example.*;

import java.io.ByteArrayOutputStream;
import java.io.PrintStream;
import java.util.HashSet;
import java.util.Set;

public class MagicalSquareTest {

    @Test
    public void testGenerateMagicSquare_ValidInput15() {
        final int testValue = 15;
        int[][] magicSquare = MagicalSquare.generateMagicSquare(testValue);
        int expectedSum = testValue;

        // Test, zda magicSquare má správné rozměry (3x3)
        assertEquals("Magický čtverec nemá 3 řádky.", 3, magicSquare.length);
        assertEquals("Magický čtverec nemá 3 sloupce.", 3, magicSquare[0].length);

        // Test, zda součet hodnot v řádcích, sloupcích a diagonálách je očekávaný
        assertEquals("Nesedí součet prvního řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 0));
        assertEquals("Nesedí součet druhého řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 1));
        assertEquals("Nesedí součet třetího řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 2));
        assertEquals("Nesedí součet prvního sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 0));
        assertEquals("Nesedí součet druhého sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 1));
        assertEquals("Nesedí součet třetího sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 2));
        assertEquals("Nesedí součet první diagonály při magickém součtu " + testValue + ".", expectedSum, diagonalSum(magicSquare, true));
        assertEquals("Nesedí součet druhé diagonály při magickém součtu " + testValue + ".", expectedSum, diagonalSum(magicSquare, false));

        // Test, zda jsou čísla v čtverci různá
        Set<Integer> uniqueNumbers = new HashSet<>();
        for (int i = 0; i < magicSquare.length; i++) {
            for (int j = 0; j < magicSquare[i].length; j++) {
                uniqueNumbers.add(magicSquare[i][j]);
            }
        }
        assertEquals("Čísla v magickém čtverci nejsou různá.", 9, uniqueNumbers.size());
    }

    @Test
    public void testGenerateMagicSquare_ValidInput18() {
        final int testValue = 18;
        int[][] magicSquare = MagicalSquare.generateMagicSquare(testValue);
        int expectedSum = testValue;

        // Test, zda magicSquare má správné rozměry (3x3)
        assertEquals("Magický čtverec nemá 3 řádky.", 3, magicSquare.length);
        assertEquals("Magický čtverec nemá 3 sloupce.", 3, magicSquare[0].length);

        // Test, zda součet hodnot v řádcích, sloupcích a diagonálách je očekávaný
        assertEquals("Nesedí součet prvního řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 0));
        assertEquals("Nesedí součet druhého řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 1));
        assertEquals("Nesedí součet třetího řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 2));
        assertEquals("Nesedí součet prvního sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 0));
        assertEquals("Nesedí součet druhého sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 1));
        assertEquals("Nesedí součet třetího sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 2));
        assertEquals("Nesedí součet první diagonály při magickém součtu " + testValue + ".", expectedSum, diagonalSum(magicSquare, true));
        assertEquals("Nesedí součet druhé diagonály při magickém součtu " + testValue + ".", expectedSum, diagonalSum(magicSquare, false));

        // Test, zda jsou čísla v čtverci různá
        Set<Integer> uniqueNumbers = new HashSet<>();
        for (int i = 0; i < magicSquare.length; i++) {
            for (int j = 0; j < magicSquare[i].length; j++) {
                uniqueNumbers.add(magicSquare[i][j]);
            }
        }
        assertEquals("Čísla v magickém čtverci nejsou různá.", 9, uniqueNumbers.size());
    }

    @Test
    public void testGenerateMagicSquare_ValidInput72() {
        final int testValue = 72;
        int[][] magicSquare = MagicalSquare.generateMagicSquare(testValue);
        int expectedSum = testValue;

        // Test, zda magicSquare má správné rozměry (3x3)
        assertEquals("Magický čtverec nemá 3 řádky.", 3, magicSquare.length);
        assertEquals("Magický čtverec nemá 3 sloupce.", 3, magicSquare[0].length);

        // Test, zda součet hodnot v řádcích, sloupcích a diagonálách je očekávaný
        assertEquals("Nesedí součet prvního řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 0));
        assertEquals("Nesedí součet druhého řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 1));
        assertEquals("Nesedí součet třetího řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 2));
        assertEquals("Nesedí součet prvního sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 0));
        assertEquals("Nesedí součet druhého sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 1));
        assertEquals("Nesedí součet třetího sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 2));
        assertEquals("Nesedí součet první diagonály při magickém součtu " + testValue + ".", expectedSum, diagonalSum(magicSquare, true));
        assertEquals("Nesedí součet druhé diagonály při magickém součtu " + testValue + ".", expectedSum, diagonalSum(magicSquare, false));

        // Test, zda jsou čísla v čtverci různá
        Set<Integer> uniqueNumbers = new HashSet<>();
        for (int i = 0; i < magicSquare.length; i++) {
            for (int j = 0; j < magicSquare[i].length; j++) {
                uniqueNumbers.add(magicSquare[i][j]);
            }
        }
        assertEquals("Čísla v magickém čtverci nejsou různá.", 9, uniqueNumbers.size());
    }

    @Test
    public void testGenerateMagicSquare_ValidInput333() {
        final int testValue = 333;
        int[][] magicSquare = MagicalSquare.generateMagicSquare(testValue);
        int expectedSum = testValue;

        // Test, zda magicSquare má správné rozměry (3x3)
        assertEquals("Magický čtverec nemá 3 řádky.", 3, magicSquare.length);
        assertEquals("Magický čtverec nemá 3 sloupce.", 3, magicSquare[0].length);

        // Test, zda součet hodnot v řádcích, sloupcích a diagonálách je očekávaný
        assertEquals("Nesedí součet prvního řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 0));
        assertEquals("Nesedí součet druhého řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 1));
        assertEquals("Nesedí součet třetího řádku při magickém součtu " + testValue + ".", expectedSum, rowSum(magicSquare, 2));
        assertEquals("Nesedí součet prvního sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 0));
        assertEquals("Nesedí součet druhého sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 1));
        assertEquals("Nesedí součet třetího sloupce při magickém součtu " + testValue + ".", expectedSum, colSum(magicSquare, 2));
        assertEquals("Nesedí součet první diagonály při magickém součtu " + testValue + ".", expectedSum, diagonalSum(magicSquare, true));
        assertEquals("Nesedí součet druhé diagonály při magickém součtu " + testValue + ".", expectedSum, diagonalSum(magicSquare, false));
    }

    @Test
    public void testGenerateMagicSquare_InvalidInput10() {
        // Test, zda metoda vyhodí IllegalArgumentException pro neplatný vstup
        assertThrows("Při vstupu 10 má program vyhodit IllegalArgumentException - lze tvořit pouze magické součty větší nebo rovno než 15", IllegalArgumentException.class, () -> {
            int[][] magicSquare = MagicalSquare.generateMagicSquare(10);
        });
    }

    @Test
    public void testGenerateMagicSquare_InvalidInput16() {
        // Test, zda metoda vyhodí IllegalArgumentException pro neplatný vstup
        assertThrows("Při vstupu 16 má program vyhodit IllegalArgumentException - lze tvořit pouze magické součty dělitelné 3", IllegalArgumentException.class, () -> {
            int[][] magicSquare = MagicalSquare.generateMagicSquare(16);
        });
    }

    @Test
    public void testPrintMagicSquare() {
        int[][] magicSquare = MagicalSquare.generateMagicSquare(15);
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        System.setOut(new PrintStream(outputStream));

        MagicalSquare.printMagicSquare(magicSquare);

        String expectedOutput = "8 1 6 \r\n3 5 7 \r\n4 9 2 \r\n";
        String actualOutput = outputStream.toString();
        assertEquals("Nesprávná implementace funkce printMagicSquare().", expectedOutput, actualOutput);
    }

    // Pomocné metody pro výpočet součtu řádků, sloupců a diagonál
    private int rowSum(int[][] square, int row) {
        int sum = 0;
        for (int col = 0; col < square[row].length; col++) {
            sum += square[row][col];
        }
        return sum;
    }

    private int colSum(int[][] square, int col) {
        int sum = 0;
        for (int row = 0; row < square.length; row++) {
            sum += square[row][col];
        }
        return sum;
    }

    private int diagonalSum(int[][] square, boolean mainDiagonal) {
        int sum = 0;
        int size = square.length;
        for (int i = 0; i < size; i++) {
            if (mainDiagonal) {
                sum += square[i][i];
            } else {
                sum += square[i][size - 1 - i];
            }
        }
        return sum;
    }
}