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

IMPLEMENT_DYNAMIC(CExpandableListCtrl, CListCtrl)

CExpandableListCtrl::CExpandableListCtrl() {}
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();
}

void CExpandableListCtrl::Expand(Node* n)
{
    if (!n || n->children.empty()) return;
    if (!n->expanded) { n->expanded = true; RebuildVisible(); }
}

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

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

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);
            UpdateWindow();
            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);
                UpdateWindow();
                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);
    UpdateWindow();
}

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);
            }
        }
    }
    *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;
}