js/ui/tfliteViewer.js

/**
 * TFLiteViewer - Displays TFLite file content
 * Shows summary, operator list, and tensor list for parsed TFLite files.
 * Supports search, sort, and dtype color badges.
 * Requirements: 3.1, 3.2, 3.3, 3.4, 4.1, 4.2, 4.3, 5.1, 5.2, 5.3, 5.4, 5.5, 5.6
 */

class TFLiteViewer {
  /**
   * @param {string} containerId - ID of the main container (modelDetailsContainer)
   */
  constructor(containerId) {
    this._containerId = containerId;
    this._container = document.getElementById(containerId);
    this._data = null;
    this._fileName = null;
    this._isVisible = false;

    // Sort state
    this._sortField = null;  // 'name' or 'size'
    this._sortAsc = true;

    // Search state
    this._searchQuery = '';
    this._debounceTimer = null;

    // Created panel elements
    this._summaryCard = null;
    this._operatorsCard = null;
    this._tensorListCard = null;

    // Bound handlers for cleanup
    this._boundHandlers = [];

    // ONNX panel IDs to hide/show
    this._onnxPanelSelectors = [
      '#graphContainer',
      '#nodeDetailContainer',
      '#layerStatsContainer',
      '#modelComplexityContainer',
      '#opsetCheckerContainer',
      '#flopsEstimatorContainer'
    ];

    // IDs of parent cards/wrappers that contain ONNX panels
    this._onnxCardSelectors = [
      '#metadataContainer',
      '#inputOutputContainer',
      '#initializerContainer'
    ];

    // Dtype color map for TFLite types
    this._dtypeColors = {
      'FLOAT32':    '#198754',   // green
      'FLOAT16':    '#fd7e14',   // orange
      'BFLOAT16':   '#e35d6a',   // pink
      'FLOAT64':    '#0d6efd',   // blue
      'INT8':       '#6f42c1',   // purple
      'INT16':      '#6f42c1',
      'INT32':      '#6f42c1',
      'INT64':      '#6f42c1',
      'INT4':       '#6f42c1',
      'UINT8':      '#20c997',   // teal
      'UINT16':     '#20c997',
      'UINT32':     '#20c997',
      'UINT64':     '#20c997',
      'BOOL':       '#6c757d',   // gray
      'STRING':     '#dc3545',   // red
      'COMPLEX64':  '#0dcaf0',   // cyan
      'COMPLEX128': '#0dcaf0'
    };

    if (!this._container) {
      console.warn(`[TFLiteViewer] Container #${containerId} not found`);
    }
  }

  // ─── Public API ───────────────────────────────────────────────────────────

  /**
   * Render TFLite data into the container.
   * @param {TFLiteModelData} data - Parsed TFLite model data
   * @param {string} fileName
   */
  render(data, fileName) {
    if (!this._container) return;
    if (!data) {
      console.warn('[TFLiteViewer] No data provided');
      return;
    }

    this._data = data;
    this._fileName = fileName;
    this._isVisible = true;
    this._searchQuery = '';
    this._sortField = null;
    this._sortAsc = true;

    // Hide all ONNX-specific panels (Req 3.3)
    this._hideOnnxPanels();

    // Also hide SafeTensors panels if visible
    this._hideSafeTensorsPanels();

    // Remove previous TFLite panels
    this._removePanels();

    const tensors = data.tensors || [];
    const operators = data.operators || [];

    // Create Summary card (Req 3.1, 3.2)
    this._summaryCard = this._createSummaryCard(data, tensors, operators);
    this._container.prepend(this._summaryCard);

    // Create Operator List card (Req 4.1, 4.2, 4.3)
    this._operatorsCard = this._createOperatorsCard(operators);
    this._summaryCard.after(this._operatorsCard);

    // Create Tensor List card (Req 5.1, 5.2, 5.3, 5.4, 5.5, 5.6)
    this._tensorListCard = this._createTensorListCard(tensors);
    this._operatorsCard.after(this._tensorListCard);
  }

