#include "stdafx.h"
#include "CExpandableListCtrl.h"

IMPLEMENT_DYNAMIC(CExpandableListCtrl, CListCtrl)

CExpandableListCtrl::CExpandableListCtrl()
{
    m_popupCols = { };
}

CExpandableListCtrl::~CExpandableListCtrl() {}

BEGIN_MESSAGE_MAP(CExpandableListCtrl, CListCtrl)
    ON_WM_CREATE()
    ON_NOTIFY_REFLECT(NM_CLICK, &CExpandableListCtrl::OnClick)
    ON_NOTIFY_REFLECT(NM_CUSTOMDRAW, &CExpandableListCtrl::OnCustomDraw)
END_MESSAGE_MAP()

int CExpandableListCtrl::OnCreate(LPCREATESTRUCT lpCreateStruct)
{
    if (CListCtrl::OnCreate(lpCreateStruct) == -1)
        return -1;

    SetExtendedStyle(GetExtendedStyle()
        | LVS_EX_FULLROWSELECT | LVS_EX_HEADERDRAGDROP | LVS_EX_GRIDLINES | LVS_EX_DOUBLEBUFFER);

    return 0;
}

void CExpandableListCtrl::PreSubclassWindow()
{
    SetExtendedStyle(GetExtendedStyle()
        | LVS_EX_FULLROWSELECT | LVS_EX_HEADERDRAGDROP | LVS_EX_GRIDLINES | LVS_EX_DOUBLEBUFFER);
    CListCtrl::PreSubclassWindow();
}

// ===== 树 API =====
CExpandableListCtrl::Node* CExpandableListCtrl::InsertRoot(const std::vector<CString>& cols)
{
    auto n = std::make_unique<Node>((int)max(1, (int)cols.size()));
    for (size_t i = 0; i < cols.size(); ++i) n->cols[i] = cols[i];
    n->level = 0;
    Node* raw = n.get();
    m_roots.emplace_back(std::move(n));
    return raw;
}

CExpandableListCtrl::Node* CExpandableListCtrl::InsertChild(Node* parent, const std::vector<CString>& cols)
{
    ASSERT(parent);
    auto n = std::make_unique<Node>((int)max(1, (int)cols.size()));
    for (size_t i = 0; i < cols.size(); ++i) n->cols[i] = cols[i];
    n->parent = parent;
    n->level = parent->level + 1;
    Node* raw = n.get();
    parent->children.emplace_back(std::move(n));
    return raw;
}

void CExpandableListCtrl::appendVisible(Node* n)
{
    m_visible.push_back(n);
    if (n->expanded) {
        for (auto& ch : n->children) {
            appendVisible(ch.get());
        }
    }
}

void CExpandableListCtrl::RebuildVisible()
{
    m_visible.clear();
    for (auto& r : m_roots) appendVisible(r.get());

    SetRedraw(FALSE);
    DeleteAllItems();

    for (int i = 0; i < (int)m_visible.size(); ++i) {
        Node* n = m_visible[i];
        LVITEM lvi{};
        lvi.mask = LVIF_TEXT;
        lvi.iItem = i;
        lvi.iSubItem = 0;
        lvi.pszText = const_cast<LPTSTR>((LPCTSTR)(n->cols.empty() ? _T("") : n->cols[0]));
        InsertItem(&lvi);

        const int colCount = GetHeaderCtrl() ? GetHeaderCtrl()->GetItemCount() : 1;
        for (int col = 1; col < colCount; ++col) {
            CString txt = (col < (int)n->cols.size()) ? n->cols[col] : _T("");
            SetItemText(i, col, txt);
        }
    }

    // 重建后，按行号颜色数组对齐
    m_rowColors.resize(GetItemCount());

    SetRedraw(TRUE);
    Invalidate(FALSE);
}

