Aegis/frontend/src/pages/TestCreatePage.tsx

import { useNavigate, useSearchParams } from "react-router-dom";
import { useMutation, useQueryClient, useQuery } from "@tanstack/react-query";
import { ArrowLeft, FlaskConical } from "lucide-react";
import TestForm, { type TestFormData } from "../components/TestForm";
import { createTest } from "../api/tests";
import { getTechniqueByMitreId } from "../api/techniques";

export default function TestCreatePage() {
  const navigate = useNavigate();
  const queryClient = useQueryClient();
  const [searchParams] = useSearchParams();
  
  // Get technique ID from URL query param (UUID format)
  const techniqueId = searchParams.get("technique");

  // If we have a technique ID, try to get its mitre_id for the back link
  const { data: technique } = useQuery({
    queryKey: ["techniqueById", techniqueId],
    queryFn: async () => {
      // We need to find the mitre_id from the technique list
      // This is a workaround since we get UUID but need mitre_id
      const apiBase = import.meta.env.VITE_API_URL || "/api/v1";
      const response = await fetch(`${apiBase}/techniques`, {
        headers: {
          Authorization: `Bearer ${localStorage.getItem("token")}`,
        },
      });
      const techniques = await response.json();
      return techniques.find((t: { id: string }) => t.id === techniqueId);
    },
    enabled: !!techniqueId,
  });

  const createMutation = useMutation({
    mutationFn: createTest,
    onSuccess: (data) => {
      queryClient.invalidateQueries({ queryKey: ["techniques"] });
      queryClient.invalidateQueries({ queryKey: ["technique"] });
      // Navigate back to the technique detail if we came from there
      if (technique?.mitre_id) {
        navigate(`/techniques/${technique.mitre_id}`);
      } else {
        navigate("/tests");
      }
    },
  });

  const handleSubmit = (data: TestFormData) => {
    createMutation.mutate({
      technique_id: data.technique_id,
      name: data.name,
      description: data.description || undefined,
      platform: data.platform || undefined,
      procedure_text: data.procedure_text || undefined,
      tool_used: data.tool_used || undefined,
    });
  };

  const handleBack = () => {
    if (technique?.mitre_id) {
      navigate(`/techniques/${technique.mitre_id}`);
    } else {
      navigate("/tests");
    }
  };

  return (
    <div className="space-y-6">
      {/* Back button */}
      <button
        onClick={handleBack}
        className="flex items-center gap-1 text-sm text-gray-400 hover:text-cyan-400 transition-colors"
      >
        <ArrowLeft className="h-4 w-4" />
        {technique ? `Back to ${technique.mitre_id}` : "Back to tests"}
      </button>

      {/* Header */}
      <div className="flex items-center gap-4">
        <div className="rounded-lg bg-cyan-500/10 p-3">
          <FlaskConical className="h-8 w-8 text-cyan-400" />
        </div>
        <div>
          <h1 className="text-2xl font-bold text-white">Create New Test</h1>
          <p className="mt-1 text-sm text-gray-400">
            Document a security test for technique validation
          </p>
        </div>
      </div>

      {/* Error display */}
      {createMutation.isError && (
        <div className="rounded-lg border border-red-500/30 bg-red-900/20 p-4">
          <p className="text-sm text-red-400">
            Failed to create test: {(createMutation.error as Error)?.message || "Unknown error"}
          </p>
        </div>
      )}

      {/* Form */}
      <div className="rounded-xl border border-gray-800 bg-gray-900 p-6">
        <TestForm
          preselectedTechniqueId={techniqueId || undefined}
          onSubmit={handleSubmit}
          isSubmitting={createMutation.isPending}
        />
      </div>
    </div>
  );
}