  /**
   * Hide TFLite panels and restore ONNX panels. (Req 3.4)
   */
  hide() {
    if (!this._isVisible) return;
    this._isVisible = false;
    this._removePanels();
    this._showOnnxPanels();
  }

  /**
   * Destroy: clean up event listeners and panels.
   */
  destroy() {
    this._clearDebounce();
    this._removePanels();
    this._boundHandlers.forEach(({ el, event, fn }) => {
      el.removeEventListener(event, fn);
    });
    this._boundHandlers = [];
    this._data = null;
    this._fileName = null;
    this._isVisible = false;
  }

  // ─── Private: Panel Management ────────────────────────────────────────────

  _removePanels() {
    [this._summaryCard, this._operatorsCard, this._tensorListCard].forEach(el => {
      if (el && el.parentNode) el.parentNode.removeChild(el);
    });
    this._summaryCard = null;
    this._operatorsCard = null;
    this._tensorListCard = null;
    this._clearDebounce();
  }

  _hideOnnxPanels() {
    this._onnxPanelSelectors.forEach(sel => {
      const el = document.querySelector(sel);
      if (el) {
        const wrapper = el.closest('.col-12, .col-md-6, .col-md-4, .col-md-8, .card');
        if (wrapper) {
          wrapper.setAttribute('data-tflite-hidden', 'true');
          wrapper.style.display = 'none';
        }
      }
    });

    this._onnxCardSelectors.forEach(sel => {
      const el = document.querySelector(sel);
      if (el) {
        const wrapper = el.closest('.col-12, .col-md-6');
        if (wrapper) {
          wrapper.setAttribute('data-tflite-hidden', 'true');
          wrapper.style.display = 'none';
        }
      }
    });
  }

  _showOnnxPanels() {
    const hidden = document.querySelectorAll('[data-tflite-hidden="true"]');
    hidden.forEach(el => {
      el.removeAttribute('data-tflite-hidden');
      el.style.display = '';
    });
  }

  _hideSafeTensorsPanels() {
    const stPanels = document.querySelectorAll('[data-st-panel]');
    stPanels.forEach(el => {
      el.style.display = 'none';
    });
  }

  // ─── Private: Summary Card (Req 3.1, 3.2) ─────────────────────────────────

  _createSummaryCard(data, tensors, operators) {
    const wrapper = document.createElement('div');
    wrapper.className = 'col-12 mb-4';
    wrapper.setAttribute('data-tflite-panel', 'summary');

    const totalOperators = operators.length;
    const totalTensors = tensors.length;
    const totalParams = tensors.reduce((sum, t) => {
      if (!t.shape || t.shape.length === 0) return sum;
      return sum + t.shape.reduce((p, d) => p * d, 1);
    }, 0);
    const totalBytes = tensors.reduce((sum, t) => sum + (t.byteSize || 0), 0);

    const version = data.version != null ? data.version : 'N/A';
    const description = data.description || 'N/A';

    wrapper.innerHTML = `
      <div class="card" style="border-left: 4px solid #0d6efd;">
        <div class="card-header" style="background-color: #f0f4ff;">
          <h5 class="mb-0"><i class="fas fa-chart-pie me-2"></i>TFLite Model Summary — ${this._escapeHtml(this._fileName || '')}</h5>
        </div>
        <div class="card-body">
          <div class="row text-center">
            <div class="col-md-2 mb-2">
              <div class="fw-bold fs-4">${totalOperators}</div>
              <div class="text-muted small">Operators</div>
            </div>
            <div class="col-md-2 mb-2">
              <div class="fw-bold fs-4">${totalTensors}</div>
              <div class="text-muted small">Tensors</div>
            </div>
            <div class="col-md-2 mb-2">
              <div class="fw-bold fs-4">${this._formatParamCount(totalParams)}</div>
              <div class="text-muted small">Parameters</div>
            </div>
            <div class="col-md-2 mb-2">
              <div class="fw-bold fs-4">${this._formatBytes(totalBytes)}</div>
              <div class="text-muted small">Memory Footprint</div>
            </div>
            <div class="col-md-2 mb-2">
              <div class="fw-bold fs-4">${this._escapeHtml(String(version))}</div>
              <div class="text-muted small">Version</div>
            </div>
            <div class="col-md-2 mb-2">
              <div class="fw-bold text-break" style="font-size:0.9rem;" title="${this._escapeHtml(String(description))}">${this._escapeHtml(String(description))}</div>
              <div class="text-muted small">Description</div>
            </div>
          </div>
        </div>
      </div>`;

    return wrapper;
  }

