Aegis/frontend/src/components/AddTestToCampaignModal.tsx

import { useState, useMemo } from "react";
import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import {
  X,
  Search,
  Plus,
  Loader2,
  CheckCircle,
  FlaskConical,
} from "lucide-react";
import { getTests } from "../api/tests";
import { addTestToCampaign } from "../api/campaigns";
import type { Test, TestState } from "../types/models";

const stateBadge: Record<TestState, string> = {
  draft: "bg-gray-800/50 text-gray-400 border-gray-600/30",
  red_executing: "bg-orange-900/50 text-orange-400 border-orange-500/30",
  blue_evaluating: "bg-indigo-900/50 text-indigo-400 border-indigo-500/30",
  in_review: "bg-blue-900/50 text-blue-400 border-blue-500/30",
  validated: "bg-green-900/50 text-green-400 border-green-500/30",
  rejected: "bg-red-900/50 text-red-400 border-red-500/30",
};

interface AddTestToCampaignModalProps {
  campaignId: string;
  existingTestIds: string[];
  open: boolean;
  onClose: () => void;
  onSuccess: () => void;
}

export default function AddTestToCampaignModal({
  campaignId,
  existingTestIds,
  open,
  onClose,
  onSuccess,
}: AddTestToCampaignModalProps) {
  const queryClient = useQueryClient();
  const [searchText, setSearchText] = useState("");
  const [addedIds, setAddedIds] = useState<Set<string>>(new Set());

  const { data: allTests, isLoading } = useQuery({
    queryKey: ["tests", "for-campaign-picker"],
    queryFn: () => getTests({ limit: 200 }),
    enabled: open,
  });

  const filteredTests = useMemo(() => {
    if (!allTests) return [];
    const alreadyIn = new Set([...existingTestIds, ...addedIds]);
    let results = allTests.filter((t) => !alreadyIn.has(t.id));

    if (searchText.trim()) {
      const q = searchText.toLowerCase();
      results = results.filter(
        (t) =>
          t.name.toLowerCase().includes(q) ||
          (t.technique_mitre_id && t.technique_mitre_id.toLowerCase().includes(q)) ||
          (t.technique_name && t.technique_name.toLowerCase().includes(q))
      );
    }

    return results;
  }, [allTests, searchText, existingTestIds, addedIds]);

  const addMutation = useMutation({
    mutationFn: (testId: string) =>
      addTestToCampaign(campaignId, { test_id: testId }),
    onSuccess: (_data, testId) => {
      setAddedIds((prev) => new Set(prev).add(testId));
      queryClient.invalidateQueries({ queryKey: ["campaign", campaignId] });
      onSuccess();
    },
  });

  if (!open) return null;

  return (
    <div className="fixed inset-0 z-50 flex items-center justify-center bg-black/60 backdrop-blur-sm">
      <div className="w-full max-w-2xl rounded-xl border border-gray-700 bg-gray-900 shadow-2xl">
        {/* Header */}
        <div className="flex items-center justify-between border-b border-gray-800 px-6 py-4">
          <h2 className="text-lg font-semibold text-white">
            Add Tests to Campaign
          </h2>
          <button
            onClick={onClose}
            className="rounded-lg p-1.5 text-gray-400 hover:bg-gray-800 hover:text-white transition-colors"
          >
            <X className="h-5 w-5" />
          </button>
        </div>

        {/* Search */}
        <div className="border-b border-gray-800 px-6 py-3">
          <div className="relative">
            <Search className="absolute left-3 top-1/2 h-4 w-4 -translate-y-1/2 text-gray-500" />
            <input
              type="text"
              value={searchText}
              onChange={(e) => setSearchText(e.target.value)}
              placeholder="Search tests by name or technique..."
              autoFocus
              className="w-full rounded-lg border border-gray-700 bg-gray-800 pl-9 pr-3 py-2.5 text-sm text-gray-300 placeholder-gray-500 focus:border-cyan-500 focus:outline-none"
            />
          </div>
        </div>

        {/* Test list */}
        <div className="max-h-[400px] overflow-y-auto px-6 py-3">
          {isLoading ? (
            <div className="flex items-center justify-center py-12">
              <Loader2 className="h-6 w-6 animate-spin text-cyan-400" />
            </div>
          ) : filteredTests.length === 0 ? (
            <div className="flex flex-col items-center justify-center py-12 text-gray-400">
              <FlaskConical className="mb-2 h-8 w-8 text-gray-600" />
              <p className="text-sm">
                {searchText
                  ? "No tests match your search."
                  : "All available tests are already in this campaign."}
              </p>
            </div>
          ) : (
            <div className="space-y-1">
              {filteredTests.map((test: Test) => (
                <div
                  key={test.id}
                  className="flex items-center justify-between rounded-lg border border-gray-800 px-4 py-3 hover:border-gray-700 hover:bg-gray-800/50 transition-colors"
                >
                  <div className="flex-1 min-w-0">
                    <div className="flex items-center gap-2">
                      <span className="text-sm font-medium text-gray-200 truncate">
                        {test.name}
                      </span>
                      <span
                        className={`inline-flex shrink-0 rounded-full border px-2 py-0.5 text-[10px] font-medium ${
                          stateBadge[test.state]
                        }`}
                      >
                        {test.state.replace(/_/g, " ")}
                      </span>
                    </div>
                    <div className="mt-0.5 flex items-center gap-3 text-xs text-gray-500">
                      {test.technique_mitre_id && (
                        <span className="font-mono text-cyan-400/70">
                          {test.technique_mitre_id}
                        </span>
                      )}
                      {test.technique_name && (
                        <span className="truncate">{test.technique_name}</span>
                      )}
                      {test.platform && (
                        <span className="capitalize">{test.platform}</span>
                      )}
                    </div>
                  </div>
                  <button
                    onClick={() => addMutation.mutate(test.id)}
                    disabled={addMutation.isPending && addMutation.variables === test.id}
                    className="ml-3 flex shrink-0 items-center gap-1.5 rounded-lg border border-cyan-500/30 bg-cyan-500/10 px-3 py-1.5 text-xs font-medium text-cyan-400 hover:bg-cyan-500/20 disabled:opacity-50 transition-colors"
                  >
                    {addMutation.isPending && addMutation.variables === test.id ? (
                      <Loader2 className="h-3.5 w-3.5 animate-spin" />
                    ) : (
                      <Plus className="h-3.5 w-3.5" />
                    )}
                    Add
                  </button>
                </div>
              ))}
            </div>
          )}
        </div>

        {/* Footer */}
        <div className="flex items-center justify-between border-t border-gray-800 px-6 py-4">
          <div className="text-xs text-gray-500">
            {addedIds.size > 0 && (
              <span className="flex items-center gap-1 text-green-400">
                <CheckCircle className="h-3.5 w-3.5" />
                {addedIds.size} test{addedIds.size !== 1 ? "s" : ""} added
              </span>
            )}
          </div>
          <button
            onClick={onClose}
            className="rounded-lg border border-gray-700 bg-gray-800 px-4 py-2 text-sm font-medium text-gray-300 hover:bg-gray-700 transition-colors"
          >
            Done
          </button>
        </div>
      </div>
    </div>
  );
}