/*
 * Copyright (C) 2025  Tim Zimmermann <tim@aglais.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package org.aglais.zms.ui.component;

import org.aglais.zms.db.core.model.AncestryData;

import javax.swing.*;
import java.awt.*;
import java.awt.event.MouseListener;
import java.util.*;
import java.util.List;

/**
 * A swing component that will draw an ancestry diagram.
 */
public class AncestryDiagram extends JPanel {

    private List<List<AncestryData>> layers;
    private final Map<AncestryData, AncestryDiagramPoint> points = new HashMap<>();
    private int columns;
    private int rows;
    private MouseListener pointMouseListener;

    /**
     * Create a new {@link AncestryDiagram} object.
     */
    public AncestryDiagram() {
        setLayout(new GridBagLayout());
    }

    /**
     * Get the current layers
     *
     * @return Currently painted layers
     */
    public List<List<AncestryData>> getLayers() {
        return layers;
    }

    /**
     * Set the data to paint
     *
     * @param data Data to paint
     */
    public void setData(AncestryData... data) {
        this.layers = createLayers(List.of(data));

        removeAll();
        addComponents();
        revalidate();
        repaint();
    }

    /**
     * Sets the mouse listener for points in the diagram.
     *
     * @param pointMouseListener The MouseListener to be set for interacting with points in the diagram
     */
    public void setPointMouseListener(MouseListener pointMouseListener) {
        points.forEach((data, point) -> {
            point.removeMouseListener(this.pointMouseListener);
            point.addMouseListener(pointMouseListener);
        });

        this.pointMouseListener = pointMouseListener;
    }

    /**
     * Retrieves a list of points currently represented in the ancestry diagram.
     *
     * @return A list of {@link AncestryDiagramPoint} objects, representing the points in the diagram.
     */
    public List<AncestryDiagramPoint> getPoints() {
        return points.values().stream().toList();
    }

    /**
     * Calculates the starting column index for centering a layer within total columns.
     *
     * @param totalColumns   Total number of columns available
     * @param columnsInLayer Number of columns in the current layer
     * @return The starting column index to center the layer
     */
    private int getColumnStart(int totalColumns, int columnsInLayer) {
        return (totalColumns - columnsInLayer) / 2;
    }

    /**
     * Sorts each layer of ancestry data to improve visual alignment based on parent positions.
     *
     * @param layers  The list of layers to sort
     * @param columns The total number of columns across layers
     */
    private void sortLayers(List<List<AncestryData>> layers, int columns) {
        for (int i = 0; i < layers.size(); i++) {
            if (i > 0) {
                List<AncestryData> currentLayer = layers.get(i);
                AncestryData[] replacementLayer = new AncestryData[currentLayer.size()];
                int columnStart = getColumnStart(columns, currentLayer.size());

                for (int j = 0; j < currentLayer.size(); j++) {
                    AncestryData data = currentLayer.get(j);
                    int totalDistance = 0;
                    int totalCount = 0;

                    // Search layers containing parents and calculate offset to index
                    for (int parent : data.getParents()) {
                        for (int k = 0; k < i; k++) {
                            List<AncestryData> previousLayer = layers.get(k);
                            AncestryData parentData = null;
                            for (AncestryData previousData : previousLayer) {
                                if (previousData.getID() == parent) {
                                    parentData = previousData;
                                    break;
                                }
                            }

                            if (parentData != null) {
                                int previousColumnStart = getColumnStart(columns, previousLayer.size());

                                totalDistance = (previousColumnStart + previousLayer.indexOf(parentData))
                                        - (columnStart + j);
                                totalCount++;
                                break;
                            }
                        }
                    }
                    float avgDistance = (float) totalDistance / totalCount;
                    int newPos = (int) Math.floor(j + avgDistance);
                    for (int offset = 0; newPos + offset < replacementLayer.length || newPos - offset >= 0; offset++) {
                        if (newPos + offset >= 0 && newPos + offset < replacementLayer.length) {
                            if (replacementLayer[newPos + offset] == null) {
                                replacementLayer[newPos + offset] = data;
                                break;
                            }
                        }
                        if (newPos - offset >= 0 && newPos - offset < replacementLayer.length) {
                            if (replacementLayer[newPos - offset] == null) {
                                replacementLayer[newPos - offset] = data;
                                break;
                            }
                        }
                    }
                }

                layers.set(i, List.of(replacementLayer));
            }
        }
    }

