AX-Copilot-Codex/src/AxCopilot/Views/ChatWindow.TranscriptHost.cs

using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Media;

namespace AxCopilot.Views;

public partial class ChatWindow
{
    private readonly ObservableCollection<TranscriptVisualItem> _transcriptElements = [];
    private readonly Dictionary<UIElement, TranscriptVisualItem> _transcriptElementMap = new();
    private ScrollViewer? _transcriptScrollViewer;
    private bool _transcriptScrollHooked;

    private void InitializeTranscriptHost()
    {
        MessageList.ItemsSource = _transcriptElements;
        AttachTranscriptScrollChanged();
        MessageList.Loaded += (_, _) => AttachTranscriptScrollChanged();
    }

    private void AttachTranscriptScrollChanged()
    {
        var scrollViewer = FindVisualChild<ScrollViewer>(MessageList);
        if (scrollViewer == null || ReferenceEquals(scrollViewer, _transcriptScrollViewer))
            return;

        if (_transcriptScrollViewer != null && _transcriptScrollHooked)
        {
            _transcriptScrollViewer.ScrollChanged -= MessageScroll_ScrollChanged;
        }

        _transcriptScrollViewer = scrollViewer;
        _transcriptScrollViewer.ScrollChanged += MessageScroll_ScrollChanged;
        _transcriptScrollHooked = true;
    }

    private static T? FindVisualChild<T>(DependencyObject? parent)
        where T : DependencyObject
    {
        if (parent == null)
            return null;

        var childCount = VisualTreeHelper.GetChildrenCount(parent);
        for (var i = 0; i < childCount; i++)
        {
            var child = VisualTreeHelper.GetChild(parent, i);
            if (child is T match)
                return match;

            var descendant = FindVisualChild<T>(child);
            if (descendant != null)
                return descendant;
        }

        return null;
    }

    private static string CreateTranscriptElementKey() => $"ui_{Guid.NewGuid():N}";

    private int GetTranscriptElementCount() => _transcriptElements.Count;

    private UIElement? GetTranscriptElementAt(int index)
    {
        if (index < 0 || index >= _transcriptElements.Count)
            return null;

        return _transcriptElements[index].GetOrCreateElement();
    }

    private int GetTranscriptElementIndex(UIElement element)
    {
        if (_transcriptElementMap.TryGetValue(element, out var item))
            return _transcriptElements.IndexOf(item);

        for (var i = 0; i < _transcriptElements.Count; i++)
        {
            var candidate = _transcriptElements[i];
            if (!candidate.IsMaterialized)
                continue;

            if (ReferenceEquals(candidate.Element, element))
            {
                _transcriptElementMap[element] = candidate;
                return i;
            }
        }

        return -1;
    }

    private bool ContainsTranscriptElement(UIElement element) => GetTranscriptElementIndex(element) >= 0;

    private TranscriptVisualItem WrapTranscriptElement(string key, UIElement element)
    {
        var item = new TranscriptVisualItem(key, element);
        _transcriptElementMap[element] = item;
        return item;
    }

    private void AddTranscriptElement(UIElement element)
        => _transcriptElements.Add(WrapTranscriptElement(CreateTranscriptElementKey(), element));

    private void AddDeferredTranscriptElement(string key, Func<UIElement> elementFactory)
    {
        TranscriptVisualItem? item = null;
        item = new TranscriptVisualItem(
            key,
            elementFactory,
            element => _transcriptElementMap[element] = item!);
        _transcriptElements.Add(item);
    }

    private void ClearTranscriptElements()
    {
        _transcriptElements.Clear();
        _transcriptElementMap.Clear();
        _lastGroupedProcessFeedKey = null;
        _lastGroupedProcessFeedIndex = -1;
        _processFeedAppendCount = 0;
        _processFeedMergeCount = 0;
        _transcriptRowKindCounts.Clear();
    }

    private void RemoveTranscriptElement(UIElement element)
    {
        var index = GetTranscriptElementIndex(element);
        if (index >= 0)
            RemoveTranscriptElementAt(index);
    }

    private void RemoveTranscriptElementAt(int index)
    {
        if (index < 0 || index >= _transcriptElements.Count)
            return;

        var item = _transcriptElements[index];
        if (item.Element != null)
            _transcriptElementMap.Remove(item.Element);
        _transcriptElements.RemoveAt(index);
    }

    private void ReplaceTranscriptElement(int index, UIElement element)
    {
        if (index < 0 || index >= _transcriptElements.Count)
            return;

        var previous = _transcriptElements[index];
        if (previous.Element != null)
            _transcriptElementMap.Remove(previous.Element);

        _transcriptElements[index] = WrapTranscriptElement(CreateTranscriptElementKey(), element);
    }

    private double GetTranscriptScrollableHeight() => _transcriptScrollViewer?.ScrollableHeight ?? 0;

    private double GetTranscriptVerticalOffset() => _transcriptScrollViewer?.VerticalOffset ?? 0;

    private void ScrollTranscriptToEnd() => _transcriptScrollViewer?.ScrollToEnd();

    private void ScrollTranscriptToVerticalOffset(double offset) => _transcriptScrollViewer?.ScrollToVerticalOffset(offset);
}