frontend/src/components/shared/modals/MethodSelectionModal.tsx

import React, { useState, useEffect } from "react";
import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogHeader,
  DialogTitle,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Card, CardContent } from "@/components/ui/card";
import { Badge } from "@/components/ui/badge";
import { Brain, Zap, Star, Cpu, FlaskConical } from "lucide-react";
import { api } from "@/lib/api";

export type MethodType = "production" | string; // Allow for baseline method names

interface MethodOption {
  id: string;
  name: string;
  description: string;
  method_type: "production" | "baseline";
  schema_type: "reference_based" | "direct_based";
  supported_features: string[];
  processing_type: "async_crew" | "direct_call";
  recommended?: boolean;
}

interface MethodSelectionModalProps {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  onConfirm: (methodName: string) => void;
  onBack: () => void;
  isLoading?: boolean;
  selectedSplitter: string;
}

export function MethodSelectionModal({
  open,
  onOpenChange,
  onConfirm,
  onBack,
  isLoading = false,
  selectedSplitter,
}: MethodSelectionModalProps) {
  const [selectedMethod, setSelectedMethod] = useState<string>("openai_structured");
  const [methods, setMethods] = useState<MethodOption[]>([]);
  const [loadingMethods, setLoadingMethods] = useState(false);

  // Fetch available methods when modal opens
  useEffect(() => {
    if (open) {
      fetchMethods();
    }
  }, [open]);

  const fetchMethods = async () => {
    setLoadingMethods(true);
    try {
      const response = await api.methods.getAvailable();
      const methodList = Object.entries(response.methods).map(
        ([id, method]: [string, any]) => ({
          id,
          ...method,
          recommended: id === "openai_structured",
        })
      );
      // HuggingFace UI: hide all baseline methods from selection
      const productionOnly = methodList.filter(
        (m: any) => m.method_type === "production"
      );
      setMethods(productionOnly);
    } catch (error) {
      console.error("Failed to fetch methods:", error);
      // Fallback to default method
      setMethods([
        {
          id: "openai_structured",
          name: "OpenAI Structured Outputs",
          description:
            "Simple OpenAI structured outputs extractor using Pydantic models",
          method_type: "production",
          schema_type: "reference_based",
          supported_features: ["structured_outputs", "direct_extraction"],
          processing_type: "direct_call",
          recommended: true,
        },
      ]);
    } finally {
      setLoadingMethods(false);
    }
  };

  const handleConfirm = () => {
    const finalMethodName = selectedMethod || "openai_structured";
    console.log("MethodSelectionModal: selectedMethod state:", selectedMethod);
    console.log("MethodSelectionModal: finalMethodName:", finalMethodName);
    console.log(
      "MethodSelectionModal: About to call onConfirm with:",
      finalMethodName
    );
    onConfirm(finalMethodName);
  };

  const getMethodIcon = (method: MethodOption) => {
    if (method.method_type === "production") {
      return Brain;
    }
    if (method.supported_features.includes("clustering")) {
      return Cpu;
    }
    if (method.supported_features.includes("llm")) {
      return Zap;
    }
    return FlaskConical;
  };

  const getSplitterDisplayName = (splitterType: string) => {
    switch (splitterType) {
      case "agent_semantic":
        return "Agent Semantic Splitter";
      case "json":
        return "JSON Splitter";
      case "prompt_interaction":
        return "Prompt Interaction Splitter";
      default:
        return splitterType;
    }
  };

  if (loadingMethods) {
    return (
      <Dialog open={open} onOpenChange={onOpenChange}>
        <DialogContent className="sm:max-w-2xl">
          <div className="flex items-center justify-center p-8">
            <div className="w-6 h-6 border-2 border-primary border-t-transparent rounded-full animate-spin mr-3" />
            Loading extraction methods...
          </div>
        </DialogContent>
      </Dialog>
    );
  }

  // Group methods by type
  const productionMethods = methods.filter(
    (m) => m.method_type === "production"
  );
  // Ensure display order: recommended (openai_structured) first
  const sortedProductionMethods = [...productionMethods].sort((a, b) => {
    if (!!a.recommended === !!b.recommended) return 0;
    return a.recommended ? -1 : 1;
  });
  // Hide baseline methods in the HuggingFace version
  const baselineMethods: MethodOption[] = [];

  return (
    <Dialog open={open} onOpenChange={onOpenChange}>
      <DialogContent className="sm:max-w-3xl max-h-[90vh] overflow-y-auto">
        <DialogHeader>
          <DialogTitle className="flex items-center gap-2">
            <Zap className="h-5 w-5" />
            Select Extraction Method
          </DialogTitle>
          <DialogDescription>
            Choose the extraction method for knowledge graph generation:
          </DialogDescription>
        </DialogHeader>

        <div className="space-y-6 py-4">
          {/* Production Methods */}
          {productionMethods.length > 0 && (
            <div>
              <h3 className="font-semibold text-sm text-muted-foreground mb-3 uppercase tracking-wide">
                Production Methods
              </h3>
              <div className="space-y-3">
                {sortedProductionMethods.map((method) => {
                  const Icon = getMethodIcon(method);
                  const isSelected = selectedMethod === method.id;

                  return (
                    <Card
                      key={method.id}
                      className={`cursor-pointer transition-all duration-200 hover:shadow-md ${
                        isSelected
                          ? "ring-2 ring-primary border-primary bg-primary/5"
                          : "border-border hover:border-primary/50"
                      }`}
                      onClick={() => {
                        console.log(
                          "MethodSelectionModal: Clicking baseline method:",
                          method.id,
                          method.name
                        );
                        setSelectedMethod(method.id);
                      }}
                    >
                      <CardContent className="p-4">
                        <div className="flex items-start gap-4">
                          <div
                            className={`p-3 rounded-full ${
                              isSelected
                                ? "bg-primary text-primary-foreground"
                                : "bg-muted text-muted-foreground"
                            }`}
                          >
                            <Icon className="h-5 w-5" />
                          </div>

                          <div className="flex-1 space-y-2">
                            <div className="flex items-center justify-between">
                              <h4 className="font-semibold">{method.name}</h4>
                              {method.recommended && (
                                <Badge className="bg-amber-100 text-amber-800 border-amber-200">
                                  <Star className="h-3 w-3 mr-1" />
                                  Recommended
                                </Badge>
                              )}
                            </div>

                            <p className="text-sm text-muted-foreground leading-relaxed">
                              {method.description}
                            </p>

                            <div className="flex flex-wrap gap-2">
                              {method.supported_features.map((feature) => (
                                <Badge
                                  key={feature}
                                  variant="secondary"
                                  className="text-xs"
                                >
                                  {feature.replace(/_/g, " ")}
                                </Badge>
                              ))}
                            </div>
                          </div>
                        </div>
                      </CardContent>
                    </Card>
                  );
                })}
              </div>
            </div>
          )}

          {/* Baseline Methods */}
          {baselineMethods.length > 0 && (
            <div>
              <h3 className="font-semibold text-sm text-muted-foreground mb-3 uppercase tracking-wide">
                Baseline Methods (For Comparison)
              </h3>
              <div className="space-y-3 max-h-80 overflow-y-auto">
                {baselineMethods.map((method) => {
                  const Icon = getMethodIcon(method);
                  const isSelected = selectedMethod === method.id;

                  return (
                    <Card
                      key={method.id}
                      className={`cursor-pointer transition-all duration-200 hover:shadow-md ${
                        isSelected
                          ? "ring-2 ring-primary border-primary bg-primary/5"
                          : "border-border hover:border-primary/50"
                      }`}
                      onClick={() => {
                        console.log(
                          "MethodSelectionModal: Clicking method:",
                          method.id,
                          method.name
                        );
                        setSelectedMethod(method.id);
                      }}
                    >
                      <CardContent className="p-3">
                        <div className="flex items-center gap-3">
                          <div
                            className={`p-2 rounded-full ${
                              isSelected
                                ? "bg-primary text-primary-foreground"
                                : "bg-muted text-muted-foreground"
                            }`}
                          >
                            <Icon className="h-4 w-4" />
                          </div>

                          <div className="flex-1 min-w-0">
                            <h4 className="font-medium text-sm">
                              {method.name}
                            </h4>
                            <p className="text-xs text-muted-foreground truncate">
                              {method.description}
                            </p>
                          </div>

                          <div className="flex gap-1">
                            {method.supported_features
                              .slice(0, 2)
                              .map((feature) => (
                                <Badge
                                  key={feature}
                                  variant="outline"
                                  className="text-xs"
                                >
                                  {feature.replace(/_/g, " ")}
                                </Badge>
                              ))}
                          </div>
                        </div>
                      </CardContent>
                    </Card>
                  );
                })}
              </div>
            </div>
          )}
        </div>

        <div className="flex justify-between gap-3 pt-4 border-t">
          <Button variant="outline" onClick={onBack} disabled={isLoading}>
            Back
          </Button>
          <div className="flex gap-2">
            <Button
              variant="outline"
              onClick={() => onOpenChange(false)}
              disabled={isLoading}
            >
              Cancel
            </Button>
            <Button onClick={handleConfirm} disabled={isLoading}>
              {isLoading ? (
                <>
                  <div className="w-4 h-4 border-2 border-primary border-t-transparent rounded-full animate-spin mr-2" />
                  Generating...
                </>
              ) : (
                <>
                  <Brain className="h-4 w-4 mr-2" />
                  Generate Agent Graph
                </>
              )}
            </Button>
          </div>
        </div>
      </DialogContent>
    </Dialog>
  );
}