// —— 优化后的展开/收起：局部插入/删除，不全量 RebuildVisible —— //
void CExpandableListCtrl::Expand(Node* n)
{
    if (!n || n->children.empty()) return;
    if (n->expanded) return;

    // 本地工具：找节点在 m_visible 中的行号
    auto VisibleIndexOf = [&](Node* x)->int {
        for (int i = 0; i < (int)m_visible.size(); ++i)
            if (m_visible[i] == x) return i;
        return -1;
    };
    // 递归收集“应当可见”的子树（受 expanded 影响）
    std::vector<Node*> toInsert;
    std::function<void(Node*)> CollectExpandedSubtree = [&](Node* x) {
        if (!x) return;
        for (auto& up : x->children) {
            Node* ch = up.get();
            toInsert.push_back(ch);
            if (ch->expanded && !ch->children.empty())
                CollectExpandedSubtree(ch);
        }
    };
    // 从 pos 起插入 nodes，对齐 m_visible / ListCtrl / m_rowColors
    auto InsertRowsAt = [&](int pos, const std::vector<Node*>& nodes) {
        if (nodes.empty()) return;
        const int colCount = GetHeaderCtrl() ? GetHeaderCtrl()->GetItemCount() : 1;

        SetRedraw(FALSE);

        // 1) 先插 m_visible
        m_visible.insert(m_visible.begin() + pos, nodes.begin(), nodes.end());

        // 2) 再插 ListCtrl
        for (int i = 0; i < (int)nodes.size(); ++i) {
            Node* cur = nodes[i];
            LVITEM lvi{}; lvi.mask = LVIF_TEXT;
            lvi.iItem = pos + i;
            lvi.iSubItem = 0;
            lvi.pszText = const_cast<LPTSTR>((LPCTSTR)(cur->cols.empty() ? _T("") : cur->cols[0]));
            InsertItem(&lvi);

            for (int col = 1; col < colCount; ++col) {
                CString txt = (col < (int)cur->cols.size()) ? cur->cols[col] : _T("");
                SetItemText(pos + i, col, txt);
            }
        }

        // 3) 行号颜色数组同步插入默认色
        m_rowColors.insert(m_rowColors.begin() + pos, nodes.size(), RowColor{});

        SetRedraw(TRUE);
        Invalidate(FALSE);
    };

    // —— 标记展开
    n->expanded = true;

    // —— 在 UI 里插入其“应当可见”的子树
    const int pos = VisibleIndexOf(n);
    if (pos < 0) { RebuildVisible(); return; }

    CollectExpandedSubtree(n);
    InsertRowsAt(pos + 1, toInsert);
}

void CExpandableListCtrl::Collapse(Node* n)
{
    if (!n || n->children.empty()) return;
    if (!n->expanded) return;

    // 本地工具：找节点行号
    auto VisibleIndexOf = [&](Node* x)->int {
        for (int i = 0; i < (int)m_visible.size(); ++i)
            if (m_visible[i] == x) return i;
        return -1;
    };
    // 计算“当前可见的所有后代数量”（基于 level 递减判断）
    auto CountDescendantsInVisible = [&](Node* x)->int {
        if (!x) return 0;
        const int start = VisibleIndexOf(x);
        if (start < 0) return 0;
        const int baseLevel = x->level;
        int cnt = 0;
        for (int i = start + 1; i < (int)m_visible.size(); ++i) {
            if (!m_visible[i]) break;
            if (m_visible[i]->level <= baseLevel) break;
            ++cnt;
        }
        return cnt;
    };
    // 从 UI 删除 pos 开始的 count 行，并同步 m_visible/m_rowColors
    auto DeleteRowsAt = [&](int pos, int count) {
        if (count <= 0) return;

        SetRedraw(FALSE);

        // 删 ListCtrl：一直删 pos，因为删一行后后续上移
        for (int i = 0; i < count; ++i) {
            DeleteItem(pos);
        }
        // 删 m_visible
        m_visible.erase(m_visible.begin() + pos, m_visible.begin() + pos + count);
        // 删颜色
        if (pos >= 0 && pos <= (int)m_rowColors.size()) {
            int end = min((int)m_rowColors.size(), pos + count);
            m_rowColors.erase(m_rowColors.begin() + pos, m_rowColors.begin() + end);
        }

        SetRedraw(TRUE);
        Invalidate(FALSE);
    };

    // —— 标记收起
    n->expanded = false;

    // —— 只删除其“当前可见”的所有后代
    const int pos = VisibleIndexOf(n);
    if (pos < 0) { RebuildVisible(); return; }

    const int cnt = CountDescendantsInVisible(n);
    if (cnt > 0) {
        DeleteRowsAt(pos + 1, cnt);
    }
}