    /**
     * Retrieves all children of the given parent ancestry data.
     * <p>
     *
     * @param parent The parent ancestry data for which to find all children
     * @return A list of all children ancestry data
     */
    private List<AncestryData> getAllChildren(AncestryData parent) {
        List<AncestryData> children = new ArrayList<>();
        List<Integer> allParentIds = new ArrayList<>(List.of(parent.getID()));

        for (List<AncestryData> layer : layers) {
            for (AncestryData data : layer) {
                for (int parentId : data.getParents()) {
                    if (allParentIds.contains(parentId)) {
                        allParentIds.add(data.getID());
                        children.add(data);
                        break;
                    }
                }
            }
        }

        return children;
    }

    /**
     * Determines if the given ancestry data is extinct.
     *
     * @param data The ancestry data to check for extinction
     * @return True if the ancestry data is extinct, false otherwise
     */
    private boolean isExtinct(AncestryData data) {
        List<AncestryData> allData = getAllChildren(data);
        allData.add(data);

        boolean extinct = true;

        for (AncestryData ancestry : allData) {
            if (ancestry.getDateEnd() == null || ancestry.getDateEnd().getTime() > System.currentTimeMillis()) {
                extinct = false;
                break;
            }
        }

        return extinct;
    }

    /**
     * Processes the given hierarchical ancestry data into a layered structure that shows the
     * relationships across multiple levels, starting from the current generation and traversing
     * to the ancestors. The resulting layers are organized from top (current generation) to bottom
     * (oldest layer of ancestors).
     *
     * @param flatOriginalList A list of {@link AncestryData} objects representing the primary nodes
     *                         for which the ancestry hierarchy is to be processed.
     * @return A list of layers, where each layer is a list of {@link AncestryData} objects
     * representing a single generation, ordered from the current generation down
     * to the oldest.
     */
    private List<List<AncestryData>> createLayers(List<AncestryData> flatOriginalList) {
        columns = 0;

        List<AncestryData> flatList = new ArrayList<>(flatOriginalList);
        Map<Integer, AncestryData> idMap = new HashMap<>();
        Map<Integer, List<Integer>> childMap = new HashMap<>();
        List<List<AncestryData>> layers = new ArrayList<>();
        List<AncestryData> currentLayer = new ArrayList<>();

        for (AncestryData data : flatList) {
            idMap.put(data.getID(), data);
            if (!childMap.containsKey(data.getID())) {
                childMap.put(data.getID(), new ArrayList<>());
            }

            if (data.getParents().length == 0) {
                currentLayer.add(data);
            } else {
                for (int parent : data.getParents()) {
                    List<Integer> childList = childMap.getOrDefault(parent, new ArrayList<>());

                    childList.add(data.getID());
                    childMap.put(parent, childList);
                }
            }
        }

        Set<Integer> addedNodes = new HashSet<>();
        while (!currentLayer.isEmpty()) {
            flatList.removeAll(currentLayer);
            layers.add(currentLayer);
            columns = Math.max(columns, currentLayer.size());

            List<AncestryData> nextLayer = new ArrayList<>();
            for (AncestryData data : currentLayer) {
                addedNodes.add(data.getID());

                for (int child : childMap.get(data.getID())) {
                    if (!addedNodes.contains(child)) {
                        addedNodes.add(child);
                        nextLayer.add(idMap.get(child));
                    }
                }
            }

            currentLayer = nextLayer;
        }

        // Each row needs one extra row for connections except for the last one
        rows = layers.size() * 2 - 1;

        sortLayers(layers, columns);

        return layers;
    }

