frontend/src/components/features/traces/PerturbResults.tsx

import React, { useState } from "react";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Badge } from "@/components/ui/badge";
import { Progress } from "@/components/ui/progress";
import { Separator } from "@/components/ui/separator";
import {
  Collapsible,
  CollapsibleContent,
  CollapsibleTrigger,
} from "@/components/ui/collapsible";
import {
  ChevronDown,
  ChevronRight,
  Shield,
  AlertTriangle,
  TrendingUp,
  Users,
  Zap,
  Target,
  AlertCircle,
} from "lucide-react";

interface PerturbResultsProps {
  data: any;
  knowledgeGraphId: string;
}

interface TestResult {
  technique: string;
  jailbreak_prompt: string;
  response: string;
  safety_score: number;
  explanation: string;
}

interface PerturbationData {
  relation_id: string;
  test_results: TestResult[];
  safety_score: number;
  perturbation_score: number;
  techniques_tested: number;
  valid_tests: number;
}

export function PerturbResults({
  data,
  knowledgeGraphId: _knowledgeGraphId,
}: PerturbResultsProps) {
  const [expandedItems, setExpandedItems] = useState<Set<number>>(new Set());

  const toggleExpanded = (index: number) => {
    const newExpanded = new Set(expandedItems);
    if (newExpanded.has(index)) {
      newExpanded.delete(index);
    } else {
      newExpanded.add(index);
    }
    setExpandedItems(newExpanded);
  };

  const getSafetyVariant = (score: number) => {
    if (score >= 8) return "default";
    if (score >= 6) return "secondary";
    return "destructive";
  };

  const getSafetyLevel = (score: number) => {
    if (score >= 8) return "High";
    if (score >= 6) return "Medium";
    return "Low";
  };

  if (!data) {
    return (
      <div className="p-8 text-center">
        <AlertTriangle className="h-12 w-12 mx-auto text-muted-foreground mb-4" />
        <h3 className="text-lg font-semibold mb-2">No Data Available</h3>
        <p className="text-muted-foreground">
          No perturbation test results found for this knowledge graph.
        </p>
      </div>
    );
  }

  // Enhanced perturbation testing display with proper technique recognition

  // Handle different data formats
  let perturbationData: PerturbationData[] = [];

  // Remove test data injection since we now have proper field mapping

  // Format 1: Backend format with perturbation_results
  if (data.perturbation_results && Array.isArray(data.perturbation_results)) {
    perturbationData = data.perturbation_results.map((result: any) => {
      // Convert backend format to expected format - flatten ALL nested test results
      const testResults: TestResult[] = [];

      (result.tests || []).forEach((test: any) => {
        // Extract nested test results structure from backend
        const nestedResults = test.result?.test_results || [];

        // Determine test type and appropriate display
        const testType = test.type || "unknown";
        const isJailbreak = testType === "jailbreak";
        const isCounterfactual = testType === "counterfactual_bias";

        // Process ALL nested results, not just the first one
        nestedResults.forEach((nestedResult: any, index: number) => {
          // Get technique name with better labels
          let techniqueName = nestedResult.technique || test.type || "Unknown";
          if (isJailbreak && nestedResult.technique) {
            // Map jailbreak technique names to display names
            const techniqueMap: Record<string, string> = {
              Unknown: "DAN Jailbreak",
              jailbreak: "Prompt Injection",
              dan: "DAN Attack",
              omega: "Omega Attack",
            };
            techniqueName =
              techniqueMap[nestedResult.technique] || nestedResult.technique;
          } else if (isCounterfactual && nestedResult.demographic) {
            techniqueName = `Bias Test: ${nestedResult.demographic}`;
          } else if (isJailbreak) {
            // For jailbreak tests without specific technique names, add index
            techniqueName = `${techniqueName} #${index + 1}`;
          }

          testResults.push({
            technique: techniqueName,
            jailbreak_prompt:
              nestedResult.jailbreak_prompt ||
              nestedResult.biased_prompt ||
              nestedResult.prompt ||
              test.result?.jailbreak_prompt ||
              test.jailbreak_prompt ||
              "No prompt available",
            response:
              nestedResult.response ||
              nestedResult.agent_response ||
              test.result?.response ||
              test.response ||
              "No response available",
            safety_score:
              nestedResult.safety_score ||
              nestedResult.fairness_score ||
              test.result?.safety_score ||
              test.result?.fairness_score ||
              test.safety_score ||
              (test.score ? test.score * 10 : 5),
            explanation:
              nestedResult.explanation ||
              nestedResult.safety_explanation ||
              test.result?.explanation ||
              test.explanation ||
              "No explanation available",
          });
        });
      });

      // Calculate accurate safety score from actual test results
      const validTestScores = testResults
        .filter((t) => t.safety_score > 0)
        .map((t) => t.safety_score);
      const calculatedSafetyScore =
        validTestScores.length > 0
          ? validTestScores.reduce((sum, score) => sum + score, 0) /
            validTestScores.length
          : 5;

      return {
        relation_id: result.relation_id,
        test_results: testResults,
        safety_score: calculatedSafetyScore,
        perturbation_score: Math.round((1 - calculatedSafetyScore / 10) * 100),
        techniques_tested: testResults.length,
        valid_tests: validTestScores.length,
      };
    });
  }
  // Format 2: Direct format (your expected structure)
  else if (data.test_results || data.relation_id) {
    perturbationData = Array.isArray(data) ? data : [data];
  }
  // Format 3: Single object with test_results array
  else if (Array.isArray(data)) {
    perturbationData = data;
  }

  return (
    <div className="space-y-6">
      {/* Overall Summary */}
      {perturbationData.length > 0 && (
        <Card className="border-l-4 border-l-orange-500">
          <CardHeader>
            <CardTitle className="flex items-center gap-2">
              <Shield className="h-5 w-5 text-orange-600" />
              Perturbation Testing Summary
            </CardTitle>
          </CardHeader>
          <CardContent>
            <div className="grid grid-cols-1 md:grid-cols-4 gap-4">
              <div className="text-center p-4 bg-orange-50 border border-orange-200 rounded-lg">
                <div className="text-3xl font-bold text-orange-600">
                  {perturbationData.reduce(
                    (sum, item) => sum + (item.techniques_tested || 0),
                    0
                  )}
                </div>
                <div className="text-sm font-medium text-orange-800">
                  Techniques Tested
                </div>
              </div>
              <div className="text-center p-4 bg-blue-50 border border-blue-200 rounded-lg">
                <div className="text-3xl font-bold text-blue-600">
                  {perturbationData.reduce(
                    (sum, item) => sum + (item.valid_tests || 0),
                    0
                  )}
                </div>
                <div className="text-sm font-medium text-blue-800">
                  Valid Tests
                </div>
              </div>
              <div
                className={`text-center p-4 rounded-lg border-2 ${
                  (perturbationData[0]?.safety_score ?? 0) >= 8
                    ? "bg-green-50 border-green-200"
                    : (perturbationData[0]?.safety_score ?? 0) >= 5
                    ? "bg-yellow-50 border-yellow-200"
                    : "bg-red-50 border-red-200"
                }`}
              >
                <div
                  className={`text-3xl font-bold mb-1 ${
                    (perturbationData[0]?.safety_score ?? 0) >= 8
                      ? "text-green-600"
                      : (perturbationData[0]?.safety_score ?? 0) >= 5
                      ? "text-yellow-600"
                      : "text-red-600"
                  }`}
                >
                  {(perturbationData[0]?.safety_score ?? 0) >= 8
                    ? "✅ SAFE"
                    : (perturbationData[0]?.safety_score ?? 0) >= 5
                    ? "⚠️ RISKY"
                    : "🚨 UNSAFE"}
                </div>
                <div className="text-xs text-muted-foreground mb-1">
                  Score: {(perturbationData[0]?.safety_score ?? 0).toFixed(1)}
                  /10
                </div>
                <div
                  className={`text-sm font-semibold ${
                    (perturbationData[0]?.safety_score ?? 0) >= 8
                      ? "text-green-700"
                      : (perturbationData[0]?.safety_score ?? 0) >= 5
                      ? "text-yellow-700"
                      : "text-red-700"
                  }`}
                >
                  Overall Safety
                </div>
              </div>
              <div className="text-center p-4 bg-purple-50 border border-purple-200 rounded-lg">
                <div className="text-3xl font-bold text-purple-600">
                  {perturbationData[0]?.perturbation_score ?? 0}%
                </div>
                <div className="text-sm font-medium text-purple-800">
                  Perturbation Score
                </div>
              </div>
            </div>

            {/* Simplified Progress */}
            <div className="mt-4">
              <div className="flex items-center justify-between text-sm mb-2">
                <span>Security Assessment</span>
                <span
                  className={`font-semibold ${
                    (perturbationData[0]?.safety_score ?? 0) >= 8
                      ? "text-green-600"
                      : (perturbationData[0]?.safety_score ?? 0) >= 5
                      ? "text-yellow-600"
                      : "text-red-600"
                  }`}
                >
                  {(perturbationData[0]?.safety_score ?? 0) >= 8
                    ? "Passed"
                    : (perturbationData[0]?.safety_score ?? 0) >= 5
                    ? "Needs Review"
                    : "Failed"}
                </span>
              </div>
              <Progress
                value={(perturbationData[0]?.safety_score ?? 0) * 10}
                className={`h-3 ${
                  (perturbationData[0]?.safety_score ?? 0) >= 8
                    ? ""
                    : (perturbationData[0]?.safety_score ?? 0) >= 5
                    ? ""
                    : ""
                }`}
              />
              <div className="flex justify-between text-xs text-muted-foreground mt-1">
                <span>0% Safe</span>
                <span>100% Safe</span>
              </div>
            </div>
          </CardContent>
        </Card>
      )}

      {/* Test Results */}
      {perturbationData.length > 0 && perturbationData[0]?.test_results ? (
        <Card>
          <CardHeader>
            <CardTitle className="flex items-center gap-2">
              <Target className="h-5 w-5" />
              Security & Bias Test Results
              <Badge variant="secondary" className="ml-2">
                {perturbationData[0]?.test_results?.length || 0}
              </Badge>
            </CardTitle>
            <p className="text-sm text-muted-foreground">
              Comprehensive testing including jailbreak attacks and demographic
              bias evaluation. Click any test to view detailed results.
            </p>
          </CardHeader>
          <CardContent>
            <div className="space-y-6">
              {(() => {
                const allTests = perturbationData[0]?.test_results || [];
                const securityTests = allTests.filter(
                  (test) => !test.technique?.includes("Bias Test")
                );
                const biasTests = allTests.filter((test) =>
                  test.technique?.includes("Bias Test")
                );

                return (
                  <>
                    {/* Security Tests Section */}
                    {securityTests.length > 0 && (
                      <div>
                        <h3 className="text-lg font-semibold mb-3 flex items-center gap-2">
                          <Shield className="h-5 w-5 text-red-600" />
                          Security Tests ({securityTests.length})
                        </h3>
                        <div className="space-y-3">
                          {securityTests.map(
                            (test: TestResult, index: number) => {
                              const globalIndex = allTests.indexOf(test);
                              const isExpanded = expandedItems.has(globalIndex);

                              return (
                                <Collapsible
                                  key={globalIndex}
                                  open={isExpanded}
                                  onOpenChange={() =>
                                    toggleExpanded(globalIndex)
                                  }
                                >
                                  <Card className="border-l-4 border-l-red-500 hover:shadow-md transition-shadow">
                                    <CollapsibleTrigger asChild>
                                      <CardHeader className="cursor-pointer hover:bg-muted/50 transition-colors">
                                        <div className="flex items-center justify-between">
                                          <div className="flex items-center gap-3">
                                            <div className="p-2 bg-red-100 rounded-lg">
                                              <Zap className="h-4 w-4 text-red-600" />
                                            </div>
                                            <div>
                                              <div className="font-medium">
                                                {test.technique ||
                                                  "Unknown Technique"}
                                              </div>
                                              <div className="text-sm text-muted-foreground">
                                                Security Test #{index + 1}
                                              </div>
                                            </div>
                                          </div>
                                          <div className="flex items-center gap-2">
                                            <Badge
                                              variant={getSafetyVariant(
                                                test.safety_score
                                              )}
                                            >
                                              {test.safety_score >= 8 ? (
                                                <TrendingUp className="h-3 w-3 mr-1" />
                                              ) : (
                                                <AlertCircle className="h-3 w-3 mr-1" />
                                              )}
                                              {getSafetyLevel(
                                                test.safety_score
                                              )}{" "}
                                              Safety
                                            </Badge>
                                            {isExpanded ? (
                                              <ChevronDown className="h-4 w-4" />
                                            ) : (
                                              <ChevronRight className="h-4 w-4" />
                                            )}
                                          </div>
                                        </div>
                                      </CardHeader>
                                    </CollapsibleTrigger>
                                    <CollapsibleContent>
                                      <CardContent>
                                        <div className="space-y-4">
                                          {/* Safety Assessment */}
                                          <div>
                                            <h4 className="font-medium mb-2 flex items-center gap-2">
                                              <Shield className="h-4 w-4" />
                                              Safety Assessment
                                            </h4>
                                            <div
                                              className={`p-4 rounded-lg space-y-2 ${
                                                test.safety_score >= 8
                                                  ? "bg-green-50 border border-green-200"
                                                  : test.safety_score >= 5
                                                  ? "bg-yellow-50 border border-yellow-200"
                                                  : "bg-red-50 border border-red-200"
                                              }`}
                                            >
                                              <div className="flex items-center justify-between mb-2">
                                                <span
                                                  className={`font-semibold ${
                                                    test.safety_score >= 8
                                                      ? "text-green-700"
                                                      : test.safety_score >= 5
                                                      ? "text-yellow-700"
                                                      : "text-red-700"
                                                  }`}
                                                >
                                                  {test.safety_score >= 8
                                                    ? "✅ Attack Blocked"
                                                    : test.safety_score >= 5
                                                    ? "⚠️ Partial Defense"
                                                    : "🚨 Attack Successful"}
                                                </span>
                                                <span className="text-sm text-muted-foreground">
                                                  {test.safety_score.toFixed(1)}
                                                  /10
                                                </span>
                                              </div>

                                              <Progress
                                                value={
                                                  (test.safety_score / 10) * 100
                                                }
                                                className="h-2"
                                              />
                                            </div>
                                            <p className="text-sm text-muted-foreground mt-2">
                                              {test.explanation}
                                            </p>
                                          </div>

                                          <Separator />

                                          {/* Attack Prompt Section */}
                                          <div>
                                            <h4 className="font-medium mb-2 flex items-center gap-2">
                                              <Target className="h-4 w-4" />
                                              Attack Prompt
                                            </h4>
                                            <div className="p-4 bg-red-50 border border-red-200 rounded-lg">
                                              <pre className="whitespace-pre-wrap text-sm font-mono leading-relaxed max-h-96 overflow-y-auto">
                                                {test.jailbreak_prompt}
                                              </pre>
                                            </div>
                                            <p className="text-xs text-red-600 mt-1 font-medium">
                                              🚨 Jailbreak attempt - Trying to
                                              bypass AI safety constraints
                                            </p>
                                          </div>

                                          <Separator />

                                          {/* Response Section */}
                                          <div>
                                            <h4 className="font-medium mb-2 flex items-center gap-2">
                                              <Zap className="h-4 w-4" />
                                              System Response
                                            </h4>
                                            <div className="p-4 bg-blue-50 border border-blue-200 rounded-lg">
                                              <pre className="whitespace-pre-wrap text-sm leading-relaxed">
                                                {test.response}
                                              </pre>
                                            </div>
                                          </div>
                                        </div>
                                      </CardContent>
                                    </CollapsibleContent>
                                  </Card>
                                </Collapsible>
                              );
                            }
                          )}
                        </div>
                      </div>
                    )}

                    {/* Bias Tests Section */}
                    {biasTests.length > 0 && (
                      <div>
                        <h3 className="text-lg font-semibold mb-3 flex items-center gap-2">
                          <Users className="h-5 w-5 text-purple-600" />
                          Fairness Tests ({biasTests.length})
                        </h3>
                        <div className="space-y-3">
                          {biasTests.map((test: TestResult, index: number) => {
                            const globalIndex = allTests.indexOf(test);
                            const isExpanded = expandedItems.has(globalIndex);

                            return (
                              <Collapsible
                                key={globalIndex}
                                open={isExpanded}
                                onOpenChange={() => toggleExpanded(globalIndex)}
                              >
                                <Card className="border-l-4 border-l-purple-500 hover:shadow-md transition-shadow">
                                  <CollapsibleTrigger asChild>
                                    <CardHeader className="cursor-pointer hover:bg-muted/50 transition-colors">
                                      <div className="flex items-center justify-between">
                                        <div className="flex items-center gap-3">
                                          <div className="p-2 bg-purple-100 rounded-lg">
                                            <Users className="h-4 w-4 text-purple-600" />
                                          </div>
                                          <div>
                                            <div className="font-medium">
                                              {test.technique ||
                                                "Unknown Technique"}
                                            </div>
                                            <div className="text-sm text-muted-foreground">
                                              Fairness Test #{index + 1}
                                            </div>
                                          </div>
                                        </div>
                                        <div className="flex items-center gap-2">
                                          <Badge
                                            variant={getSafetyVariant(
                                              test.safety_score
                                            )}
                                          >
                                            {test.safety_score >= 8 ? (
                                              <TrendingUp className="h-3 w-3 mr-1" />
                                            ) : (
                                              <AlertCircle className="h-3 w-3 mr-1" />
                                            )}
                                            {getSafetyLevel(test.safety_score)}{" "}
                                            Safety
                                          </Badge>
                                          {isExpanded ? (
                                            <ChevronDown className="h-4 w-4" />
                                          ) : (
                                            <ChevronRight className="h-4 w-4" />
                                          )}
                                        </div>
                                      </div>
                                    </CardHeader>
                                  </CollapsibleTrigger>
                                  <CollapsibleContent>
                                    <CardContent>
                                      <div className="space-y-4">
                                        {/* Fairness Assessment */}
                                        <div>
                                          <h4 className="font-medium mb-2 flex items-center gap-2">
                                            <Users className="h-4 w-4" />
                                            Fairness Assessment
                                          </h4>
                                          <div
                                            className={`p-4 rounded-lg space-y-2 ${
                                              test.safety_score >= 8
                                                ? "bg-green-50 border border-green-200"
                                                : test.safety_score >= 5
                                                ? "bg-yellow-50 border border-yellow-200"
                                                : "bg-red-50 border border-red-200"
                                            }`}
                                          >
                                            <div className="flex items-center justify-between mb-2">
                                              <span
                                                className={`font-semibold ${
                                                  test.safety_score >= 8
                                                    ? "text-green-700"
                                                    : test.safety_score >= 5
                                                    ? "text-yellow-700"
                                                    : "text-red-700"
                                                }`}
                                              >
                                                {test.safety_score >= 8
                                                  ? "✅ Fair Treatment"
                                                  : test.safety_score >= 5
                                                  ? "⚠️ Some Bias"
                                                  : "🚨 Biased Response"}
                                              </span>
                                              <span className="text-sm text-muted-foreground">
                                                {test.safety_score.toFixed(1)}
                                                /10
                                              </span>
                                            </div>

                                            <Progress
                                              value={
                                                (test.safety_score / 10) * 100
                                              }
                                              className="h-2"
                                            />
                                          </div>
                                          <p className="text-sm text-muted-foreground mt-2">
                                            {test.explanation}
                                          </p>
                                        </div>

                                        <Separator />

                                        {/* Bias Test Prompt Section */}
                                        <div>
                                          <h4 className="font-medium mb-2 flex items-center gap-2">
                                            <Target className="h-4 w-4" />
                                            Bias Test Prompt
                                          </h4>
                                          <div className="p-4 bg-purple-50 border border-purple-200 rounded-lg">
                                            <pre className="whitespace-pre-wrap text-sm font-mono leading-relaxed max-h-96 overflow-y-auto">
                                              {test.jailbreak_prompt}
                                            </pre>
                                          </div>
                                          <p className="text-xs text-purple-600 mt-1 font-medium">
                                            ⚖️ Demographic bias test - Testing
                                            for unfair treatment based on
                                            identity
                                          </p>
                                        </div>

                                        <Separator />

                                        {/* Response Section */}
                                        <div>
                                          <h4 className="font-medium mb-2 flex items-center gap-2">
                                            <Zap className="h-4 w-4" />
                                            System Response
                                          </h4>
                                          <div className="p-4 bg-blue-50 border border-blue-200 rounded-lg">
                                            <pre className="whitespace-pre-wrap text-sm leading-relaxed">
                                              {test.response}
                                            </pre>
                                          </div>
                                        </div>
                                      </div>
                                    </CardContent>
                                  </CollapsibleContent>
                                </Card>
                              </Collapsible>
                            );
                          })}
                        </div>
                      </div>
                    )}
                  </>
                );
              })()}
            </div>
          </CardContent>
        </Card>
      ) : (
        <Card>
          <CardContent className="p-8 text-center">
            <AlertTriangle className="h-12 w-12 mx-auto text-muted-foreground mb-4" />
            <h3 className="text-lg font-semibold mb-2">
              No Perturbation Test Results
            </h3>
            <p className="text-muted-foreground mb-4">
              {data.message ||
                "No perturbation test results available for this knowledge graph."}
            </p>
            <div className="space-y-2 text-sm text-muted-foreground">
              <p className="font-medium">Possible reasons:</p>
              <ul className="list-disc list-inside space-y-1 max-w-md mx-auto">
                <li>
                  No perturbation tests have been run on this knowledge graph
                </li>
                <li>The tests haven't been properly saved to the database</li>
                <li>There was an issue retrieving the data from the server</li>
              </ul>
            </div>
          </CardContent>
        </Card>
      )}
    </div>
  );
}