void CExpandableListCtrl::Toggle(Node* n)
{
    if (!n || n->children.empty()) return;
    if (n->expanded) Collapse(n);
    else             Expand(n);
}

CExpandableListCtrl::Node* CExpandableListCtrl::GetNodeByVisibleIndex(int i) const
{
    if (i < 0 || i >= (int)m_visible.size()) return nullptr;
    return m_visible[i];
}

// ===== 颜色 API =====
void CExpandableListCtrl::SetNodeColor(Node* n, COLORREF text, COLORREF bk)
{
    if (!n) return;
    RowColor rc{};
    rc.text = text; rc.bk = bk;
    rc.hasText = (text != CLR_DEFAULT);
    rc.hasBk = (bk != CLR_DEFAULT);
    m_colorByNode[n] = rc;

    for (int i = 0; i < (int)m_visible.size(); ++i) {
        if (m_visible[i] == n) {
            RedrawItems(i, i);
            return;
        }
    }
}

void CExpandableListCtrl::ClearNodeColor(Node* n)
{
    if (!n) return;
    auto it = m_colorByNode.find(n);
    if (it != m_colorByNode.end()) {
        m_colorByNode.erase(it);
        for (int i = 0; i < (int)m_visible.size(); ++i) {
            if (m_visible[i] == n) {
                RedrawItems(i, i);
                return;
            }
        }
    }
}

void CExpandableListCtrl::ClearAllColors()
{
    m_colorByNode.clear();
    m_rowColors.clear();
    Invalidate(FALSE);
}

// 兼容旧接口：按“可见行号”着色
void CExpandableListCtrl::SetItemColor(DWORD_PTR iItem, COLORREF TextColor, COLORREF TextBkColor)
{
    SetItemColorByVisibleIndex((int)iItem, TextColor, TextBkColor);
}
void CExpandableListCtrl::SetItemColorByVisibleIndex(int row, COLORREF text, COLORREF bk)
{
    if (row < 0) return;
    if (row >= (int)m_rowColors.size())
        m_rowColors.resize(row + 1);

    RowColor rc{};
    rc.text = text; rc.bk = bk;
    rc.hasText = (text != CLR_DEFAULT);
    rc.hasBk = (bk != CLR_DEFAULT);
    m_rowColors[row] = rc;

    RedrawItems(row, row);
}

CRect CExpandableListCtrl::expanderRectForRow(int row) const
{
    CRect rcLabel;
    if (!const_cast<CExpandableListCtrl*>(this)->GetSubItemRect(row, 0, LVIR_LABEL, rcLabel))
        return CRect(0, 0, 0, 0);

    Node* n = const_cast<CExpandableListCtrl*>(this)->GetNodeByVisibleIndex(row);
    if (!n || n->children.empty())
        return CRect(0, 0, 0, 0);

    const int indent = n->level;
    const int left = rcLabel.left + m_expanderPadding + indent * 16;

    return CRect(
        left,
        rcLabel.CenterPoint().y - m_expanderSize / 2,
        left + m_expanderSize,
        rcLabel.CenterPoint().y + m_expanderSize / 2
    );
}

