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

import React, { useState, useEffect, useMemo } from "react";
import { ExampleTraceLite, ExampleTrace } from "@/types";
import { api } from "@/lib/api";
import { Button } from "@/components/ui/button";
import { useToast } from "@/hooks/use-toast";
import { useAgentGraph } from "@/context/AgentGraphContext";
import { CheckCircle, AlertTriangle } from "lucide-react";
import ExampleTraceDetailModal from "./ExampleTraceDetailModal";
import { AgentMultiSelect } from "@/components/ui/agent-multi-select";

interface Props {
  data?: any;
  onClose: () => void;
}

// Remove subset concept - load all traces as "Who_and_When"
const SUBSETS = ["Who_and_When"] as const;

type CountFilter = 1 | 2 | 3 | null;

export const ExampleTraceModal: React.FC<Props> = ({ onClose: _onClose }) => {
  const [subset, setSubset] = useState<(typeof SUBSETS)[number]>(SUBSETS[0]);
  const [examples, setExamples] = useState<ExampleTraceLite[]>([]);
  const [loading, setLoading] = useState(false);

  const [selectedAgents, setSelectedAgents] = useState<string[]>([]);
  const [countFilter, setCountFilter] = useState<CountFilter>(null);

  const [detailOpen, setDetailOpen] = useState(false);
  const [selectedExample, setSelectedExample] = useState<ExampleTrace | null>(
    null
  );

  const { toast } = useToast();
  const { actions } = useAgentGraph();

  const load = async (s: string) => {
    setLoading(true);
    try {
      // For "Who_and_When", load all traces (no subset filter)
      const subsetParam = s === "Who_and_When" ? undefined : s;
      const list = await api.exampleTraces.list(subsetParam);
      setExamples(list);
    } catch (e: any) {
      toast({ title: "Error", description: e.message });
    } finally {
      setLoading(false);
    }
  };

  useEffect(() => {
    load(subset);
  }, [subset]);

  // Unique agent options for current examples
  const agentOptions = useMemo(() => {
    const set = new Set<string>();
    examples.forEach((ex) => ex.agents?.forEach((a) => set.add(a)));
    return Array.from(set).sort();
  }, [examples]);

  // Apply filters
  const filteredExamples = useMemo(() => {
    return examples.filter((ex) => {
      // agent multi-select filter
      if (selectedAgents.length) {
        const hasAll = selectedAgents.every((a) => ex.agents?.includes(a));
        if (!hasAll) return false;
      }

      // count filter
      if (countFilter) {
        const c = ex.agents?.length || 0;
        if (countFilter === 3) {
          if (c < 3) return false;
        } else if (c !== countFilter) {
          return false;
        }
      }

      return true;
    });
  }, [examples, selectedAgents, countFilter]);

  const handleImport = async (ex: ExampleTrace) => {
    setLoading(true);
    try {
      await api.exampleTraces.import(ex.subset, ex.id);
      toast({ title: "Imported", description: "Trace added to workspace" });
      const tracesData = await api.traces.list();
      actions.setTraces(Array.isArray(tracesData) ? tracesData : []);
    } catch (e: any) {
      toast({ title: "Import failed", description: e.message });
    } finally {
      setLoading(false);
    }
  };

  const clearFilters = () => {
    setSelectedAgents([]);
    setCountFilter(null);
  };

  // UI helpers
  const countBtnVariant = (val: CountFilter) =>
    val === countFilter ? "default" : "outline";

  return (
    <div className="flex flex-col h-[90vh]">
      {/* Toolbar */}
      <div className="flex flex-wrap items-center gap-4 pb-4 border-b p-4">
        {/* Subset buttons */}
        <div className="flex gap-2">
          {SUBSETS.map((s) => (
            <Button
              key={s}
              variant={s === subset ? "default" : "outline"}
              size="sm"
              onClick={() => setSubset(s)}
            >
              {s}
            </Button>
          ))}
        </div>

        {/* Agent multiselect */}
        <AgentMultiSelect
          options={agentOptions}
          value={selectedAgents}
          onChange={setSelectedAgents}
        />

        {/* Count filter */}
        <div className="flex items-center gap-2 ml-auto">
          <span className="text-sm text-muted-foreground">Agents:</span>
          {[1, 2, 3].map((n) => (
            <Button
              key={n}
              variant={countBtnVariant(n as CountFilter)}
              size="sm"
              onClick={() =>
                setCountFilter((prev) =>
                  prev === n ? null : (n as CountFilter)
                )
              }
            >
              {n === 3 ? "3+" : n}
            </Button>
          ))}
          {(selectedAgents.length > 0 || countFilter) && (
            <Button variant="ghost" size="sm" onClick={clearFilters}>
              Clear
            </Button>
          )}
        </div>
      </div>

      {/* Content */}
      <div className="flex-1 overflow-y-auto p-4">
        {loading ? (
          <p className="p-4 text-muted-foreground">Loading...</p>
        ) : filteredExamples.length === 0 ? (
          <p className="p-4 text-muted-foreground">No examples found</p>
        ) : (
          <div className="grid gap-6 grid-cols-1 sm:grid-cols-2 xl:grid-cols-3">
            {filteredExamples.map((ex) => (
              <div
                key={ex.id}
                className="border rounded-lg p-5 hover:shadow-md transition cursor-pointer flex flex-col gap-3 min-h-[150px] hover:border-primary bg-white"
                onClick={async () => {
                  try {
                    const full = await api.exampleTraces.get(ex.subset, ex.id);
                    setSelectedExample(full);
                    setDetailOpen(true);
                  } catch (err: any) {
                    toast({ title: "Error", description: err.message });
                  }
                }}
              >
                {/* Header with correctness indicator */}
                <div className="flex items-start justify-between gap-2">
                  <div className="font-medium line-clamp-3 flex-1">
                    {ex.question}
                  </div>
                  <div className="flex items-center gap-1 flex-shrink-0">
                    {ex.is_correct === true && (
                      <CheckCircle className="h-4 w-4 text-green-600" />
                    )}
                  </div>
                </div>

                {/* Agent info */}
                <div className="text-xs text-muted-foreground">
                  <div className="flex flex-wrap gap-1">
                    <span className="text-muted-foreground">agents:</span>
                    {ex.agents?.map((agent, index) => (
                      <span
                        key={agent}
                        className={`${
                          agent === ex.mistake_agent
                            ? "text-orange-700 font-medium"
                            : "text-gray-700"
                        }`}
                      >
                        {agent}
                        {agent === ex.mistake_agent ? " ⚠️" : ""}
                        {index < (ex.agents?.length || 0) - 1 ? "," : ""}
                      </span>
                    )) || <span>?</span>}
                  </div>
                </div>

                {/* Failure reason preview */}
                {ex.mistake_reason && (
                  <div className="flex items-start gap-2 bg-blue-50/50 rounded-md p-2 border border-blue-200">
                    <AlertTriangle className="h-3 w-3 text-blue-600 mt-0.5 flex-shrink-0" />
                    <p className="text-xs text-blue-700 line-clamp-2">
                      {ex.mistake_reason}
                    </p>
                  </div>
                )}
              </div>
            ))}
          </div>
        )}
      </div>

      {/* Detail Drawer */}
      {selectedExample && (
        <ExampleTraceDetailModal
          open={detailOpen}
          example={selectedExample}
          onOpenChange={(o) => setDetailOpen(o)}
          onImport={() => handleImport(selectedExample)}
        />
      )}
    </div>
  );
};

export default ExampleTraceModal;