  // ─── Private: Operators Card (Req 4.1, 4.2, 4.3) ──────────────────────────

  _createOperatorsCard(operators) {
    const wrapper = document.createElement('div');
    wrapper.className = 'col-12 mb-4';
    wrapper.setAttribute('data-tflite-panel', 'operators');

    // Count occurrences of each operator
    const opCounts = {};
    operators.forEach(op => {
      const name = op.opcodeName || 'UNKNOWN';
      opCounts[name] = (opCounts[name] || 0) + 1;
    });

    const uniqueOps = Object.keys(opCounts).length;

    let bodyHtml;
    if (operators.length === 0) {
      // Req 4.3: empty operators
      bodyHtml = '<p class="text-muted mb-0">Không tìm thấy operator nào</p>';
    } else {
      const rows = Object.entries(opCounts)
        .sort((a, b) => b[1] - a[1])
        .map(([name, count]) => {
          return `<tr><td>${this._escapeHtml(name)}</td><td class="text-center">${count}</td></tr>`;
        }).join('');

      bodyHtml = `
        <div class="mb-2 text-muted small">${uniqueOps} unique operator type${uniqueOps !== 1 ? 's' : ''}</div>
        <div style="max-height:400px;overflow-y:auto;">
          <table class="table table-sm table-hover mb-0">
            <thead class="table-light sticky-top">
              <tr>
                <th>Operator</th>
                <th class="text-center">Count</th>
              </tr>
            </thead>
            <tbody>${rows}</tbody>
          </table>
        </div>`;
    }

    wrapper.innerHTML = `
      <div class="card" style="border-left: 4px solid #fd7e14;">
        <div class="card-header" style="background-color: #fff8f0;">
          <h5 class="mb-0"><i class="fas fa-cogs me-2"></i>Operator List</h5>
        </div>
        <div class="card-body">${bodyHtml}</div>
      </div>`;

    return wrapper;
  }

  // ─── Private: Tensor List Card (Req 5.1, 5.2, 5.3, 5.4, 5.5, 5.6) ────────

