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

import React, { useState, useCallback } from "react";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";

import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";

import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogHeader,
  DialogTitle,
} from "@/components/ui/dialog";
import {
  AlertCircle,
  CheckCircle,
  Clock,
  Play,
  Eye,
  Zap,
  Shield,
  GitBranch,
  Sparkles,
  Settings,
} from "lucide-react";
import { KnowledgeGraph, PerturbationConfig } from "@/types";
import { useNotification } from "@/context/NotificationContext";
import { api } from "@/lib/api";
import { EnrichResults } from "./EnrichResults";
import { PerturbResults } from "./PerturbResults";
import { PerturbationTestConfig } from "./PerturbationTestConfig";

interface OptionalPipelineSectionProps {
  knowledgeGraph: KnowledgeGraph;
  onStageComplete?: (stage: string) => void;
}

interface StageState {
  status: "idle" | "running" | "completed" | "error";
  progress: number;
  taskId?: string;
  error?: string;
}

interface PipelineStageConfig {
  id: string;
  name: string;
  description: string;
  icon: React.ComponentType<{ className?: string }>;
  apiEndpoint: (kgId: string) => string;
  benefit: string;
}

const PIPELINE_STAGES: PipelineStageConfig[] = [
  {
    id: "enrich",
    name: "Prompt Reconstruction",
    description: "Reconstruct original prompts from conversation patterns",
    icon: Sparkles,
    apiEndpoint: (kgId) => `/knowledge-graphs/${kgId}/enrich`,
    benefit: "Understand conversation context and intent",
  },
  {
    id: "perturb",
    name: "Perturbation Testing",
    description: "Test graph robustness through systematic variations",
    icon: Shield,
    apiEndpoint: (kgId) => `/knowledge-graphs/${kgId}/perturb`,
    benefit: "Validate relationship reliability and stability",
  },
  {
    id: "causal",
    name: "Causal Analysis",
    description: "Analyze causal relationships and dependencies",
    icon: GitBranch,
    apiEndpoint: (kgId) => `/knowledge-graphs/${kgId}/analyze`,
    benefit: "Discover cause-and-effect patterns",
  },
];

