mimic/frontend/src/components/MitreMatrixModal.tsx

import { useEffect, useRef, useState } from 'react';
import { LoadingState } from './LoadingState';
import { ErrorState } from './ErrorState';
import { extractApiError } from '@/api/client';
import { useMitreMatrix } from '@/hooks/useMitre';
import type { MitreTechnique, MitreTacticRef } from '@/api/types';

export interface MatrixSelection {
  techniques: MitreTechnique[];
  tactics: MitreTacticRef[];
}

interface MitreMatrixModalProps {
  isOpen: boolean;
  initialTechniques: MitreTechnique[];
  initialTactics: MitreTacticRef[];
  onApply: (selection: MatrixSelection) => void;
  onCancel: () => void;
}

function countSelected(
  techniques: { id: string; subtechniques: { id: string }[] }[],
  techMap: Set<string>,
  tacticId: string,
  tacticMap: Set<string>,
): number {
  let count = tacticMap.has(tacticId) ? 1 : 0;
  for (const t of techniques) {
    if (techMap.has(t.id)) count++;
    for (const s of t.subtechniques) {
      if (techMap.has(s.id)) count++;
    }
  }
  return count;
}

export function MitreMatrixModal({
  isOpen,
  initialTechniques,
  initialTactics,
  onApply,
  onCancel,
}: MitreMatrixModalProps): JSX.Element | null {
  const { data: matrix, isLoading, isError, error } = useMitreMatrix(isOpen);

  const [selectedTechMap, setSelectedTechMap] = useState<Map<string, { id: string; name: string }>>(
    () => new Map(initialTechniques.map((t) => [t.id, { id: t.id, name: t.name }])),
  );
  const [selectedTacticSet, setSelectedTacticSet] = useState<Set<string>>(
    () => new Set(initialTactics.map((t) => t.id)),
  );
  const [expandedTechniques, setExpandedTechniques] = useState<Set<string>>(new Set());
  const [search, setSearch] = useState('');

  const containerRef = useRef<HTMLDivElement>(null);
  const searchInputRef = useRef<HTMLInputElement>(null);

  useEffect(() => {
    if (isOpen) {
      setSelectedTechMap(new Map(initialTechniques.map((t) => [t.id, { id: t.id, name: t.name }])));
      setSelectedTacticSet(new Set(initialTactics.map((t) => t.id)));
      setExpandedTechniques(new Set());
      setSearch('');
    }
  }, [isOpen]); // eslint-disable-line react-hooks/exhaustive-deps

  useEffect(() => {
    if (isOpen) {
      setTimeout(() => searchInputRef.current?.focus(), 0);
    }
  }, [isOpen]);

  useEffect(() => {
    if (!isOpen) return;
    const handler = (e: KeyboardEvent) => {
      if (e.key === 'Escape') onCancel();
    };
    document.addEventListener('keydown', handler);
    return () => document.removeEventListener('keydown', handler);
  }, [isOpen, onCancel]);

  const getFocusableElements = () => {
    if (!containerRef.current) return [];
    return Array.from(
      containerRef.current.querySelectorAll<HTMLElement>(
        'a, button, input, [tabindex]:not([tabindex="-1"])',
      ),
    ).filter((el) => !(el as HTMLButtonElement | HTMLInputElement).disabled && !el.hidden && el.tabIndex !== -1);
  };

  const handleKeyDown = (e: React.KeyboardEvent) => {
    if (e.key !== 'Tab') return;
    const focusables = getFocusableElements();
    if (focusables.length === 0) return;
    const first = focusables[0];
    const last = focusables[focusables.length - 1];
    if (e.shiftKey) {
      if (document.activeElement === first) { e.preventDefault(); last.focus(); }
    } else {
      if (document.activeElement === last) { e.preventDefault(); first.focus(); }
    }
  };

  if (!isOpen) return null;

  const toggleTechnique = (id: string, name: string) => {
    setSelectedTechMap((prev) => {
      const next = new Map(prev);
      if (next.has(id)) next.delete(id); else next.set(id, { id, name });
      return next;
    });
  };

  const toggleTactic = (tacticId: string) => {
    setSelectedTacticSet((prev) => {
      const next = new Set(prev);
      if (next.has(tacticId)) next.delete(tacticId); else next.add(tacticId);
      return next;
    });
  };

  const toggleExpand = (id: string) => {
    setExpandedTechniques((prev) => {
      const next = new Set(prev);
      if (next.has(id)) next.delete(id); else next.add(id);
      return next;
    });
  };

  const searchLower = search.toLowerCase().trim();

  const autoExpanded = new Set<string>();
  if (searchLower && matrix) {
    for (const tactic of matrix) {
      for (const tech of tactic.techniques) {
        const subMatch = tech.subtechniques.some(
          (s) => s.id.toLowerCase().includes(searchLower) || s.name.toLowerCase().includes(searchLower),
        );
        if (subMatch) autoExpanded.add(tech.id);
      }
    }
  }

  const handleApply = () => {
    const techniques: MitreTechnique[] = Array.from(selectedTechMap.values()).map((t) => ({
      id: t.id,
      name: t.name,
      tactics: [],
    }));
    // Reconstruct tactic refs from matrix data
    const tactics: MitreTacticRef[] = matrix
      ? matrix
          .filter((t) => selectedTacticSet.has(t.tactic_id))
          .map((t) => ({ id: t.tactic_id, name: t.tactic_name }))
      : Array.from(selectedTacticSet).map((id) => ({ id, name: id }));
    onApply({ techniques, tactics });
  };

  const totalTechSelected = selectedTechMap.size;
  const totalTacticSelected = selectedTacticSet.size;
  const totalSelected = totalTechSelected + totalTacticSelected;
  const hasInitial = initialTechniques.length + initialTactics.length > 0;

  return (
    <div className="fixed inset-0 z-50 flex items-center justify-center">
      <div className="absolute inset-0 bg-ink/60" onClick={onCancel} aria-hidden="true" />

      <div
        ref={containerRef}
        role="dialog"
        aria-modal="true"
        aria-labelledby="matrix-modal-title"
        className="relative bg-canvas rounded-xl shadow-floating max-w-[98vw] max-h-[80vh] overflow-hidden flex flex-col"
        style={{ width: '1400px' }}
        onKeyDown={handleKeyDown}
      >
        {/* Header */}
        <div className="flex items-center justify-between px-xl py-md border-b border-hairline flex-shrink-0">
          <h2 id="matrix-modal-title" className="text-[18px] font-medium text-ink">
            MITRE ATT&amp;CK Matrix
          </h2>
          <input
            ref={searchInputRef}
            type="text"
            placeholder="Filter techniques…"
            value={search}
            onChange={(e) => setSearch(e.target.value)}
            className="text-input w-56 h-9 text-[14px]"
            aria-label="Filter techniques"
          />
        </div>

        {/* Body — overflow-y-auto, NO overflow-x */}
        <div className="flex-1 overflow-y-auto overflow-x-hidden px-md py-md">
          {isLoading && <LoadingState label="Loading MITRE matrix…" />}
          {isError && (
            <ErrorState message={extractApiError(error, 'Could not load MITRE matrix')} />
          )}
          {!isLoading && !isError && matrix && (
            <div
              className="grid gap-xxs"
              style={{ gridTemplateColumns: `repeat(${matrix.length}, minmax(0, 1fr))` }}
            >
              {matrix.map((tactic) => {
                const tacticSelected = selectedTacticSet.has(tactic.tactic_id);
                const selectedCount = countSelected(
                  tactic.techniques,
                  new Set(selectedTechMap.keys()),
                  tactic.tactic_id,
                  selectedTacticSet,
                );

                const visibleTechniques = tactic.techniques.filter((tech) => {
                  if (!searchLower) return true;
                  const techMatch =
                    tech.id.toLowerCase().includes(searchLower) ||
                    tech.name.toLowerCase().includes(searchLower);
                  const subMatch = tech.subtechniques.some(
                    (s) =>
                      s.id.toLowerCase().includes(searchLower) ||
                      s.name.toLowerCase().includes(searchLower),
                  );
                  return techMatch || subMatch;
                });

                return (
                  <div key={tactic.tactic_id} className="flex flex-col min-w-0">
                    {/* Tactic header — clickable to toggle tactic selection */}
                    <button
                      type="button"
                      onClick={() => toggleTactic(tactic.tactic_id)}
                      title={`${tactic.tactic_name} (${tactic.tactic_id}) — click to tag this tactic`}
                      className={`w-full text-left px-xs py-xxs rounded-t-sm border border-b-0 border-hairline transition-colors ${
                        tacticSelected
                          ? 'bg-primary border-primary'
                          : 'bg-cloud hover:bg-fog'
                      }`}
                    >
                      <div className={`text-[10px] uppercase tracking-[0.6px] font-semibold leading-tight truncate ${
                        tacticSelected ? 'text-white' : 'text-graphite'
                      }`}>
                        {tactic.tactic_name}
                      </div>
                      {selectedCount > 0 && (
                        <div className={`text-[10px] font-medium leading-none mt-[2px] ${
                          tacticSelected ? 'text-white/80' : 'text-primary-deep'
                        }`}>
                          {selectedCount} sel.
                        </div>
                      )}
                    </button>

                    {/* Techniques */}
                    <div className="border border-hairline rounded-b-sm overflow-hidden flex flex-col">
                      {visibleTechniques.map((tech, techIdx) => {
                        const isSelected = selectedTechMap.has(tech.id);
                        const isExpanded = expandedTechniques.has(tech.id) || autoExpanded.has(tech.id);
                        const hasSubtechniques = tech.subtechniques.length > 0;
                        const isLast = techIdx === visibleTechniques.length - 1;

                        const visibleSubs = searchLower
                          ? tech.subtechniques.filter(
                              (s) =>
                                s.id.toLowerCase().includes(searchLower) ||
                                s.name.toLowerCase().includes(searchLower),
                            )
                          : tech.subtechniques;

                        return (
                          <div key={tech.id} className={!isLast ? 'border-b border-hairline' : ''}>
                            <div
                              className={`flex items-start px-xs py-xxs text-[11px] ${
                                isSelected ? 'bg-primary' : 'bg-canvas hover:bg-cloud'
                              }`}
                            >
                              {hasSubtechniques ? (
                                <button
                                  type="button"
                                  aria-label={isExpanded ? `Collapse ${tech.id}` : `Expand ${tech.id}`}
                                  onClick={() => toggleExpand(tech.id)}
                                  className={`mr-[2px] flex-shrink-0 text-[9px] w-3 leading-none mt-[1px] ${
                                    isSelected ? 'text-white' : 'text-graphite'
                                  }`}
                                >
                                  {isExpanded ? '▾' : '▸'}
                                </button>
                              ) : (
                                <span className="mr-[2px] w-3 flex-shrink-0" />
                              )}

                              <button
                                type="button"
                                onClick={() => toggleTechnique(tech.id, tech.name)}
                                title={`${tech.id} — ${tech.name}`}
                                className={`text-left leading-tight flex-1 min-w-0 ${
                                  isSelected ? 'text-white' : 'text-ink'
                                }`}
                              >
                                <span className="font-semibold block truncate">{tech.id}</span>
                                <span className={`block truncate text-[10px] ${isSelected ? 'text-white/80' : 'text-charcoal'}`}>
                                  {tech.name}
                                </span>
                              </button>
                            </div>

                            {isExpanded &&
                              visibleSubs.map((sub) => {
                                const isSubSelected = selectedTechMap.has(sub.id);
                                return (
                                  <button
                                    key={sub.id}
                                    type="button"
                                    onClick={() => toggleTechnique(sub.id, sub.name)}
                                    title={`${sub.id} — ${sub.name}`}
                                    className={`w-full text-left pl-[14px] pr-xs py-[2px] text-[10px] border-t border-hairline leading-tight ${
                                      isSubSelected
                                        ? 'bg-primary-soft text-primary-deep'
                                        : 'bg-cloud text-charcoal hover:bg-fog'
                                    }`}
                                  >
                                    <span className="font-semibold block truncate">{sub.id}</span>
                                    <span className="block truncate">{sub.name}</span>
                                  </button>
                                );
                              })}
                          </div>
                        );
                      })}
                      {visibleTechniques.length === 0 && searchLower && (
                        <div className="px-xs py-xxs text-[10px] text-graphite italic">No match</div>
                      )}
                    </div>
                  </div>
                );
              })}
            </div>
          )}
        </div>

        {/* Footer */}
        <div className="flex items-center justify-end gap-md px-xl py-md border-t border-hairline flex-shrink-0">
          <button type="button" className="btn-outline-ink" onClick={onCancel}>
            Cancel
          </button>
          <button
            type="button"
            className="btn-primary"
            onClick={handleApply}
            disabled={isLoading || isError || (totalSelected === 0 && !hasInitial)}
          >
            {totalSelected === 0
              ? 'Clear all'
              : `Apply ${totalSelected} item${totalSelected !== 1 ? 's' : ''}`}
          </button>
        </div>
      </div>
    </div>
  );
}