﻿Public Class TreeList

#Region "Events"
    Public Event SelectedNodeChanged(ByVal sender As TreeList, ByVal selection As TreeListNode)
    Public Event NodeExpanded(ByVal sender As TreeList, ByVal node As TreeListNode)
    Public Event NodeCollapsed(ByVal sender As TreeList, ByVal node As TreeListNode)
#End Region

#Region "Variables"
    Private m_Nodes As New TreeListNodeCollection(Me)
    Private m_SelectedNode As TreeListNode
    Private m_NodeIndent As Integer = 20
    Private m_Updating As Boolean = False
#End Region

#Region "Properties"
    Public ReadOnly Property Nodes() As TreeListNodeCollection
        Get
            Return m_Nodes
        End Get
    End Property

    Public Property SelectedNode() As TreeListNode
        Get
            Return m_SelectedNode
        End Get
        Set(ByVal value As TreeListNode)
            m_SelectedNode = value
            Me.Invalidate()
        End Set
    End Property

    Public Property NodeIndent() As Integer
        Get
            Return m_NodeIndent
        End Get
        Set(ByVal value As Integer)
            m_NodeIndent = value
            Me.Invalidate()
        End Set
    End Property
#End Region

#Region "Drawing"
    Protected Overrides Sub OnPaint(ByVal e As System.Windows.Forms.PaintEventArgs)
        If Not m_Updating Then
            Dim g As Graphics = e.Graphics
            Dim y As Integer = 2
            For Each n As TreeListNode In Me.Nodes
                DrawRecursive(g, n, y, 0)
            Next
            If Me.Nodes.Count > 0 Then g.DrawLine(Pens.Black, 0, y, Me.Width, y)
            g.DrawRectangle(Pens.Black, New Rectangle(0, 0, Me.Width - 1, Me.Height - 1))
            MyBase.OnPaint(e)
        End If
    End Sub

    Private Sub DrawRecursive(ByVal g As Graphics, ByVal n As TreeListNode, ByRef y As Integer, ByVal indent As Integer)
        n.RenderTo(g, Me.SelectedNode Is n, indent, y)
        If n.Expanded Then
            For Each node As TreeListNode In n.ChildNodes
                DrawRecursive(g, node, y, indent + m_NodeIndent)
            Next
            If n.ChildList IsNot Nothing Then
                n.ChildList.Location = New Point(0, y)
                n.ChildList.Width = Me.Width
                Me.Controls.Add(n.ChildList)
                y += n.ChildList.Height
            End If
        Else
            If n.ChildList IsNot Nothing Then _
                Me.Controls.Remove(n.ChildList) 'Will just return False if not found.
        End If
    End Sub
#End Region

#Region "Utility Methods"
    Private Function GetNodeAt(ByVal y As Integer) As TreeListNode
        Dim g As Graphics = Me.CreateGraphics()
        Dim r As TreeListNode = Nothing
        Dim yR As Integer = 0
        For Each n As TreeListNode In Me.Nodes
            RetrieveRecursive(g, n, yR, y, r)
            If r IsNot Nothing Then Return r
        Next
        Return r
    End Function

    Private Sub RetrieveRecursive(ByVal g As Graphics, ByVal n As TreeListNode, ByRef y As Integer, ByVal destY As Integer, ByRef found As TreeListNode)
        Dim h As Integer = n.GetHeight(g)
        If y + h > destY AndAlso y < destY Then
            'Found it!
            found = n
            Exit Sub
        End If
        y += h
        If n.Expanded Then
            For Each node As TreeListNode In n.ChildNodes
                RetrieveRecursive(g, node, y, destY, found)
                If found IsNot Nothing Then Exit Sub
            Next
            If n.ChildList IsNot Nothing Then y += n.ChildList.Height
        End If
    End Sub
#End Region