    /**
     * Updates and re-renders the ancestry diagram by populating the components and layout
     * as per the provided ancestry data and the internal structure of the diagram.
     * <p>
     * The method clears all existing points in the diagram and iterates through the defined
     * ancestry layers. For each row in the layers, it calculates the appropriate column start
     * position to center the content horizontally and creates {@link AncestryDiagramPoint}
     * objects corresponding to each `AncestryData` element. These points are positioned
     * dynamically using a {@link GridBagLayout} configuration.
     * <p>
     * Additionally, the method incorporates invisible spacer panels between rows in odd-indexed
     * positions to maintain adequate spacing in the UI.
     */
    private void addComponents() {
        points.clear();

        for (int row = 0; row < layers.size(); row++) {
            List<AncestryData> layer = layers.get(row);
            int columnStart = getColumnStart(columns, layer.size());

            for (int column = 0; column < layer.size(); column++) {
                int columnDisplay = columnStart + column;
                GridBagConstraints constraints = new GridBagConstraints();
                constraints.gridx = columnDisplay;
                constraints.gridy = row * 2;
                constraints.fill = GridBagConstraints.BOTH;
                constraints.weightx = 1.0;
                constraints.weighty = 1.0;

                AncestryDiagramPoint point = new AncestryDiagramPoint(layer.get(column));
                point.setExtinct(isExtinct(layer.get(column)));

                add(point, constraints);

                point.addMouseListener(pointMouseListener);

                points.put(layer.get(column), point);
            }
        }

        for (int row = 0; row < rows; row++) {
            // Add spacers to ensure buttons don't get too close
            if (row % 2 != 0) {
                GridBagConstraints constraints = new GridBagConstraints();
                constraints.gridx = 0;
                constraints.gridy = row;
                constraints.fill = GridBagConstraints.BOTH;
                constraints.weightx = 1.0;
                constraints.weighty = 1.0;

                JPanel spacer = new JPanel();
                spacer.setOpaque(false);

                add(spacer, constraints);
            }
        }
    }

    /**
     * Calculates and returns the coordinates of the center point
     * along the top edge of the given {@link JComponent}.
     *
     * @param comp the {@link JComponent} for which the top edge center is calculated
     * @return a {@link Point} representing the center of the top edge of the given component
     */
    private Point getTopEdgeCenter(JComponent comp) {
        Point location = comp.getLocation();
        int x = location.x + comp.getWidth() / 2;

        return new Point(x, location.y);
    }

    /**
     * Calculates the center point of the bottom edge of a given component.
     *
     * @param comp the component for which the bottom edge center point is to be determined
     * @return a Point representing the center of the bottom edge of the specified component
     */
    private Point getBottomEdgeCenter(JComponent comp) {
        Point topEdgeCenter = getTopEdgeCenter(comp);
        int y = topEdgeCenter.y + comp.getHeight();

        return new Point(topEdgeCenter.x, y);
    }

    @Override
    protected void paintComponent(Graphics g) {
        super.paintComponent(g);

        Graphics2D g2 = (Graphics2D) g;
        g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
        g2.setColor(UIManager.getColor("Label.foreground"));
        g2.setStroke(new BasicStroke(2));

        for (Component component : getComponents()) {
            if (component instanceof AncestryDiagramPoint point) {
                Point p1 = getTopEdgeCenter(point);

                for (int parent : point.getData().getParents()) {
                    for (AncestryData data : points.keySet()) {
                        if (data.getID() == parent) {
                            AncestryDiagramPoint parentPoint = points.get(data);
                            Point p2 = getBottomEdgeCenter(parentPoint);

                            g2.drawLine(p1.x, p1.y, p2.x, p2.y);
                            break;
                        }
                    }
                }
            }
        }
    }
}