Aegis/frontend/src/pages/CompliancePage.tsx

import { useState } from "react";
import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import { Loader2, AlertCircle, Download, FileText, Plus } from "lucide-react";
import {
  getComplianceFrameworks,
  getFrameworkStatus,
  downloadComplianceCSV,
  importNistMappings,
  importCisMappings,
  type ComplianceFrameworkSummary,
} from "../api/compliance";
import { useAuth } from "../context/AuthContext";
import ComplianceGauge from "../components/compliance/ComplianceGauge";
import ControlsTable from "../components/compliance/ControlsTable";

export default function CompliancePage() {
  const [selectedFrameworkId, setSelectedFrameworkId] = useState<string | null>(null);
  const queryClient = useQueryClient();
  const { user } = useAuth();
  const isAdmin = user?.role === "admin";

  // Fetch available frameworks
  const {
    data: frameworks,
    isLoading: loadingFrameworks,
  } = useQuery({
    queryKey: ["compliance-frameworks"],
    queryFn: getComplianceFrameworks,
  });

  // Auto-select first framework
  const activeFrameworkId = selectedFrameworkId || frameworks?.[0]?.id || null;

  // Fetch framework status
  const {
    data: frameworkStatus,
    isLoading: loadingStatus,
  } = useQuery({
    queryKey: ["compliance-status", activeFrameworkId],
    queryFn: () => getFrameworkStatus(activeFrameworkId!),
    enabled: !!activeFrameworkId,
  });

  const isLoading = loadingFrameworks || loadingStatus;
  const summary = frameworkStatus?.summary;
  const controls = frameworkStatus?.controls || [];

  const handleExportCSV = async () => {
    if (activeFrameworkId) {
      await downloadComplianceCSV(activeFrameworkId);
    }
  };

  const handleExportJSON = async () => {
    if (!frameworkStatus) return;
    const json = JSON.stringify(frameworkStatus, null, 2);
    const blob = new Blob([json], { type: "application/json" });
    const url = URL.createObjectURL(blob);
    const a = document.createElement("a");
    a.href = url;
    a.download = `compliance_${frameworkStatus.framework.name.replace(/\s+/g, "_")}.json`;
    document.body.appendChild(a);
    a.click();
    document.body.removeChild(a);
    URL.revokeObjectURL(url);
  };

  const importNist = useMutation({
    mutationFn: importNistMappings,
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["compliance-frameworks"] });
      queryClient.invalidateQueries({ queryKey: ["compliance-status"] });
    },
  });

  const importCis = useMutation({
    mutationFn: importCisMappings,
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["compliance-frameworks"] });
      queryClient.invalidateQueries({ queryKey: ["compliance-status"] });
    },
  });

  const isImporting = importNist.isPending || importCis.isPending;

  if (isLoading && !frameworkStatus) {
    return (
      <div className="flex h-64 items-center justify-center">
        <Loader2 className="h-8 w-8 animate-spin text-cyan-400" />
      </div>
    );
  }

  return (
    <div className="space-y-6">
      {/* Header */}
      <div className="flex items-center justify-between">
        <div>
          <h1 className="text-2xl font-bold text-white">Compliance</h1>
          <p className="mt-1 text-sm text-gray-400">
            Map ATT&CK coverage to compliance framework controls
          </p>
        </div>

        <div className="flex items-center gap-2">
          {/* Framework selector */}
          <select
            value={activeFrameworkId || ""}
            onChange={(e) => setSelectedFrameworkId(e.target.value)}
            className="rounded-lg border border-gray-700 bg-gray-800 px-3 py-2 text-sm text-gray-200 focus:border-cyan-500 focus:outline-none"
          >
            {(frameworks || []).map((fw) => (
              <option key={fw.id} value={fw.id}>
                {fw.name}
              </option>
            ))}
          </select>

          {/* Export buttons */}
          <button
            onClick={handleExportCSV}
            className="flex items-center gap-1.5 rounded-lg border border-gray-700 bg-gray-800 px-3 py-2 text-xs font-medium text-gray-300 hover:border-cyan-500/50 hover:text-white"
          >
            <Download className="h-3.5 w-3.5" />
            Export CSV
          </button>
          <button
            onClick={handleExportJSON}
            className="flex items-center gap-1.5 rounded-lg border border-gray-700 bg-gray-800 px-3 py-2 text-xs font-medium text-gray-300 hover:border-cyan-500/50 hover:text-white"
          >
            <FileText className="h-3.5 w-3.5" />
            Export JSON
          </button>
        </div>
      </div>

      {/* Summary cards */}
      {summary && (
        <div className="grid grid-cols-2 gap-4 lg:grid-cols-5">
          {/* Gauge */}
          <div className="rounded-xl border border-gray-800 bg-gray-900 p-4 flex flex-col items-center justify-center">
            <ComplianceGauge percentage={summary.compliance_percentage} size="md" />
            <p className="mt-2 text-xs text-gray-500">Overall Compliance</p>
          </div>

          {/* Covered */}
          <div className="rounded-xl border border-gray-800 bg-gray-900 p-4">
            <p className="text-xs text-gray-500 uppercase tracking-wider">Covered</p>
            <p className="mt-2 text-3xl font-bold text-green-400">{summary.covered}</p>
            <div className="mt-2 h-1.5 w-full rounded-full bg-gray-800">
              <div
                className="h-full rounded-full bg-green-500"
                style={{ width: `${summary.total_controls > 0 ? (summary.covered / summary.total_controls) * 100 : 0}%` }}
              />
            </div>
          </div>

          {/* Partial */}
          <div className="rounded-xl border border-gray-800 bg-gray-900 p-4">
            <p className="text-xs text-gray-500 uppercase tracking-wider">Partial</p>
            <p className="mt-2 text-3xl font-bold text-yellow-400">
              {summary.partially_covered}
            </p>
            <div className="mt-2 h-1.5 w-full rounded-full bg-gray-800">
              <div
                className="h-full rounded-full bg-yellow-500"
                style={{ width: `${summary.total_controls > 0 ? (summary.partially_covered / summary.total_controls) * 100 : 0}%` }}
              />
            </div>
          </div>

          {/* Not Covered */}
          <div className="rounded-xl border border-gray-800 bg-gray-900 p-4">
            <p className="text-xs text-gray-500 uppercase tracking-wider">Not Covered</p>
            <p className="mt-2 text-3xl font-bold text-red-400">{summary.not_covered}</p>
            <div className="mt-2 h-1.5 w-full rounded-full bg-gray-800">
              <div
                className="h-full rounded-full bg-red-500"
                style={{ width: `${summary.total_controls > 0 ? (summary.not_covered / summary.total_controls) * 100 : 0}%` }}
              />
            </div>
          </div>

          {/* Not Evaluated */}
          <div className="rounded-xl border border-gray-800 bg-gray-900 p-4">
            <p className="text-xs text-gray-500 uppercase tracking-wider">Not Evaluated</p>
            <p className="mt-2 text-3xl font-bold text-gray-400">{summary.not_evaluated}</p>
            <div className="mt-2 h-1.5 w-full rounded-full bg-gray-800">
              <div
                className="h-full rounded-full bg-gray-500"
                style={{ width: `${summary.total_controls > 0 ? (summary.not_evaluated / summary.total_controls) * 100 : 0}%` }}
              />
            </div>
          </div>
        </div>
      )}

      {/* Import buttons for admin */}
      {isAdmin && (
        <div className="flex flex-wrap items-center gap-2">
          <span className="text-xs text-gray-500">Import frameworks:</span>
          <button
            onClick={() => importNist.mutate()}
            disabled={isImporting}
            className="flex items-center gap-1.5 rounded-lg border border-gray-700 bg-gray-800 px-3 py-1.5 text-xs font-medium text-gray-300 hover:border-cyan-500/50 hover:text-white transition-colors disabled:opacity-50"
          >
            {importNist.isPending ? <Loader2 className="h-3 w-3 animate-spin" /> : <Plus className="h-3 w-3" />}
            NIST 800-53
          </button>
          <button
            onClick={() => importCis.mutate()}
            disabled={isImporting}
            className="flex items-center gap-1.5 rounded-lg border border-gray-700 bg-gray-800 px-3 py-1.5 text-xs font-medium text-gray-300 hover:border-cyan-500/50 hover:text-white transition-colors disabled:opacity-50"
          >
            {importCis.isPending ? <Loader2 className="h-3 w-3 animate-spin" /> : <Plus className="h-3 w-3" />}
            CIS Controls v8
          </button>
          {(importNist.isSuccess || importCis.isSuccess) && (
            <span className="text-xs text-green-400">Import complete</span>
          )}
          {(importNist.isError || importCis.isError) && (
            <span className="text-xs text-red-400">Import failed</span>
          )}
        </div>
      )}

      {/* Controls table */}
      {controls.length > 0 ? (
        <ControlsTable controls={controls} />
      ) : (
        !isLoading && (
          <div className="rounded-xl border border-gray-800 bg-gray-900 p-8 text-center">
            <AlertCircle className="mx-auto h-10 w-10 text-gray-600" />
            <p className="mt-3 text-gray-400">
              No compliance data available. Use the import buttons above to load a compliance framework.
            </p>
          </div>
        )
      )}
    </div>
  );
}