#Region "Public Methods"
    Public Sub BeginUpdate()
        m_Updating = True
    End Sub

    Public Sub EndUpdate()
        m_Updating = False
        Me.Invalidate()
    End Sub
#End Region

#Region "Mouse Events"
    Protected Overrides Sub OnMouseClick(ByVal e As System.Windows.Forms.MouseEventArgs)
        Dim n As TreeListNode = GetNodeAt(e.Y) 'TODO: Add scrolling.
        If n IsNot Nothing Then
            n.Clicked()
            Me.SelectedNode = n
        End If
        Me.Focus()
        MyBase.OnMouseClick(e)
    End Sub
#End Region

#Region "Constructor"
    Public Sub New()
        MyBase.New()
        Me.InitializeComponent()
        Me.DoubleBuffered = True
        Me.SetStyle(ControlStyles.AllPaintingInWmPaint Or ControlStyles.UserPaint Or _
                    ControlStyles.OptimizedDoubleBuffer Or ControlStyles.Selectable Or _
                    ControlStyles.CacheText Or ControlStyles.ResizeRedraw Or _
                    ControlStyles.StandardClick Or ControlStyles.StandardDoubleClick Or _
                    ControlStyles.SupportsTransparentBackColor, True)
    End Sub
#End Region

#Region "Custom Collection"
    Public Class TreeListNodeCollection
        Implements IList(Of TreeListNode)

        Private m_InnerList As New List(Of TreeListNode)
        Private m_ParentTreeList As TreeList

        Public Sub New(ByVal parent As TreeList)
            m_ParentTreeList = parent
        End Sub

        Private Sub Update()
            If m_ParentTreeList IsNot Nothing Then m_ParentTreeList.Invalidate()
        End Sub

        Public Sub Add(ByVal item As TreeListNode) Implements System.Collections.Generic.ICollection(Of TreeListNode).Add
            m_InnerList.Add(item)
            Me.Update()
        End Sub

        Public Sub Clear() Implements System.Collections.Generic.ICollection(Of TreeListNode).Clear
            m_InnerList.Clear()
            Me.Update()
        End Sub

        Public Function Contains(ByVal item As TreeListNode) As Boolean Implements System.Collections.Generic.ICollection(Of TreeListNode).Contains
            Return m_InnerList.Contains(item)
        End Function

        Public Sub CopyTo(ByVal array() As TreeListNode, ByVal arrayIndex As Integer) Implements System.Collections.Generic.ICollection(Of TreeListNode).CopyTo
            m_InnerList.CopyTo(array, arrayIndex)
        End Sub

        Public ReadOnly Property Count() As Integer Implements System.Collections.Generic.ICollection(Of TreeListNode).Count
            Get
                Return m_InnerList.Count
            End Get
        End Property

        Public ReadOnly Property IsReadOnly() As Boolean Implements System.Collections.Generic.ICollection(Of TreeListNode).IsReadOnly
            Get
                Return False
            End Get
        End Property

        Public Function Remove(ByVal item As TreeListNode) As Boolean Implements System.Collections.Generic.ICollection(Of TreeListNode).Remove
            Remove = m_InnerList.Remove(item)
            If (item.ChildList IsNot Nothing) AndAlso item.Expanded Then
                m_ParentTreeList.Controls.Remove(item.ChildList)
            End If
            Me.Update()
        End Function

        Public Function GetEnumerator() As System.Collections.Generic.IEnumerator(Of TreeListNode) Implements System.Collections.Generic.IEnumerable(Of TreeListNode).GetEnumerator
            Return m_InnerList.GetEnumerator()
        End Function

        Public Function IndexOf(ByVal item As TreeListNode) As Integer Implements System.Collections.Generic.IList(Of TreeListNode).IndexOf
            Return m_InnerList.IndexOf(item)
        End Function

        Public Sub Insert(ByVal index As Integer, ByVal item As TreeListNode) Implements System.Collections.Generic.IList(Of TreeListNode).Insert
            m_InnerList.Insert(index, item)
            Me.Update()
        End Sub

        Default Public Property Item(ByVal index As Integer) As TreeListNode Implements System.Collections.Generic.IList(Of TreeListNode).Item
            Get
                Return m_InnerList(index)
            End Get
            Set(ByVal value As TreeListNode)
                m_InnerList(index) = value
                Me.Update()
            End Set
        End Property

        Public Sub RemoveAt(ByVal index As Integer) Implements System.Collections.Generic.IList(Of TreeListNode).RemoveAt
            m_InnerList.RemoveAt(index)
            If (Me(index).ChildList IsNot Nothing) AndAlso Me(index).Expanded Then
                m_ParentTreeList.Controls.Remove(Me(index).ChildList)
            End If
            Me.Update()
        End Sub

        Private Function GetEnumerator1() As System.Collections.IEnumerator Implements System.Collections.IEnumerable.GetEnumerator
            Return Me.GetEnumerator()
        End Function
    End Class
