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

import React, { useState, useEffect } from "react";
import {
  Dialog,
  DialogContent,
  DialogHeader,
  DialogTitle,
  DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Checkbox } from "@/components/ui/checkbox";
import {
  FileText,
  Database,
  Activity,
  Clock,
  HardDrive,
  Globe,
  Settings,
  TreePine,
  Zap,
  MessageSquare,
  Target,
  BarChart3,
  Layers,
  DollarSign,
  Upload,
} from "lucide-react";
import { MetadataCardConfig } from "../CompactMetadataCard";

interface MetadataCardSelectorProps {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  selectedCards: string[];
  onSave: (selectedCards: string[]) => void;
  trace: {
    character_count?: number;
    trace_type?: string;
    upload_timestamp?: string;
    filename?: string;
    metadata?: Record<string, any>;
  };
  graphCount: number;
}

export const AVAILABLE_CARDS: Omit<MetadataCardConfig, "value">[] = [
  // Basic metadata cards
  {
    id: "size",
    label: "Size",
    icon: FileText,
    color: "text-blue-600",
    tooltip: "Total character count of the trace content",
  },
  {
    id: "type",
    label: "Type",
    icon: Database,
    color: "text-green-600",
    tooltip: "Type of trace data (e.g., example, production)",
  },
  {
    id: "uploaded",
    label: "Uploaded",
    icon: Upload,
    color: "text-purple-600",
    tooltip: "Date when the trace was uploaded",
  },
  {
    id: "graphs",
    label: "Graphs",
    icon: Activity,
    color: "text-amber-600",
    tooltip: "Number of generated agent graphs",
  },
  {
    id: "modified",
    label: "Modified",
    icon: Clock,
    color: "text-orange-600",
    tooltip: "Last modification date",
  },
  {
    id: "filesize",
    label: "File Size",
    icon: HardDrive,
    color: "text-gray-600",
    tooltip: "Approximate file size on disk",
  },
  {
    id: "source",
    label: "Source",
    icon: Globe,
    color: "text-indigo-600",
    tooltip: "Source platform or system",
  },
  {
    id: "method",
    label: "Method",
    icon: Settings,
    color: "text-pink-600",
    tooltip: "Processing method used for extraction",
  },

  // Schema analytics cards
  {
    id: "depth",
    label: "Depth",
    icon: TreePine,
    color: "text-emerald-600",
    tooltip: "Maximum depth of component hierarchy in the trace",
  },
  {
    id: "execution_time",
    label: "Exec Time",
    icon: Zap,
    color: "text-yellow-600",
    tooltip: "Total execution time of all components",
  },
  {
    id: "total_tokens",
    label: "Tokens",
    icon: MessageSquare,
    color: "text-cyan-600",
    tooltip: "Total token count (input + output)",
  },
  {
    id: "prompt_calls",
    label: "Prompts",
    icon: Target,
    color: "text-rose-600",
    tooltip: "Number of LLM/prompt calls detected",
  },
  {
    id: "components",
    label: "Components",
    icon: Layers,
    color: "text-violet-600",
    tooltip: "Total number of components in the trace",
  },
  {
    id: "success_rate",
    label: "Success Rate",
    icon: BarChart3,
    color: "text-teal-600",
    tooltip: "Percentage of components that executed successfully",
  },

  // Cost cards
  {
    id: "total_cost",
    label: "Total Cost",
    icon: DollarSign,
    color: "text-purple-600",
    tooltip: "Total estimated cost in USD",
  },
  {
    id: "cost_per_call",
    label: "Cost per Call",
    icon: DollarSign,
    color: "text-blue-600",
    tooltip: "Average cost per LLM call",
  },
  {
    id: "input_cost",
    label: "Input Cost",
    icon: DollarSign,
    color: "text-green-600",
    tooltip: "Cost for input tokens",
  },
  {
    id: "output_cost",
    label: "Output Cost",
    icon: DollarSign,
    color: "text-red-600",
    tooltip: "Cost for output tokens",
  },

  // Token analytics cards
  {
    id: "avg_input_tokens",
    label: "Avg Input Tokens",
    icon: MessageSquare,
    color: "text-purple-600",
    tooltip: "Average input tokens per call",
  },
  {
    id: "avg_output_tokens",
    label: "Avg Output Tokens",
    icon: MessageSquare,
    color: "text-blue-600",
    tooltip: "Average output tokens per call",
  },
];

export const DEFAULT_CARDS = ["uploaded", "graphs", "size", "type"];