  _createTensorListCard(tensors) {
    const wrapper = document.createElement('div');
    wrapper.className = 'col-12 mb-4';
    wrapper.setAttribute('data-tflite-panel', 'tensors');

    wrapper.innerHTML = `
      <div class="card" style="border-left: 4px solid #198754;">
        <div class="card-header" style="background-color: #f0faf4;">
          <h5 class="mb-0"><i class="fas fa-list me-2"></i>Tensor List</h5>
        </div>
        <div class="card-body">
          <div class="d-flex flex-wrap align-items-center gap-2 mb-3">
            <div class="flex-grow-1">
              <input type="text" class="form-control form-control-sm tflite-tensor-search" placeholder="Search tensor by name...">
            </div>
            <div class="d-flex gap-1">
              <button class="btn btn-outline-secondary btn-sm tflite-sort-name" title="Sort by name">
                <i class="fas fa-sort-alpha-down me-1"></i>Name
              </button>
              <button class="btn btn-outline-secondary btn-sm tflite-sort-size" title="Sort by size">
                <i class="fas fa-sort-amount-down me-1"></i>Size
              </button>
            </div>
          </div>
          <div class="tflite-tensor-count text-muted small mb-2"></div>
          <div class="tflite-tensor-table-wrapper" style="max-height:500px;overflow-y:auto;">
            <table class="table table-sm table-hover mb-0">
              <thead class="table-light sticky-top">
                <tr>
                  <th>Name</th>
                  <th>Shape</th>
                  <th>Dtype</th>
                  <th>Size</th>
                </tr>
              </thead>
              <tbody class="tflite-tensor-tbody"></tbody>
            </table>
          </div>
        </div>
      </div>`;

    // Render initial tensor rows
    this._renderTensorRows(wrapper, tensors);

    // Attach search handler (Req 5.2 — 300ms debounce)
    const searchInput = wrapper.querySelector('.tflite-tensor-search');
    if (searchInput) {
      const handler = () => {
        this._clearDebounce();
        this._debounceTimer = setTimeout(() => {
          this._searchQuery = searchInput.value;
          this._renderTensorRows(wrapper, this._getFilteredSortedTensors());
        }, 300);
      };
      searchInput.addEventListener('input', handler);
      this._boundHandlers.push({ el: searchInput, event: 'input', fn: handler });
    }

    // Attach sort handlers (Req 5.3)
    const sortNameBtn = wrapper.querySelector('.tflite-sort-name');
    const sortSizeBtn = wrapper.querySelector('.tflite-sort-size');

    if (sortNameBtn) {
      const handler = () => {
        if (this._sortField === 'name') {
          this._sortAsc = !this._sortAsc;
        } else {
          this._sortField = 'name';
          this._sortAsc = true;
        }
        this._updateSortButtons(wrapper);
        this._renderTensorRows(wrapper, this._getFilteredSortedTensors());
      };
      sortNameBtn.addEventListener('click', handler);
      this._boundHandlers.push({ el: sortNameBtn, event: 'click', fn: handler });
    }

    if (sortSizeBtn) {
      const handler = () => {
        if (this._sortField === 'size') {
          this._sortAsc = !this._sortAsc;
        } else {
          this._sortField = 'size';
          this._sortAsc = false; // default: largest first
        }
        this._updateSortButtons(wrapper);
        this._renderTensorRows(wrapper, this._getFilteredSortedTensors());
      };
      sortSizeBtn.addEventListener('click', handler);
      this._boundHandlers.push({ el: sortSizeBtn, event: 'click', fn: handler });
    }

    return wrapper;
  }

  _getFilteredSortedTensors() {
    let tensors = (this._data && this._data.tensors) || [];

    // Filter by search query (Req 5.2)
    if (this._searchQuery.trim()) {
      const q = this._searchQuery.trim().toLowerCase();
      tensors = tensors.filter(t => (t.name || '').toLowerCase().includes(q));
    }

    // Sort (Req 5.3)
    if (this._sortField === 'name') {
      tensors = tensors.slice().sort((a, b) => {
        const cmp = (a.name || '').localeCompare(b.name || '');
        return this._sortAsc ? cmp : -cmp;
      });
    } else if (this._sortField === 'size') {
      tensors = tensors.slice().sort((a, b) => {
        const cmp = (a.byteSize || 0) - (b.byteSize || 0);
        return this._sortAsc ? cmp : -cmp;
      });
    }

    return tensors;
  }

