HawkeyeVision/Hawkeye.VisionBuilder/Features/ImageStatistics/ImageStatisticsDialog.cs

using System;
using System.Windows.Forms;
using System.Collections.Generic;
using System.Linq;

namespace Hawkeye.VisionBuilder.Features.ImageStatistics
{
    public partial class ImageStatisticsDialog : Form
    {
        private TreeNode _goodNode;
        private TreeNode _badNode;
        private HashSet<string> _allFiles = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
        private TreeNode? _selectedLeafNode = null;
        private Dictionary<string, int> _badDefectCounts = new(StringComparer.OrdinalIgnoreCase);

        public event Action<string>? FileSelected;

        public ImageStatisticsDialog()
        {
            InitializeComponent();
            InitializeTree();
            treeView.AfterSelect += treeView_AfterSelect;
            UpdateNavigationButtons();
        }

        private void InitializeTree()
        {
            treeView.Nodes.Clear();
            _goodNode = new TreeNode(GetBranchText("Good", 0));
            _badNode = new TreeNode(GetBranchText("Bad", 0));
            treeView.Nodes.Add(_goodNode);
            treeView.Nodes.Add(_badNode);
            treeView.ExpandAll();
            _allFiles.Clear();
            _badDefectCounts.Clear();
        }

        public void AddGoodImage(string filePath)
        {
            if (_allFiles.Contains(filePath))
                return;
            string nodeText = System.IO.Path.GetFileName(filePath);
            _goodNode.Nodes.Add(new TreeNode(nodeText) { Tag = filePath });
            _allFiles.Add(filePath);
            UpdateBranchText(_goodNode, "Good");
            //treeView.ExpandAll();
        }

        public void AddBadImage(string filePath, string[]? defects = null)
        {
            if (_allFiles.Contains(filePath))
                return;
            string nodeText = System.IO.Path.GetFileName(filePath);
            if (defects != null && defects.Length > 0)
            {
                nodeText += " (" + string.Join(", ", defects) + ")";
                foreach (var defect in defects)
                {
                    if (string.IsNullOrWhiteSpace(defect)) continue;
                    if (_badDefectCounts.ContainsKey(defect))
                        _badDefectCounts[defect]++;
                    else
                        _badDefectCounts[defect] = 1;
                }
            }
            _badNode.Nodes.Add(new TreeNode(nodeText) { Tag = filePath });
            _allFiles.Add(filePath);
            UpdateBranchText(_badNode, "Bad");
            //treeView.ExpandAll();
        }

        public void AddExpected(string filePath, string[]? defects)
        {
            // get folder name as expected defect
            string expectedDefect = System.IO.Path.GetFileName(System.IO.Path.GetDirectoryName(filePath) ?? string.Empty) ?? string.Empty;
            if (defects.Contains(expectedDefect, StringComparer.OrdinalIgnoreCase))
            {
                AddGoodImage(filePath);
            }
            else
            {
                AddBadImage(filePath, [expectedDefect]);
            }
        }

        private void UpdateBranchText(TreeNode node, string name)
        {
            if (node == _badNode)
            {
                node.Text = GetBranchText(name, node.Nodes.Count);
            }
            else
            {
                node.Text = GetBranchText(name, node.Nodes.Count);
            }
        }

        private string GetBranchText(string name, int count)
        {
            if (name == "Bad")
            {
                if (_badDefectCounts.Count > 0)
                {
                    var defectSummary = string.Join(", ",
                        _badDefectCounts
                            .OrderByDescending(kv => kv.Value)
                            .ThenBy(kv => kv.Key)
                            .Select(kv => $"{kv.Key} {kv.Value}")
                    );
                    return $"{name} ({defectSummary})";
                }
            }
            return $"{name} ({count})";
        }

        private void btnReset_Click(object sender, EventArgs e)
        {
            InitializeTree();
            _badDefectCounts.Clear();
        }

        private void btnClose_Click(object sender, EventArgs e)
        {
            this.Close();
        }

        private void treeView_NodeMouseClick(object sender, TreeNodeMouseClickEventArgs e)
        {
            // Only fire event for leaf nodes (filenames)
            if (e.Node.Parent != null && e.Node.Tag is string filePath)
            {
                _selectedLeafNode = e.Node;
                FileSelected?.Invoke(filePath);
            }
            else
            {
                _selectedLeafNode = null;
            }
            UpdateNavigationButtons();
        }

        private void treeView_AfterSelect(object? sender, TreeViewEventArgs e)
        {
            // Only consider leaf nodes (filenames)
            if (e.Node.Parent != null && e.Node.Tag is string)
            {
                _selectedLeafNode = e.Node;
            }
            else
            {
                _selectedLeafNode = null;
            }
            UpdateNavigationButtons();
        }

        private void UpdateNavigationButtons()
        {
            if (_selectedLeafNode == null)
            {
                btnBackward.Enabled = false;
                btnForward.Enabled = false;
                return;
            }
            var parent = _selectedLeafNode.Parent;
            if (parent == null)
            {
                btnBackward.Enabled = false;
                btnForward.Enabled = false;
                return;
            }
            int idx = parent.Nodes.IndexOf(_selectedLeafNode);
            btnBackward.Enabled = idx > 0;
            btnForward.Enabled = idx < parent.Nodes.Count - 1;
        }

        private void btnBackward_Click(object sender, EventArgs e)
        {
            if (_selectedLeafNode == null) return;
            var parent = _selectedLeafNode.Parent;
            if (parent == null) return;
            int idx = parent.Nodes.IndexOf(_selectedLeafNode);
            if (idx > 0)
            {
                var prev = parent.Nodes[idx - 1];
                treeView.SelectedNode = prev;
                treeView.Focus();
                if (prev.Tag is string filePath)
                    FileSelected?.Invoke(filePath);
            }
        }

        private void btnForward_Click(object sender, EventArgs e)
        {
            if (_selectedLeafNode == null) return;
            var parent = _selectedLeafNode.Parent;
            if (parent == null) return;
            int idx = parent.Nodes.IndexOf(_selectedLeafNode);
            if (idx < parent.Nodes.Count - 1)
            {
                var next = parent.Nodes[idx + 1];
                treeView.SelectedNode = next;
                treeView.Focus();
                if (next.Tag is string filePath)
                    FileSelected?.Invoke(filePath);
            }
        }
    }
}