#End Region

End Class

#Region "TreeListNode"
Public Class TreeListNode
    Implements ICloneable

#Region "Variables"
    Private m_ParentList As TreeList
    'Private m_ParentNode As TreeListNode
    Private m_Children As TreeList.TreeListNodeCollection
    Private m_ChildList As ListView

    Private m_Text As String = "TreeListNode"
    Private m_BackColor As Color = SystemColors.Control
    Private m_SelectedBackColor As Color = SystemColors.Highlight
    Private m_ForeColor As Color = SystemColors.WindowText
    Private m_SelectedForeColor As Color = SystemColors.HighlightText
    Private m_Icon As Bitmap = Nothing
    Private m_Height As Integer = 0
    Private m_Font As Font = Nothing

    Private m_Expanded As Boolean = False
#End Region

#Region "Events"
    Public Event SelectedChanged(ByVal sender As TreeListNode, ByVal selected As Boolean)
    Public Event ExpandedChanged(ByVal sender As TreeListNode, ByVal expanded As Boolean)
#End Region

#Region "Friend Event Helpers"
    Friend Sub RaiseSelectedChanged(ByVal selected As Boolean)
        RaiseEvent SelectedChanged(Me, selected)
    End Sub

    Friend Sub Clicked()
        Me.Expanded = Not Me.Expanded
    End Sub
#End Region

#Region "Constructor"
    Public Sub New(ByVal parentlist As TreeList)
        m_ParentList = parentlist
        m_Children = New TreeList.TreeListNodeCollection(m_ParentList)
    End Sub
#End Region

#Region "Node Properties"
    Public ReadOnly Property ChildNodes() As TreeList.TreeListNodeCollection
        Get
            Return m_Children
        End Get
    End Property

    Public ReadOnly Property FirstChild() As TreeListNode
        Get
            If Me.ChildNodes.Count > 0 Then
                Return Me.ChildNodes(0)
            Else
                Return Nothing
            End If
        End Get
    End Property

    Public ReadOnly Property LastChild() As TreeListNode
        Get
            If Me.ChildNodes.Count > 0 Then
                Return Me.ChildNodes(Me.ChildNodes.Count - 1)
            Else
                Return Nothing
            End If
        End Get
    End Property

    Public Property ChildList() As ListView
        Get
            Return m_ChildList
        End Get
        Set(ByVal value As ListView)
            m_ChildList = value
            m_ParentList.Invalidate()
        End Set
    End Property
#End Region

