frontend/src/components/features/processing/AdvancedProcessingView.tsx

import React, { useState, useCallback, useEffect } from "react";
import { KnowledgeGraph } from "@/types";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { Card, CardContent } from "@/components/ui/card";

import { Button } from "@/components/ui/button";
import {
  Tooltip,
  TooltipContent,
  TooltipProvider,
  TooltipTrigger,
} from "@/components/ui/tooltip";

import {
  Sparkles,
  Shield,
  GitBranch,
  Play,
  Clock,
  CheckCircle,
  AlertCircle,
  Trash2,
} from "lucide-react";
import { useNotification } from "@/context/NotificationContext";
import { EnrichResults } from "../traces/EnrichResults";
import { PerturbResults } from "../traces/PerturbResults";
import CausalResults from "../traces/CausalResults";

const API_BASE = "/api";

async function fetchApi<T>(
  endpoint: string,
  options?: RequestInit
): Promise<T> {
  const response = await fetch(`${API_BASE}${endpoint}`, {
    headers: {
      "Content-Type": "application/json",
      ...options?.headers,
    },
    ...options,
  });

  if (!response.ok) {
    throw new Error(`API Error: ${response.statusText}`);
  }

  return response.json();
}

interface AdvancedProcessingViewProps {
  knowledgeGraph: KnowledgeGraph;
  onBack?: () => 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 const AdvancedProcessingView: React.FC<AdvancedProcessingViewProps> = ({
  knowledgeGraph: initialKnowledgeGraph,
  onBack: _onBack,
}) => {
  const { showNotification } = useNotification();
  const [knowledgeGraph, setKnowledgeGraph] = useState(initialKnowledgeGraph);
  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 [stageResults, setStageResults] = useState<Record<string, any>>({});
  const [loadingResults, setLoadingResults] = useState<Record<string, boolean>>(
    {}
  );
  const [knowledgeGraphData, setKnowledgeGraphData] = useState<any>(null);
  const [progressTimers, setProgressTimers] = useState<
    Record<string, NodeJS.Timeout>
  >({});

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

  // Estimated durations for each stage (in seconds)
  const STAGE_DURATIONS = {
    enrich: 45, // Prompt Reconstruction typically takes ~45 seconds
    perturb: 90, // Perturbation Testing takes ~90 seconds
    causal: 60, // Causal Analysis takes ~60 seconds
  };

  // Start estimated progress animation
  const startEstimatedProgress = useCallback(
    (stageId: string) => {
      // Clear any existing timer
      if (progressTimers[stageId]) {
        clearInterval(progressTimers[stageId]);
      }

      const duration =
        STAGE_DURATIONS[stageId as keyof typeof STAGE_DURATIONS] || 60;
      const interval = 1000; // Update every second
      const increment = 100 / ((duration * 1000) / interval); // Progress per interval
      let currentProgress = 0;

      const timer = setInterval(() => {
        // Add some randomness to make progress feel more realistic
        const randomVariation = (Math.random() - 0.5) * 0.5; // ±0.25% variation
        currentProgress += increment + randomVariation;

        // Cap at 95% to avoid reaching 100% before completion
        if (currentProgress >= 95) {
          currentProgress = 95;
        }

        // Ensure progress doesn't go backwards
        if (currentProgress < 0) {
          currentProgress = 0;
        }

        setStageStates((prev) => ({
          ...prev,
          [stageId]: {
            ...prev[stageId],
            progress: Math.round(currentProgress),
          } as StageState,
        }));
      }, interval);

      setProgressTimers((prev) => ({
        ...prev,
        [stageId]: timer,
      }));
    },
    [progressTimers]
  );

  // Stop progress animation
  const stopEstimatedProgress = useCallback(
    (stageId: string) => {
      if (progressTimers[stageId]) {
        clearInterval(progressTimers[stageId]);
        setProgressTimers((prev) => {
          const updated = { ...prev };
          delete updated[stageId];
          return updated;
        });
      }
    },
    [progressTimers]
  );

  // Cleanup timers on unmount
  useEffect(() => {
    return () => {
      Object.values(progressTimers).forEach((timer) => clearInterval(timer));
    };
  }, [progressTimers]);

  // Function to refresh knowledge graph data
  const refreshKnowledgeGraph = useCallback(async () => {
    try {
      console.log("🔄 Refreshing knowledge graph status...");
      const response = await fetchApi<any>(
        `/knowledge-graphs/${knowledgeGraph.kg_id}/status`
      );
      console.log("📡 KG Status API Response:", response);
      setKnowledgeGraph((prev) => {
        console.log("📝 Previous KG state:", prev);
        const updated = {
          ...prev,
          is_enriched: response.is_enriched,
          is_perturbed: response.is_perturbed,
          is_analyzed: response.is_analyzed,
          status: response.status,
        };
        console.log("✅ Updated KG state:", updated);
        return updated;
      });
    } catch (error) {
      console.error("❌ Error refreshing knowledge graph status:", error);
    }
  }, [knowledgeGraph.kg_id]);

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

      const poll = async (): Promise<void> => {
        try {
          const response = await fetchApi<any>(`/tasks/${taskId}/status`);
          const { status } = response;

          // Note: We're using estimated progress instead of backend progress
          // The progress is managed by the startEstimatedProgress function

          if (status === "COMPLETED") {
            console.log(`🎉 Stage ${stageId} completed! Updating state...`);

            // Stop estimated progress animation and set to 100%
            stopEstimatedProgress(stageId);

            setStageStates((prev) => ({
              ...prev,
              [stageId]: {
                ...prev[stageId],
                status: "completed",
                progress: 100,
              } as StageState,
            }));
            showNotification({
              title: `${
                PIPELINE_STAGES.find((s) => s.id === stageId)?.name
              } completed successfully`,
              message: "The processing stage has been completed successfully.",
              type: "success",
            });
            // Refresh the knowledge graph data to get updated status flags
            console.log(
              `🔄 Calling refreshKnowledgeGraph for stage ${stageId}...`
            );
            await refreshKnowledgeGraph();
            console.log(`✅ Finished refreshing KG for stage ${stageId}`);
            return;
          }

          if (status === "FAILED") {
            // Stop estimated progress animation on failure
            stopEstimatedProgress(stageId);

            setStageStates((prev) => ({
              ...prev,
              [stageId]: {
                ...prev[stageId],
                status: "error",
                error: response.error || "Task failed",
                progress: 0,
              } as StageState,
            }));
            showNotification({
              title: `${
                PIPELINE_STAGES.find((s) => s.id === stageId)?.name
              } failed`,
              message: "The processing stage encountered an error.",
              type: "error",
            });
            return;
          }

          attempts++;
          if (attempts < maxAttempts && status === "RUNNING") {
            setTimeout(poll, 5000);
          }
        } catch (error) {
          console.error("Error polling task status:", error);
          attempts++;
          if (attempts < maxAttempts) {
            setTimeout(poll, 5000);
          }
        }
      };

      poll();
    },
    [showNotification, refreshKnowledgeGraph]
  );

  const runStage = async (stage: PipelineStageConfig) => {
    try {
      setStageStates((prev) => ({
        ...prev,
        [stage.id]: {
          ...prev[stage.id],
          status: "running",
          progress: 0,
          error: undefined,
        },
      }));

      // Start estimated progress animation
      startEstimatedProgress(stage.id);

      const response = await fetchApi<any>(
        stage.apiEndpoint(knowledgeGraph.kg_id),
        {
          method: "POST",
        }
      );
      const { task_id } = response;

      if (task_id) {
        setStageStates((prev) => ({
          ...prev,
          [stage.id]: {
            ...prev[stage.id],
            taskId: task_id,
          } as StageState,
        }));
        pollTaskStatus(stage.id, task_id);
      }
    } catch (error: any) {
      console.error(`Error running ${stage.name}:`, error);

      // Stop progress animation on error
      stopEstimatedProgress(stage.id);

      setStageStates((prev) => ({
        ...prev,
        [stage.id]: {
          ...prev[stage.id],
          status: "error",
          error: error.message,
          progress: 0,
        } as StageState,
      }));
      showNotification({
        title: `Failed to start ${stage.name}`,
        message: "Unable to start the processing stage. Please try again.",
        type: "error",
      });
    }
  };

  const clearStage = async (stage: PipelineStageConfig) => {
    try {
      const response = await fetchApi<any>(
        `/knowledge-graphs/${knowledgeGraph.kg_id}/stage-results/${stage.id}`,
        {
          method: "DELETE",
        }
      );

      const { cleared_stages, new_status } = response;

      // Update knowledge graph status
      setKnowledgeGraph((prev) => ({
        ...prev,
        status: new_status,
        is_enriched:
          new_status === "enriched" ||
          new_status === "perturbed" ||
          new_status === "analyzed",
        is_perturbed: new_status === "perturbed" || new_status === "analyzed",
        is_analyzed: new_status === "analyzed",
      }));

      // Clear results and reset stage states for cleared stages
      const updatedStageStates: Record<string, StageState> = {};
      const updatedStageResults: Record<string, any> = {};

      cleared_stages.forEach((clearedStage: string) => {
        updatedStageStates[clearedStage] = {
          status: "idle",
          progress: 0,
          error: undefined,
        };
        updatedStageResults[clearedStage] = undefined;
      });

      setStageStates((prev) => ({ ...prev, ...updatedStageStates }));
      setStageResults((prev) => {
        const newResults = { ...prev };
        cleared_stages.forEach((clearedStage: string) => {
          delete newResults[clearedStage];
        });
        return newResults;
      });

      showNotification({
        title: `${stage.name} cleared successfully`,
        message: `Cleared ${cleared_stages.join(", ")} stage(s) as requested.`,
        type: "success",
      });
    } catch (error: any) {
      console.error(`Error clearing ${stage.name}:`, error);
      showNotification({
        title: `Failed to clear ${stage.name}`,
        message: "Unable to clear the stage results. Please try again.",
        type: "error",
      });
    }
  };

  const fetchStageResults = useCallback(
    async (stageId: string) => {
      if (loadingResults[stageId]) return;

      setLoadingResults((prev) => ({ ...prev, [stageId]: true }));
      try {
        const response = await fetchApi<any>(
          `/knowledge-graphs/${knowledgeGraph.kg_id}/stage-results/${stageId}`
        );
        setStageResults((prev) => ({ ...prev, [stageId]: response }));
      } catch (error) {
        console.error("Error fetching stage results:", error);
        showNotification({
          title: "Failed to fetch results",
          message: "Unable to retrieve the results for this stage.",
          type: "error",
        });
      } finally {
        setLoadingResults((prev) => ({ ...prev, [stageId]: false }));
      }
    },
    [loadingResults, knowledgeGraph.kg_id, showNotification]
  );

  // Fetch knowledge graph data for component name mapping
  useEffect(() => {
    const fetchKnowledgeGraphData = async () => {
      try {
        const response = await fetchApi<{
          entities: Array<{ id: string; name: string; type: string }>;
          relations: Array<{
            id: string;
            type: string;
            source: string;
            target: string;
          }>;
        }>(`/knowledge-graphs/${knowledgeGraph.kg_id}`);
        setKnowledgeGraphData(response);
      } catch (error) {
        console.error("Error fetching knowledge graph data:", error);
        // Don't show notification for this as it's not critical - causal analysis will still work with IDs
      }
    };

    fetchKnowledgeGraphData();
  }, [knowledgeGraph.kg_id]);

  // Auto-fetch results when stage completes
  useEffect(() => {
    console.log("🔍 Auto-fetch effect triggered. StageStates:", stageStates);
    PIPELINE_STAGES.forEach((stage) => {
      const state = stageStates[stage.id];
      console.log(
        `📊 Stage ${stage.id} state:`,
        state,
        `Results exist:`,
        !!stageResults[stage.id],
        `Loading:`,
        !!loadingResults[stage.id]
      );
      if (
        state?.status === "completed" &&
        !stageResults[stage.id] &&
        !loadingResults[stage.id]
      ) {
        console.log(`🚀 Fetching results for completed stage: ${stage.id}`);
        fetchStageResults(stage.id);
      }
    });
  }, [stageStates, stageResults, loadingResults, fetchStageResults]);

  const renderStageContent = (stage: PipelineStageConfig) => {
    const state = stageStates[stage.id] || {
      status: "idle" as const,
      progress: 0,
    };

    return (
      <div className="space-y-6">
        {/* Progress */}
        {state.status === "running" && (
          <Card>
            <CardContent className="pt-6">
              <div className="space-y-2">
                <div className="flex items-center justify-between">
                  <span className="font-medium">Progress</span>
                  <span className="font-mono">
                    {Math.round(state.progress)}%
                  </span>
                </div>
                <div className="w-full bg-gray-200 rounded-full h-2">
                  <div
                    className="bg-blue-600 h-2 rounded-full transition-all duration-300"
                    style={{ width: `${state.progress}%` }}
                  />
                </div>
              </div>
            </CardContent>
          </Card>
        )}

        {/* Error Display */}
        {state.error && (
          <Card>
            <CardContent className="pt-6">
              <div className="p-4 bg-red-50 border border-red-200 rounded-lg">
                <p className="text-red-700 font-medium">{state.error}</p>
              </div>
            </CardContent>
          </Card>
        )}

        {/* Results Section */}
        {state.status === "completed" && stageResults[stage.id] ? (
          <div className="space-y-6">
            {/* Results Header with Actions */}
            <Card className="border-t-2 border-t-green-500 bg-gradient-to-r from-green-50 to-emerald-50">
              <CardContent className="p-4">
                <div className="flex items-center justify-between">
                  <div className="flex items-center gap-3">
                    <div className="flex items-center justify-center w-8 h-8 bg-green-100 rounded-full">
                      <CheckCircle className="h-4 w-4 text-green-600" />
                    </div>
                    <div>
                      <h4 className="font-semibold text-green-900">
                        {stage.name} Completed
                      </h4>
                      <p className="text-sm text-green-700">
                        Results are ready to view
                      </p>
                    </div>
                  </div>

                  {/* Action Buttons */}
                  <div className="flex items-center gap-3">
                    <Button
                      onClick={() => runStage(stage)}
                      disabled={anyStageRunning}
                      size="sm"
                      variant="outline"
                      className="gap-2 border-green-200 text-green-700 hover:bg-green-100"
                    >
                      <Play className="h-4 w-4" />
                      Run Again
                    </Button>
                    <Button
                      variant="outline"
                      onClick={() => clearStage(stage)}
                      disabled={anyStageRunning}
                      size="sm"
                      className="gap-2 text-red-600 border-red-200 hover:bg-red-50 hover:border-red-300"
                    >
                      <Trash2 className="h-4 w-4" />
                      Clear
                    </Button>
                  </div>
                </div>
              </CardContent>
            </Card>

            {/* Results Content */}
            <div>
              {stage.id === "enrich" && (
                <EnrichResults
                  data={stageResults[stage.id]}
                  knowledgeGraphId={knowledgeGraph.kg_id}
                  knowledgeGraphData={knowledgeGraphData}
                />
              )}
              {stage.id === "perturb" && (
                <PerturbResults
                  data={stageResults[stage.id]}
                  knowledgeGraphId={knowledgeGraph.kg_id}
                />
              )}
              {stage.id === "causal" && (
                <CausalResults
                  data={stageResults[stage.id]}
                  knowledgeGraphData={knowledgeGraphData}
                />
              )}
            </div>
          </div>
        ) : (
          /* Actions for non-completed stages */
          <Card className="border-dashed border-2 border-gray-200">
            <CardContent className="p-6">
              <div className="text-center space-y-4">
                {state.status === "idle" && (
                  <>
                    <div className="flex items-center justify-center w-12 h-12 bg-blue-50 rounded-full mx-auto">
                      <Play className="h-6 w-6 text-blue-600" />
                    </div>
                    <div>
                      <h4 className="font-medium text-gray-900 mb-1">
                        Ready to Run {stage.name}
                      </h4>
                      <p className="text-sm text-gray-600 mb-4">
                        {stage.description}
                      </p>
                      <Button
                        onClick={() => runStage(stage)}
                        disabled={anyStageRunning}
                        size="lg"
                        className="gap-2"
                      >
                        <Play className="h-5 w-5" />
                        Run {stage.name}
                      </Button>
                    </div>
                  </>
                )}

                {state.status === "completed" && !stageResults[stage.id] && (
                  <>
                    <div className="flex items-center justify-center w-12 h-12 bg-orange-50 rounded-full mx-auto">
                      <AlertCircle className="h-6 w-6 text-orange-600" />
                    </div>
                    <div>
                      <h4 className="font-medium text-gray-900 mb-1">
                        No Results Available
                      </h4>
                      <p className="text-sm text-gray-600 mb-4">
                        The stage completed but no results were found. Try
                        running again.
                      </p>
                      <Button
                        onClick={() => runStage(stage)}
                        disabled={anyStageRunning}
                        size="lg"
                        className="gap-2"
                      >
                        <Play className="h-5 w-5" />
                        Run Again
                      </Button>
                    </div>
                  </>
                )}

                {state.status === "error" && (
                  <>
                    <div className="flex items-center justify-center w-12 h-12 bg-red-50 rounded-full mx-auto">
                      <AlertCircle className="h-6 w-6 text-red-600" />
                    </div>
                    <div>
                      <h4 className="font-medium text-red-900 mb-1">
                        Error Occurred
                      </h4>
                      <p className="text-sm text-red-600 mb-4">
                        {state.error || "An error occurred during processing"}
                      </p>
                      <Button
                        variant="outline"
                        onClick={() => runStage(stage)}
                        disabled={anyStageRunning}
                        size="lg"
                        className="gap-2 border-red-200 text-red-700 hover:bg-red-50"
                      >
                        <Play className="h-5 w-5" />
                        Retry
                      </Button>
                    </div>
                  </>
                )}

                {state.status === "running" && (
                  <>
                    <div className="flex items-center justify-center w-12 h-12 bg-blue-50 rounded-full mx-auto">
                      <Clock className="h-6 w-6 text-blue-600 animate-spin" />
                    </div>
                    <div>
                      <h4 className="font-medium text-blue-900 mb-1">
                        Running {stage.name}
                      </h4>
                      <p className="text-sm text-blue-600 mb-4">
                        Processing in progress... {Math.round(state.progress)}%
                      </p>
                      <Button disabled size="lg" className="gap-2">
                        <Clock className="h-5 w-5" />
                        Running...
                      </Button>
                    </div>
                  </>
                )}

                {loadingResults[stage.id] && (
                  <>
                    <div className="flex items-center justify-center w-12 h-12 bg-purple-50 rounded-full mx-auto">
                      <Clock className="h-6 w-6 text-purple-600 animate-pulse" />
                    </div>
                    <div>
                      <h4 className="font-medium text-purple-900 mb-1">
                        Loading Results
                      </h4>
                      <p className="text-sm text-purple-600 mb-4">
                        Fetching {stage.name.toLowerCase()} results...
                      </p>
                      <Button disabled size="lg" className="gap-2">
                        <Clock className="h-5 w-5" />
                        Loading Results...
                      </Button>
                    </div>
                  </>
                )}
              </div>
            </CardContent>
          </Card>
        )}
      </div>
    );
  };

  return (
    <TooltipProvider>
      <div className="flex flex-col h-screen bg-background">
        <div className="flex-1 overflow-hidden">
          <div className="h-full p-6">
            <Tabs defaultValue="enrich" className="h-full flex flex-col">
              <TabsList className="grid w-full grid-cols-3 mb-6">
                {PIPELINE_STAGES.map((stage) => {
                  const state = stageStates[stage.id];
                  const StageIcon = stage.icon;

                  return (
                    <Tooltip key={stage.id}>
                      <TooltipTrigger asChild>
                        <TabsTrigger
                          value={stage.id}
                          className="flex items-center gap-2 relative"
                        >
                          <StageIcon className="h-4 w-4" />
                          {stage.name}
                          {state?.status === "completed" && (
                            <CheckCircle className="h-3 w-3 text-green-600 absolute -top-1 -right-1" />
                          )}
                          {state?.status === "running" && (
                            <Clock className="h-3 w-3 text-blue-600 absolute -top-1 -right-1 animate-pulse" />
                          )}
                          {state?.status === "error" && (
                            <AlertCircle className="h-3 w-3 text-red-600 absolute -top-1 -right-1" />
                          )}
                        </TabsTrigger>
                      </TooltipTrigger>
                      <TooltipContent>
                        <p>{stage.description}</p>
                      </TooltipContent>
                    </Tooltip>
                  );
                })}
              </TabsList>

              <div className="flex-1 overflow-y-auto">
                {PIPELINE_STAGES.map((stage) => (
                  <TabsContent
                    key={stage.id}
                    value={stage.id}
                    className="h-full mt-0"
                  >
                    <div className="max-w-5xl mx-auto">
                      {renderStageContent(stage)}
                    </div>
                  </TabsContent>
                ))}
              </div>
            </Tabs>
          </div>
        </div>
      </div>
    </TooltipProvider>
  );
};