// 颜色计算：优先 Node*，其次行号；若需要则让系统高亮覆盖
void CExpandableListCtrl::computeColorsForRow(int row, COLORREF& outText, COLORREF& outBk) const
{
    outText = ListView_GetTextColor(const_cast<CExpandableListCtrl*>(this)->m_hWnd);
    outBk = ListView_GetBkColor(const_cast<CExpandableListCtrl*>(this)->m_hWnd);

    const bool selected = (const_cast<CExpandableListCtrl*>(this)->GetItemState(row, LVIS_SELECTED) & LVIS_SELECTED) != 0;
    const bool focusOnCtrl = (const_cast<CExpandableListCtrl*>(this)->GetSafeHwnd() == ::GetFocus());

    if (m_preserveSelHighlight && selected) {
        outBk = GetSysColor(focusOnCtrl ? COLOR_HIGHLIGHT : COLOR_3DFACE);
        outText = GetSysColor(focusOnCtrl ? COLOR_HIGHLIGHTTEXT : COLOR_WINDOWTEXT);
        return;
    }

    // Node* 颜色
    if (Node* n = const_cast<CExpandableListCtrl*>(this)->GetNodeByVisibleIndex(row)) {
        auto it = m_colorByNode.find(n);
        if (it != m_colorByNode.end()) {
            if (it->second.hasText) outText = it->second.text;
            if (it->second.hasBk)   outBk = it->second.bk;
            return;
        }
    }

    // 行号颜色
    if (row >= 0 && row < (int)m_rowColors.size()) {
        const RowColor& rc = m_rowColors[row];
        if (rc.hasText) outText = rc.text;
        if (rc.hasBk)   outBk = rc.bk;
    }
}

void CExpandableListCtrl::OnClick(NMHDR* pNMHDR, LRESULT* pResult)
{
    LPNMITEMACTIVATE pia = reinterpret_cast<LPNMITEMACTIVATE>(pNMHDR);
    if (pia->iItem >= 0) {
        CPoint pt = pia->ptAction;
        CRect expRc = expanderRectForRow(pia->iItem);
        if (expRc.PtInRect(pt)) {
            Node* n = GetNodeByVisibleIndex(pia->iItem);
            if (n && !n->children.empty()) {
                Toggle(n);
            }
        }
    }

    // —— 若点击到需要“全文显示”的列，则向父窗口发送自定义通知 —— //
    if (!m_popupCols.empty()) {
        LPNMITEMACTIVATE pia2 = reinterpret_cast<LPNMITEMACTIVATE>(pNMHDR);

        // 用 SubItemHitTest 更精准拿到列
        LVHITTESTINFO ht{};
        ht.pt = pia2->ptAction;
        int hit = SubItemHitTest(&ht);
        if (hit >= 0 && ht.iItem >= 0 && ht.iSubItem >= 0) {
            const int row = ht.iItem;
            const int col = ht.iSubItem;

            if (m_popupCols.count(col)) {
                CString full = GetItemText(row, col);
                if (!full.IsEmpty() && _IsCellTruncated(row, col, full)) {
                    NMC_ELC_SHOWFULLTEXT nm{};
                    nm.hdr.hwndFrom = m_hWnd;
                    nm.hdr.idFrom = GetDlgCtrlID();
                    nm.hdr.code = ELCN_SHOWFULLTEXT;
                    nm.iItem = row;
                    nm.iSubItem = col;
                    nm.text = full;

                    if (CWnd* pParent = GetParent()) {
                        pParent->SendMessage(WM_NOTIFY, nm.hdr.idFrom, reinterpret_cast<LPARAM>(&nm));
                    }
                }
            }
        }
    }

    *pResult = 0;
}