#Region "Design Properties"
    Public Property Text() As String
        Get
            Return m_Text
        End Get
        Set(ByVal value As String)
            m_Text = value
            m_ParentList.Invalidate()
        End Set
    End Property

    Public Property BackColor() As Color
        Get
            Return m_BackColor
        End Get
        Set(ByVal value As Color)
            m_BackColor = value
            m_ParentList.Invalidate()
        End Set
    End Property

    Public Property ForeColor() As Color
        Get
            Return m_ForeColor
        End Get
        Set(ByVal value As Color)
            m_ForeColor = value
            m_ParentList.Invalidate()
        End Set
    End Property

    Public Property SelectedBackColor() As Color
        Get
            Return m_SelectedBackColor
        End Get
        Set(ByVal value As Color)
            m_SelectedBackColor = value
            m_ParentList.Invalidate()
        End Set
    End Property

    Public Property SelectedForeColor() As Color
        Get
            Return m_SelectedForeColor
        End Get
        Set(ByVal value As Color)
            m_SelectedForeColor = value
            m_ParentList.Invalidate()
        End Set
    End Property

    Public Property Icon() As Bitmap
        Get
            Return m_Icon
        End Get
        Set(ByVal value As Bitmap)
            m_Icon = value
            m_ParentList.Invalidate()
        End Set
    End Property

    ''' <summary>
    ''' The height of the TreeListNode. Set to zero for auto-sizing.
    ''' </summary>
    Public Property Height() As Integer
        Get
            Return m_Height
        End Get
        Set(ByVal value As Integer)
            m_Height = value
            m_ParentList.Invalidate()
        End Set
    End Property

    Public Property Expanded() As Boolean
        Get
            Return m_Expanded
        End Get
        Set(ByVal value As Boolean)
            m_Expanded = value
            m_ParentList.Invalidate()
            RaiseEvent ExpandedChanged(Me, Me.Expanded)
        End Set
    End Property

    Public Property Font() As Font
        Get
            Return m_Font
        End Get
        Set(ByVal value As Font)
            m_Font = value
            m_ParentList.Invalidate()
        End Set
    End Property
#End Region

#Region "Rendering"
    Friend Sub RenderTo(ByVal g As Graphics, ByVal selected As Boolean, ByVal indent As Integer, ByRef y As Integer)
        'Fill in the background color.
        Dim h As Integer = GetHeight(g)
        Dim r As New Rectangle(0, y, m_ParentList.Width, h)
        g.FillRectangle(New SolidBrush(If(selected, Me.SelectedBackColor, Me.BackColor)), r)

        'Draw an overline.
        g.DrawLine(Pens.Black, r.Left, r.Top, r.Right, r.Top)

        'Draw the icon, if applicable.
        If Me.Icon IsNot Nothing Then
            g.DrawImage(Me.Icon, indent, y + 2)
            indent += Me.Icon.Width + 2
        End If

        'Draw the text
        Dim rt As New Rectangle(indent, y, m_ParentList.Width, h)
        g.DrawString(Me.Text, GetFont(), New SolidBrush(If(selected, Me.SelectedForeColor, Me.ForeColor)), rt, New StringFormat() With {.LineAlignment = StringAlignment.Center})

        'Increase the drawing location.
        y += h
    End Sub

    Friend Function GetHeight(ByVal g As Graphics) As Integer
        Dim r As Integer = Me.Height + 4
        If Me.Height = 0 Then r = CInt(g.MeasureString(Me.Text, GetFont()).Height) + 4
        If Me.Icon IsNot Nothing Then r = Math.Max(Me.Icon.Height + 4, r)
        Return r
    End Function

    Private Function GetFont() As Font
        If Me.Font IsNot Nothing Then
            Return Me.Font
        Else
            Return m_ParentList.Font
        End If
    End Function
#End Region

#Region "Clone"
    Public Function Clone() As Object Implements System.ICloneable.Clone
        Dim r As TreeListNode = DirectCast(Me.MemberwiseClone(), TreeListNode)
        r.m_Children = New TreeList.TreeListNodeCollection(m_ParentList)
        For Each cn As TreeListNode In Me.ChildNodes
            r.ChildNodes.Add(DirectCast(cn.Clone(), TreeListNode))
        Next
        Return r
    End Function
#End Region

End Class
#End Region