  _renderTensorRows(wrapper, tensors) {
    const tbody = wrapper.querySelector('.tflite-tensor-tbody');
    const countEl = wrapper.querySelector('.tflite-tensor-count');
    const allTensors = (this._data && this._data.tensors) || [];

    if (!tbody) return;

    // Update count display (Req 5.5)
    if (countEl) {
      if (this._searchQuery.trim()) {
        countEl.textContent = `Showing ${tensors.length} / ${allTensors.length} tensors`;
      } else {
        countEl.textContent = `${allTensors.length} tensor${allTensors.length !== 1 ? 's' : ''}`;
      }
    }

    // Empty state: no tensors at all
    if (allTensors.length === 0) {
      tbody.innerHTML = '<tr><td colspan="4" class="text-center text-muted py-3">Không có tensor nào</td></tr>';
      return;
    }

    // Empty state: search with no results (Req 5.4)
    if (tensors.length === 0 && this._searchQuery.trim()) {
      tbody.innerHTML = '<tr><td colspan="4" class="text-center text-muted py-3">Không tìm thấy tensor phù hợp</td></tr>';
      return;
    }

    const rows = tensors.map(t => {
      const color = this._dtypeColors[t.dtype] || '#6c757d';
      const shapeStr = t.shape ? `[${t.shape.join(', ')}]` : '[]';
      const sizeStr = this._formatBytes(t.byteSize);

      return `<tr>
        <td class="text-break" style="max-width:300px;" title="${this._escapeHtml(t.name || '')}">${this._escapeHtml(t.name || '')}</td>
        <td class="text-nowrap"><code>${this._escapeHtml(shapeStr)}</code></td>
        <td><span class="badge" style="background-color:${color}">${this._escapeHtml(t.dtype || 'UNKNOWN')}</span></td>
        <td class="text-nowrap">${sizeStr}</td>
      </tr>`;
    }).join('');

    tbody.innerHTML = rows;
  }

  _updateSortButtons(wrapper) {
    const nameBtn = wrapper.querySelector('.tflite-sort-name');
    const sizeBtn = wrapper.querySelector('.tflite-sort-size');

    [nameBtn, sizeBtn].forEach(btn => {
      if (btn) {
        btn.classList.remove('btn-secondary');
        btn.classList.add('btn-outline-secondary');
      }
    });

    if (this._sortField === 'name' && nameBtn) {
      nameBtn.classList.remove('btn-outline-secondary');
      nameBtn.classList.add('btn-secondary');
      const icon = nameBtn.querySelector('i');
      if (icon) {
        icon.className = this._sortAsc ? 'fas fa-sort-alpha-down me-1' : 'fas fa-sort-alpha-up me-1';
      }
    }

    if (this._sortField === 'size' && sizeBtn) {
      sizeBtn.classList.remove('btn-outline-secondary');
      sizeBtn.classList.add('btn-secondary');
      const icon = sizeBtn.querySelector('i');
      if (icon) {
        icon.className = this._sortAsc ? 'fas fa-sort-amount-up me-1' : 'fas fa-sort-amount-down me-1';
      }
    }
  }

  // ─── Private: Formatters ──────────────────────────────────────────────────

  /**
   * Format parameter count: 25600000 → "25.6M", 1200000000 → "1.2B"
   * @param {number} count
   * @returns {string}
   */
  _formatParamCount(count) {
    if (count == null || isNaN(count) || count === 0) return '0';
    const abs = Math.abs(count);
    if (abs >= 1e9) return (count / 1e9).toFixed(1).replace(/\.0$/, '') + 'B';
    if (abs >= 1e6) return (count / 1e6).toFixed(1).replace(/\.0$/, '') + 'M';
    if (abs >= 1e3) return (count / 1e3).toFixed(1).replace(/\.0$/, '') + 'K';
    return String(count);
  }

  /**
   * Format byte size: 512000000 → "488.28 MB"
   * @param {number} bytes
   * @returns {string}
   */
  _formatBytes(bytes) {
    if (bytes == null || isNaN(bytes) || bytes === 0) return '0 B';
    const k = 1024;
    const sizes = ['B', 'KB', 'MB', 'GB', 'TB'];
    const i = Math.min(Math.floor(Math.log(Math.abs(bytes)) / Math.log(k)), sizes.length - 1);
    return (bytes / Math.pow(k, i)).toFixed(2).replace(/\.00$/, '') + ' ' + sizes[i];
  }

  /**
   * Escape HTML special characters.
   * @param {string} str
   * @returns {string}
   */
  _escapeHtml(str) {
    const div = document.createElement('div');
    div.appendChild(document.createTextNode(String(str)));
    return div.innerHTML;
  }

  _clearDebounce() {
    if (this._debounceTimer) {
      clearTimeout(this._debounceTimer);
      this._debounceTimer = null;
    }
  }
}

window.TFLiteViewer = TFLiteViewer;