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

import React, { useState } from "react";
import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogHeader,
  DialogTitle,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { Card, CardContent } from "@/components/ui/card";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import {
  Brain,
  Code,
  MessageSquare,
  Star,
  Database,
  ArrowRightLeft,
  Settings,
} from "lucide-react";
import { MethodSelectionModal } from "./MethodSelectionModal";
import { useModelPreferences } from "@/hooks/useModelPreferences";

export type SplitterType = "agent_semantic" | "json" | "prompt_interaction";

export interface ChunkingConfig {
  min_chunk_size?: number;
  max_chunk_size?: number;
}

interface SplitterOption {
  id: SplitterType;
  name: string;
  description: string;
  icon: React.ElementType;
  badge?: {
    text: string;
    icon: React.ElementType;
  };
  recommended?: boolean;
}

interface SplitterSelectionModalProps {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  onConfirm: (
    splitterType: SplitterType,
    methodName?: string,
    modelName?: string,
    chunkingConfig?: ChunkingConfig
  ) => void;
  isLoading?: boolean;
}

// HF Spaces variant: only expose the Agent Semantic Splitter
const SPLITTER_OPTIONS: SplitterOption[] = [
  {
    id: "agent_semantic",
    name: "Agent Semantic Splitter",
    description:
      "Intelligently splits content based on semantic meaning and agent interactions. Best for conversational traces and complex content.",
    icon: Brain,
    badge: {
      text: "Recommended",
      icon: Star,
    },
    recommended: true,
  },
];