void CExpandableListCtrl::OnCustomDraw(NMHDR* pNMHDR, LRESULT* pResult)
{
    LPNMLVCUSTOMDRAW pCD = reinterpret_cast<LPNMLVCUSTOMDRAW>(pNMHDR);

    switch (pCD->nmcd.dwDrawStage)
    {
    case CDDS_PREPAINT:
        *pResult = CDRF_NOTIFYITEMDRAW | CDRF_NOTIFYSUBITEMDRAW;
        return;

    case CDDS_ITEMPREPAINT:
    {
        const int row = (int)pCD->nmcd.dwItemSpec;
        COLORREF txt, bk;
        computeColorsForRow(row, txt, bk);
        pCD->clrText = txt;
        pCD->clrTextBk = bk;
        *pResult = CDRF_NOTIFYSUBITEMDRAW;
        return;
    }

    case CDDS_ITEMPREPAINT | CDDS_SUBITEM:
    {
        const int row = (int)pCD->nmcd.dwItemSpec;
        const int col = pCD->iSubItem;
        CDC* pDC = CDC::FromHandle(pCD->nmcd.hdc);

        // 如果没有树节点（纯平列表），首列也走默认绘制
        Node* n = GetNodeByVisibleIndex(row);
        if (col != 0 || !n) {
            *pResult = CDRF_DODEFAULT;
            return;
        }

        // 首列自绘（树模式）
        CRect rc; GetSubItemRect(row, 0, LVIR_LABEL, rc);

        COLORREF bk, txt;
        computeColorsForRow(row, txt, bk);

        CBrush bkBrush(bk);
        pDC->FillRect(rc, &bkBrush);

        if (!n->children.empty())
        {
            CRect box = expanderRectForRow(row);
            const int ROFFSET = 2;
            const int WIDE = max(9, min(min(box.Width(), box.Height()), 13));
            const int WIDE2 = WIDE / 2;
            const int EXPANDED_WIDE = WIDE;

            box.DeflateRect(1, 1);

            POINT pt[3];
            if (n->expanded) {
                int nBottomOffset = (box.Height() - EXPANDED_WIDE) / 2;
                pt[0].x = box.right - ROFFSET - EXPANDED_WIDE;
                pt[0].y = box.bottom - nBottomOffset;
                pt[1].x = box.right - ROFFSET;
                pt[1].y = box.bottom - nBottomOffset;
                pt[2].x = box.right - ROFFSET;
                pt[2].y = box.bottom - nBottomOffset - EXPANDED_WIDE;
            }
            else {
                int nBottomOffset = (box.Height() - WIDE) / 2;
                pt[0].x = box.right - ROFFSET - WIDE2;
                pt[0].y = box.bottom - nBottomOffset - WIDE;
                pt[1].x = box.right - ROFFSET - WIDE2;
                pt[1].y = box.bottom - nBottomOffset;
                pt[2].x = box.right - ROFFSET;
                pt[2].y = box.bottom - nBottomOffset - WIDE2;
            }

            HGDIOBJ oldPen = pDC->SelectObject(GetStockObject(NULL_PEN));
            HBRUSH   hBrush = CreateSolidBrush(txt);
            HGDIOBJ oldBrush = pDC->SelectObject(hBrush);
            pDC->Polygon(pt, 3);
            pDC->SelectObject(oldPen);
            pDC->SelectObject(oldBrush);
            DeleteObject(hBrush);
        }

        const int indentPx = n->level * 14;
        const int baseLeft = rc.left + m_expanderPadding + indentPx;

        CRect textRc = rc;
        if (!n->children.empty()) {
            textRc.left = baseLeft + m_expanderSize + m_textGap;
        }
        else {
            constexpr int kLeafGap = 2;
            textRc.left = baseLeft + kLeafGap;
        }

        pDC->SetBkMode(TRANSPARENT);
        pDC->SetTextColor(txt);
        CString txt0 = n->cols.empty() ? _T("") : n->cols[0];
        pDC->DrawText(txt0, textRc, DT_SINGLELINE | DT_VCENTER | DT_NOPREFIX | DT_END_ELLIPSIS);

        if (GetExtendedStyle() & LVS_EX_GRIDLINES)
        {
            CRect rcRow; GetSubItemRect(row, 0, LVIR_BOUNDS, rcRow);
            const int y = rcRow.bottom - 1;
            CPen pen(PS_SOLID, 1, GetSysColor(COLOR_3DLIGHT));
            CPen* oldPen = pDC->SelectObject(&pen);
            pDC->MoveTo(rcRow.left, y);
            pDC->LineTo(rcRow.right, y);
            pDC->SelectObject(oldPen);
        }

        *pResult = CDRF_SKIPDEFAULT;
        return;
    }
    }

    *pResult = CDRF_DODEFAULT;
}