export function MetadataCardSelector({
  open,
  onOpenChange,
  selectedCards,
  onSave,
  trace,
  graphCount,
}: MetadataCardSelectorProps) {
  const [localSelection, setLocalSelection] = useState<string[]>(selectedCards);

  useEffect(() => {
    setLocalSelection(selectedCards);
  }, [selectedCards, open]);

  const handleToggleCard = (cardId: string) => {
    setLocalSelection((prev) =>
      prev.includes(cardId)
        ? prev.filter((id) => id !== cardId)
        : [...prev, cardId]
    );
  };

  const handleSave = () => {
    onSave(localSelection);
    onOpenChange(false);
  };

  const handleReset = () => {
    console.log("Reset clicked, current selection:", localSelection);
    console.log("Setting to default:", DEFAULT_CARDS);
    setLocalSelection([...DEFAULT_CARDS]); // Create a new array to ensure state update
  };

  const formatDate = (timestamp?: string) => {
    if (!timestamp) return "N/A";
    return new Date(timestamp).toLocaleDateString();
  };

  const getCardValue = (cardId: string): string | number => {
    // Get schema analytics from trace metadata
    const schemaAnalytics = trace.metadata?.schema_analytics;

    switch (cardId) {
      case "size":
        return trace.character_count || 0;
      case "type":
        return trace.trace_type || "Unknown";
      case "uploaded":
        return formatDate(trace.upload_timestamp);
      case "graphs":
        return graphCount;
      case "modified":
        return formatDate(trace.upload_timestamp); // Fallback to upload date
      case "filesize":
        return `${Math.round((trace.character_count || 0) / 1024)} KB`;
      case "source":
        return "Manual Upload"; // Default value
      case "method":
        return "Production"; // Default value

      // Schema analytics cards
      case "depth":
        return (
          schemaAnalytics?.numerical_overview?.component_stats?.max_depth ||
          "N/A"
        );
      case "execution_time": {
        const totalTimeMs =
          schemaAnalytics?.numerical_overview?.timing_analytics
            ?.total_execution_time_ms;
        if (totalTimeMs && totalTimeMs > 0) {
          if (totalTimeMs >= 1000) {
            return `${(totalTimeMs / 1000).toFixed(1)}s`;
          } else {
            return `${totalTimeMs}ms`;
          }
        }
        return "N/A";
      }
      case "total_tokens":
        return (
          schemaAnalytics?.numerical_overview?.token_analytics?.total_tokens ||
          "N/A"
        );
      case "prompt_calls":
        return (
          schemaAnalytics?.prompt_analytics?.prompt_calls_detected || "N/A"
        );
      case "components":
        return (
          schemaAnalytics?.numerical_overview?.component_stats
            ?.total_components || "N/A"
        );
      case "success_rate": {
        const successRate =
          schemaAnalytics?.numerical_overview?.component_stats?.success_rate;
        if (successRate !== undefined && successRate !== null) {
          return `${successRate.toFixed(1)}%`;
        }
        return "N/A";
      }

      // Cost cards
      case "total_cost":
        return trace.metadata?.cost_analytics?.total_cost_usd || "N/A";
      case "cost_per_call":
        return trace.metadata?.cost_analytics?.avg_cost_per_call_usd || "N/A";
      case "input_cost":
        return trace.metadata?.cost_analytics?.input_cost_usd || "N/A";
      case "output_cost":
        return trace.metadata?.cost_analytics?.output_cost_usd || "N/A";

      // Token analytics cards
      case "avg_input_tokens":
        return trace.metadata?.token_analytics?.avg_input_tokens || "N/A";
      case "avg_output_tokens":
        return trace.metadata?.token_analytics?.avg_output_tokens || "N/A";

      default:
        return "N/A";
    }
  };

  return (
    <Dialog open={open} onOpenChange={onOpenChange}>
      <DialogContent className="sm:max-w-md">
        <DialogHeader>
          <DialogTitle>Customize Metadata Cards</DialogTitle>
        </DialogHeader>

        <div className="space-y-4">
          <p className="text-sm text-muted-foreground">
            Select which metadata cards to display. You can choose up to 6
            cards.
          </p>

          <div className="space-y-3 max-h-[300px] overflow-y-auto">
            {AVAILABLE_CARDS.map((card) => {
              const Icon = card.icon;
              const isSelected = localSelection.includes(card.id);
              const value = getCardValue(card.id);

              return (
                <div
                  key={card.id}
                  className={`
                    flex items-center space-x-3 p-3 rounded-lg border transition-colors
                    ${
                      isSelected
                        ? "bg-primary/5 border-primary/20"
                        : "border-border hover:bg-muted/50"
                    }
                  `}
                >
                  <Checkbox
                    id={card.id}
                    checked={isSelected}
                    onCheckedChange={() => handleToggleCard(card.id)}
                  />
                  <div className="flex items-center gap-3 flex-1">
                    <Icon className={`h-4 w-4 ${card.color}`} />
                    <div className="flex-1">
                      <div className="flex items-center justify-between">
                        <label
                          htmlFor={card.id}
                          className="text-sm font-medium cursor-pointer"
                        >
                          {card.label}
                        </label>
                        <span className="text-sm text-muted-foreground">
                          {typeof value === "number" && value > 999
                            ? value.toLocaleString()
                            : value}
                        </span>
                      </div>
                      {card.tooltip && (
                        <p className="text-xs text-muted-foreground mt-1">
                          {card.tooltip}
                        </p>
                      )}
                    </div>
                  </div>
                </div>
              );
            })}
          </div>
        </div>

        <DialogFooter className="gap-2">
          <Button variant="outline" onClick={handleReset}>
            Reset to Default
          </Button>
          <Button variant="outline" onClick={() => onOpenChange(false)}>
            Cancel
          </Button>
          <Button onClick={handleSave}>Save Changes</Button>
        </DialogFooter>
      </DialogContent>
    </Dialog>
  );
}