export function SplitterSelectionModal({
  open,
  onOpenChange,
  onConfirm,
  isLoading = false,
}: SplitterSelectionModalProps) {
  const [selectedSplitter, setSelectedSplitter] =
    useState<SplitterType>("agent_semantic");
  const [showMethodSelection, setShowMethodSelection] = useState(false);
  const [chunkingConfig, setChunkingConfig] = useState<ChunkingConfig>({
    min_chunk_size: 100000, // 100K chars default
    max_chunk_size: 300000, // 300K chars default
  });
  const { selectedModel } = useModelPreferences();

  const handleSplitterConfirm = () => {
    setShowMethodSelection(true);
  };

  const handleMethodConfirm = (methodName: string) => {
    const finalMethodName = methodName || "production";
    console.log("SplitterSelectionModal: Final method name:", finalMethodName);
    console.log("SplitterSelectionModal: Selected model:", selectedModel);
    console.log("SplitterSelectionModal: Chunking config:", chunkingConfig);

    // Close the modal immediately when user clicks Generate Agent Graph
    setShowMethodSelection(false);
    onOpenChange(false);

    // Then trigger the generation process with model parameter and chunking config
    onConfirm(selectedSplitter, finalMethodName, selectedModel, chunkingConfig);
  };

  const handleMethodBack = () => {
    setShowMethodSelection(false);
  };

  const handleCancel = () => {
    setShowMethodSelection(false);
    onOpenChange(false);
  };

  // Reset state when modal closes
  const handleOpenChange = (open: boolean) => {
    if (!open) {
      setShowMethodSelection(false);
    }
    onOpenChange(open);
  };

  return (
    <>
      <Dialog
        open={open && !showMethodSelection}
        onOpenChange={handleOpenChange}
      >
        <DialogContent className="sm:max-w-2xl">
          <DialogHeader>
            <DialogTitle className="flex items-center gap-2">
              <Brain className="h-5 w-5" />
              Select Chunking Strategy
            </DialogTitle>
            <DialogDescription>
              Choose how you want to split your trace content for knowledge
              graph generation:
            </DialogDescription>
          </DialogHeader>

          <div className="space-y-4 py-4">
            {SPLITTER_OPTIONS.map((option) => {
              const Icon = option.icon;
              const BadgeIcon = option.badge?.icon;
              const isSelected = selectedSplitter === option.id;

              return (
                <Card
                  key={option.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={() => setSelectedSplitter(option.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">{option.name}</h4>
                          {option.badge && (
                            <Badge
                              variant={
                                option.recommended ? "default" : "secondary"
                              }
                              className={`${
                                option.recommended
                                  ? "bg-amber-100 text-amber-800 border-amber-200"
                                  : ""
                              }`}
                            >
                              {BadgeIcon && (
                                <BadgeIcon className="h-3 w-3 mr-1" />
                              )}
                              {option.badge.text}
                            </Badge>
                          )}
                        </div>

                        <p className="text-sm text-muted-foreground leading-relaxed">
                          {option.description}
                        </p>
                      </div>
                    </div>
                  </CardContent>
                </Card>
              );
            })}
          </div>

          {/* Configuration for Agent Semantic Splitter */}
          {selectedSplitter === "agent_semantic" && (
            <div className="space-y-4 py-4 border-t">
              <div className="flex items-center gap-2 mb-3">
                <Settings className="h-4 w-4 text-muted-foreground" />
                <h4 className="font-medium">Configuration</h4>
              </div>

              <div className="grid grid-cols-2 gap-4">
                <div className="space-y-2">
                  <Label htmlFor="min-chunk-size">
                    Minimum Characters
                    <span className="text-xs text-muted-foreground ml-1">
                      (≈25K tokens)
                    </span>
                  </Label>
                  <Input
                    id="min-chunk-size"
                    type="number"
                    min="10000"
                    max="1000000"
                    step="10000"
                    value={chunkingConfig.min_chunk_size || ""}
                    onChange={(e) =>
                      setChunkingConfig((prev) => ({
                        ...prev,
                        min_chunk_size: parseInt(e.target.value) || 100000,
                      }))
                    }
                    placeholder="100000"
                    className="text-sm"
                  />
                  <p className="text-xs text-muted-foreground">
                    Minimum chunk size in characters
                  </p>
                </div>

                <div className="space-y-2">
                  <Label htmlFor="max-chunk-size">
                    Maximum Characters
                    <span className="text-xs text-muted-foreground ml-1">
                      (≈75K tokens)
                    </span>
                  </Label>
                  <Input
                    id="max-chunk-size"
                    type="number"
                    min="50000"
                    max="2000000"
                    step="10000"
                    value={chunkingConfig.max_chunk_size || ""}
                    onChange={(e) =>
                      setChunkingConfig((prev) => ({
                        ...prev,
                        max_chunk_size: parseInt(e.target.value) || 300000,
                      }))
                    }
                    placeholder="300000"
                    className="text-sm"
                  />
                  <p className="text-xs text-muted-foreground">
                    Maximum chunk size in characters
                  </p>
                </div>
              </div>

              <div className="bg-muted/30 rounded-lg p-3">
                <div className="flex items-center gap-2">
                  <Brain className="h-4 w-4 text-blue-500" />
                  <p className="text-xs text-muted-foreground">
                    <span className="font-medium text-foreground">
                      Smart Chunking:
                    </span>{" "}
                    Balance context preservation with processing speed -
                    defaults optimized for most traces.
                  </p>
                </div>
              </div>
            </div>
          )}

          <div className="flex justify-end gap-3 pt-4 border-t">
            <Button
              variant="outline"
              onClick={handleCancel}
              disabled={isLoading}
            >
              Cancel
            </Button>
            <Button onClick={handleSplitterConfirm} disabled={isLoading}>
              <Brain className="h-4 w-4 mr-2" />
              Next: Select Method
            </Button>
          </div>
        </DialogContent>
      </Dialog>

      <MethodSelectionModal
        open={showMethodSelection}
        onOpenChange={handleCancel}
        onConfirm={handleMethodConfirm}
        onBack={handleMethodBack}
        isLoading={isLoading}
        selectedSplitter={selectedSplitter}
      />
    </>
  );
}