// 兼容行为：同步 SetItemText 到 Node->cols；维护行号颜色数组的插入/删除
LRESULT CExpandableListCtrl::WindowProc(UINT message, WPARAM wParam, LPARAM lParam)
{
    // 同步 SetItemText 到 Node（A/W 兼容）
    if (message == LVM_SETITEMTEXT
#ifdef LVM_SETITEMTEXTA
        || message == LVM_SETITEMTEXTA
#endif
#ifdef LVM_SETITEMTEXTW
        || message == LVM_SETITEMTEXTW
#endif
        )
    {
        int row = static_cast<int>(wParam);
        LVITEM* p = reinterpret_cast<LVITEM*>(lParam);
        if (p) {
            Node* n = GetNodeByVisibleIndex(row);
            if (n) {
                int sub = p->iSubItem;
                if (sub >= (int)n->cols.size())
                    n->cols.resize(sub + 1);
                CString newText = p->pszText ? p->pszText : _T("");
                n->cols[sub] = newText;
            }
        }
        // 继续交给基类处理
    }

    LRESULT nRet = CListCtrl::WindowProc(message, wParam, lParam);

    // 维护行号颜色数组（兼容旧 SetItemColor）
    if (message == LVM_INSERTITEM) {
        if (nRet != -1) {
            LVITEM* p = (LVITEM*)lParam;
            int pos = p ? p->iItem : (int)nRet;
            if (pos < 0) pos = (int)nRet;
            if (pos > (int)m_rowColors.size()) pos = (int)m_rowColors.size();
            m_rowColors.insert(m_rowColors.begin() + pos, RowColor{}); // 默认色
        }
    }
    else if (message == LVM_DELETEITEM) {
        if (nRet != 0) {
            int pos = (int)wParam;
            if (pos >= 0 && pos < (int)m_rowColors.size())
                m_rowColors.erase(m_rowColors.begin() + pos);
        }
    }
    else if (message == LVM_DELETEALLITEMS) {
        if (nRet != 0) {
            m_rowColors.clear();
        }
    }

    return nRet;
}

// CExpandableListCtrl.cpp 里
void CExpandableListCtrl::ClearTree()
{
    // 清数据
    m_roots.clear();
    m_visible.clear();

    // 清可见项（务必！否则旧页的行会残留）
    SetRedraw(FALSE);
    DeleteAllItems();
    SetRedraw(TRUE);

    Invalidate(FALSE);
}

void CExpandableListCtrl::SetPopupFullTextColumns(const std::vector<int>& cols)
{
    m_popupCols.clear();
    for (int c : cols) m_popupCols.insert(c);
}

bool CExpandableListCtrl::_IsCellTruncated(int row, int col, const CString& text) const
{
    if (text.IsEmpty()) return false;

    // 单元格显示区域宽度
    CRect rcCell;
    if (!const_cast<CExpandableListCtrl*>(this)->GetSubItemRect(row, col, LVIR_BOUNDS, rcCell))
        return false;

    // 用控件字体测量文本像素宽
    CClientDC dc(const_cast<CExpandableListCtrl*>(this));
    CFont* pOld = dc.SelectObject(const_cast<CExpandableListCtrl*>(this)->GetFont());
    CSize sz = dc.GetTextExtent(text);
    dc.SelectObject(pOld);

    const int kPadding = 8; // 预留一点边距/省略号余量
    return sz.cx > (rcCell.Width() - kPadding);
}