export function OptionalPipelineSection({
  knowledgeGraph,
  onStageComplete,
}: OptionalPipelineSectionProps) {
  const { showNotification } = useNotification();
  const [stageStates, setStageStates] = useState<Record<string, StageState>>({
    enrich: {
      status: knowledgeGraph.is_enriched ? "completed" : "idle",
      progress: knowledgeGraph.is_enriched ? 100 : 0,
    },
    perturb: {
      status: knowledgeGraph.is_perturbed ? "completed" : "idle",
      progress: knowledgeGraph.is_perturbed ? 100 : 0,
    },
    causal: {
      status: knowledgeGraph.is_analyzed ? "completed" : "idle",
      progress: knowledgeGraph.is_analyzed ? 100 : 0,
    },
  });

  const [selectedStageResults, setSelectedStageResults] = useState<{
    stage: string;
    data: any;
  } | null>(null);

  const [showPerturbConfig, setShowPerturbConfig] = useState(false);

  const getStageStatusColor = (status: StageState["status"]) => {
    switch (status) {
      case "completed":
        return "bg-green-500/10 text-green-700 border-green-200";
      case "running":
        return "bg-blue-500/10 text-blue-700 border-blue-200";
      case "error":
        return "bg-red-500/10 text-red-700 border-red-200";
      default:
        return "bg-gray-500/10 text-gray-700 border-gray-200";
    }
  };

  const pollTaskStatus = useCallback(
    async (stageId: string, taskId: string) => {
      const maxAttempts = 60; // 5 minutes at 5-second intervals
      let attempts = 0;

      const poll = async () => {
        try {
          const status = await api.tasks.get(taskId);

          setStageStates((prev) => {
            const currentState = prev[stageId] || {
              status: "idle" as const,
              progress: 0,
            };
            return {
              ...prev,
              [stageId]: {
                ...currentState,
                progress: status.progress || currentState.progress,
              },
            };
          });

          if (status.status === "completed" || status.status === "COMPLETED") {
            setStageStates((prev) => ({
              ...prev,
              [stageId]: { status: "completed", progress: 100 },
            }));

            onStageComplete?.(stageId);

            showNotification({
              type: "success",
              title: "Stage Completed",
              message: `${
                PIPELINE_STAGES.find((s) => s.id === stageId)?.name
              } has finished successfully.`,
            });

            return;
          } else if (status.status === "failed" || status.status === "FAILED") {
            setStageStates((prev) => ({
              ...prev,
              [stageId]: {
                status: "error",
                progress: 0,
                error: status.error || "Unknown error",
              },
            }));

            showNotification({
              type: "error",
              title: "Stage Failed",
              message: status.error || "Processing failed",
            });

            return;
          }

          // Continue polling
          attempts++;
          if (attempts < maxAttempts) {
            setTimeout(poll, 5000); // Poll every 5 seconds
          } else {
            // Timeout
            setStageStates((prev) => ({
              ...prev,
              [stageId]: {
                status: "error",
                progress: 0,
                error: "Processing timeout",
              },
            }));
          }
        } catch (error) {
          console.error("Error polling task status:", error);
          // Continue polling on error, might be temporary
          attempts++;
          if (attempts < maxAttempts) {
            setTimeout(poll, 5000);
          }
        }
      };

      // Start polling
      setTimeout(poll, 2000); // First check after 2 seconds
    },
    [onStageComplete, showNotification]
  );

  const runPerturbWithConfig = useCallback(
    async (config: PerturbationConfig) => {
      const stageConfig = PIPELINE_STAGES.find((s) => s.id === "perturb")!;

      setStageStates((prev) => ({
        ...prev,
        perturb: { status: "running", progress: 0 },
      }));

      try {
        const response = await api.knowledgeGraphs.perturb(
          knowledgeGraph.kg_id,
          config
        );

        if (response.task_id) {
          pollTaskStatus("perturb", response.task_id);
          showNotification({
            type: "info",
            title: `${stageConfig.name} Started`,
            message: "Processing has begun with custom configuration.",
          });
        }
      } catch (error) {
        console.error("Error running perturb stage:", error);
        setStageStates((prev) => ({
          ...prev,
          perturb: {
            status: "error",
            progress: 0,
            error: error instanceof Error ? error.message : "Unknown error",
          },
        }));
        showNotification({
          type: "error",
          title: `${stageConfig.name} Failed`,
          message: error instanceof Error ? error.message : "An error occurred",
        });
      }
    },
    [knowledgeGraph.kg_id, showNotification, pollTaskStatus]
  );

  const runStage = useCallback(
    async (stageConfig: PipelineStageConfig) => {
      const { id } = stageConfig;

      // For perturb stage, open the configuration dialog instead
      if (id === "perturb") {
        setShowPerturbConfig(true);
        return;
      }

      setStageStates((prev) => ({
        ...prev,
        [id]: { status: "running", progress: 0 },
      }));

      try {
        // Call the API endpoint based on stage
        let response;

        if (id === "enrich") {
          response = await api.knowledgeGraphs.enrich(knowledgeGraph.kg_id);
        } else if (id === "causal") {
          response = await api.knowledgeGraphs.analyze(knowledgeGraph.kg_id);
        } else {
          throw new Error(`Unknown stage: ${id}`);
        }

        if (response.task_id) {
          // Start polling for task status
          pollTaskStatus(id, response.task_id);

          showNotification({
            type: "info",
            title: `${stageConfig.name} Started`,
            message: "Processing has begun. This may take several minutes.",
          });
        } else {
          // Mark as completed immediately if no task ID
          setStageStates((prev) => ({
            ...prev,
            [id]: { status: "completed", progress: 100 },
          }));

          onStageComplete?.(id);
        }
      } catch (error) {
        console.error(`Error running ${id} stage:`, error);

        setStageStates((prev) => ({
          ...prev,
          [id]: {
            status: "error",
            progress: 0,
            error: error instanceof Error ? error.message : "Unknown error",
          },
        }));

        showNotification({
          type: "error",
          title: `${stageConfig.name} Failed`,
          message: error instanceof Error ? error.message : "An error occurred",
        });
      }
    },
    [knowledgeGraph.kg_id, showNotification, onStageComplete, pollTaskStatus]
  );

  const viewStageResults = useCallback(
    async (stageId: string) => {
      try {
        const results = await api.knowledgeGraphs.getStageResults(
          knowledgeGraph.kg_id,
          stageId
        );

        setSelectedStageResults({
          stage: stageId,
          data: results,
        });
      } catch (error) {
        showNotification({
          type: "error",
          title: "Failed to Load Results",
          message: "Could not retrieve stage results",
        });
      }
    },
    [knowledgeGraph.kg_id, showNotification]
  );

  const anyStageRunning = Object.values(stageStates).some(
    (state) => state.status === "running"
  );

  return (
    <>
      <div className="space-y-6">
        {/* Advanced Processing - Unified */}
        <Card>
          <CardHeader>
            <CardTitle className="flex items-center gap-2 text-xl">
              <Zap className="h-6 w-6" />
              Advanced Processing
            </CardTitle>
          </CardHeader>
          <CardContent>
            <div className="space-y-6">
              {/* Description */}
              <div className="text-muted-foreground">
                <p className="text-base">
                  Enhance your knowledge graph with advanced analysis
                  techniques. These optional stages provide deeper insights but
                  are not required for basic graph functionality.
                </p>
              </div>

              {/* Horizontal Stepper */}
              <div className="relative">
                {/* Step Circles - Evenly Distributed */}
                <div className="flex items-center justify-between">
                  {PIPELINE_STAGES.map((stage, index) => {
                    const state = stageStates[stage.id] || {
                      status: "idle" as const,
                      progress: 0,
                    };
                    const StageIcon = stage.icon;

                    return (
                      <div
                        key={stage.id}
                        className="flex flex-col items-center flex-1"
                      >
                        {/* Step Circle */}
                        <div
                          className={`
                            w-12 h-12 rounded-full border-2 flex items-center justify-center transition-all
                            ${
                              state.status === "completed"
                                ? "bg-green-500 border-green-500 text-white"
                                : state.status === "running"
                                ? "bg-blue-500 border-blue-500 text-white animate-pulse"
                                : state.status === "error"
                                ? "bg-red-500 border-red-500 text-white"
                                : "bg-background border-muted-foreground text-muted-foreground"
                            }
                          `}
                        >
                          {state.status === "completed" ? (
                            <CheckCircle className="h-6 w-6" />
                          ) : state.status === "running" ? (
                            <Clock className="h-6 w-6" />
                          ) : state.status === "error" ? (
                            <AlertCircle className="h-6 w-6" />
                          ) : (
                            <StageIcon className="h-6 w-6" />
                          )}
                        </div>

                        {/* Step Number */}
                        <span className="text-xs font-medium mt-1 text-muted-foreground">
                          Step {index + 1}
                        </span>
                      </div>
                    );
                  })}
                </div>

                {/* Connector Lines - Positioned Absolutely */}
                <div className="absolute top-6 left-0 right-0 flex items-center justify-between px-12">
                  {PIPELINE_STAGES.slice(0, -1).map((stage, _index) => {
                    const state = stageStates[stage.id] || {
                      status: "idle" as const,
                      progress: 0,
                    };

                    return (
                      <div
                        key={`connector-${stage.id}`}
                        className="flex-1 mx-8"
                      >
                        <div
                          className={`
                            h-0.5 transition-all
                            ${
                              state.status === "completed"
                                ? "bg-green-500"
                                : "bg-muted-foreground/30"
                            }
                          `}
                        />
                      </div>
                    );
                  })}
                </div>
              </div>

              {/* Stage Labels */}
              <div className="flex items-start justify-between">
                {PIPELINE_STAGES.map((stage, _index) => {
                  const state = stageStates[stage.id] || {
                    status: "idle" as const,
                    progress: 0,
                  };

                  return (
                    <div
                      key={`${stage.id}-label`}
                      className="flex-1 text-center"
                    >
                      <h3 className="font-medium text-sm">{stage.name}</h3>
                      <p className="text-xs text-muted-foreground mt-1 leading-relaxed">
                        {stage.description}
                      </p>
                      <Badge
                        variant="outline"
                        className={`${getStageStatusColor(
                          state.status
                        )} font-medium mt-2`}
                      >
                        {state.status === "idle"
                          ? "Ready"
                          : state.status === "running"
                          ? "Running"
                          : state.status === "completed"
                          ? "Completed"
                          : "Error"}
                      </Badge>
                    </div>
                  );
                })}
              </div>

              {/* Error Display */}
              {Object.values(stageStates).some((s) => s.error) && (
                <div className="p-4 bg-red-50 border border-red-200 rounded-lg">
                  <p className="text-red-700 font-medium">
                    {Object.values(stageStates).find((s) => s.error)?.error}
                  </p>
                </div>
              )}

              {/* Action Buttons - Aligned with Step Circles */}
              <div className="flex items-center justify-between">
                {PIPELINE_STAGES.map((stage) => {
                  const state = stageStates[stage.id] || {
                    status: "idle" as const,
                    progress: 0,
                  };

                  return (
                    <div
                      key={`${stage.id}-action`}
                      className="flex-1 flex justify-center"
                    >
                      {state.status === "idle" && (
                        <Button
                          onClick={() => runStage(stage)}
                          disabled={anyStageRunning}
                          className="gap-2"
                        >
                          {stage.id === "perturb" ? (
                            <>
                              <Settings className="h-4 w-4" />
                              Configure & Run
                            </>
                          ) : (
                            <>
                              <Play className="h-4 w-4" />
                              Run {stage.name}
                            </>
                          )}
                        </Button>
                      )}

                      {state.status === "completed" && (
                        <Button
                          variant="outline"
                          onClick={() => viewStageResults(stage.id)}
                          className="gap-2 bg-green-50 border-green-200 hover:bg-green-100"
                        >
                          <Eye className="h-4 w-4" />
                          View Results
                        </Button>
                      )}

                      {state.status === "error" && (
                        <Button
                          variant="outline"
                          onClick={() => runStage(stage)}
                          disabled={anyStageRunning}
                          className="gap-2"
                        >
                          <Play className="h-4 w-4" />
                          Retry
                        </Button>
                      )}

                      {state.status === "running" && (
                        <Button disabled className="gap-2">
                          <Clock className="h-4 w-4" />
                          Running...
                        </Button>
                      )}
                    </div>
                  );
                })}
              </div>
            </div>
          </CardContent>
        </Card>
      </div>

      {/* Stage Results Dialog */}
      {selectedStageResults && (
        <Dialog
          open={!!selectedStageResults}
          onOpenChange={() => setSelectedStageResults(null)}
        >
          <DialogContent className="max-w-4xl max-h-[80vh] overflow-y-auto">
            <DialogHeader>
              <DialogTitle>
                {
                  PIPELINE_STAGES.find(
                    (s) => s.id === selectedStageResults.stage
                  )?.name
                }{" "}
                Results
              </DialogTitle>
              <DialogDescription>
                Analysis results for knowledge graph {knowledgeGraph.kg_id}
              </DialogDescription>
            </DialogHeader>
            <div className="mt-4">
              {/* Render appropriate results component based on stage */}
              {selectedStageResults.stage === "enrich" && (
                <EnrichResults
                  data={selectedStageResults.data}
                  knowledgeGraphId={knowledgeGraph.kg_id}
                />
              )}
              {selectedStageResults.stage === "perturb" && (
                <PerturbResults
                  data={selectedStageResults.data}
                  knowledgeGraphId={knowledgeGraph.kg_id}
                />
              )}
              {selectedStageResults.stage === "causal" && (
                <div className="p-8 text-center">
                  <h3 className="text-lg font-semibold mb-2">
                    Causal Analysis Results
                  </h3>
                  <p className="text-muted-foreground mb-4">
                    Rich causal analysis visualization is in development.
                  </p>
                  <pre className="bg-muted p-4 rounded-lg text-sm overflow-auto text-left max-h-96">
                    {JSON.stringify(selectedStageResults.data, null, 2)}
                  </pre>
                </div>
              )}
            </div>
          </DialogContent>
        </Dialog>
      )}

      {/* Perturbation Test Configuration Dialog */}
      <PerturbationTestConfig
        open={showPerturbConfig}
        onOpenChange={setShowPerturbConfig}
        onRun={runPerturbWithConfig}
      />
    </>
  );
}