src/app/components/DataDictionaryForm.tsx

import { useState, useEffect, useCallback } from "react";
import { Save, X, Loader2, HelpCircle } from "lucide-react";
import { Button } from "./ui/button";
import { Input } from "./ui/input";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "./ui/select";
import { Label } from "./ui/label";
import { Card, CardContent, CardHeader, CardTitle } from "./ui/card";
import { Checkbox } from "./ui/checkbox";
import { toast } from "sonner";
import { getDataCatalogSource, saveDataDictionary } from "../../services/orchestrationApi";

export interface DataDictionaryColumn {
  column_id: string;
  name: string;
  data_type: string;
  nullable: boolean;
  pii_flag: boolean;
  sample_values?: any[];
  stats?: {
    min?: any;
    max?: any;
    mean?: number | null;
    median?: number | null;
    distinct_count?: number;
  };
  description?: string;
  aggregation?: string; // sum, avg, min, max, count
  classify?: string; // dimension or measure
  business_meaning?: string;
  data_classification?: string;
  is_key?: boolean;
}

export interface DataDictionaryTable {
  table_id: string;
  name: string;
  row_count?: number;
  columns: DataDictionaryColumn[];
}

export interface DataDictionarySource {
  source_id: string;
  source_type: "schema" | "tabular" | "unstructured";
  name: string;
  tables: DataDictionaryTable[];
}

interface DataDictionaryFormProps {
  open: boolean;
  onClose: () => void;
  sourceId: string;
  sourceName: string;
  userId: string;
  selectedTableId?: string;
  onTableChange?: (tableId: string) => void;
}

const AGGREGATION_OPTIONS = [
  { value: "none", label: "None" },
  { value: "sum", label: "Sum" },
  { value: "avg", label: "Average" },
  { value: "min", label: "Min" },
  { value: "max", label: "Max" },
  { value: "count", label: "Count" },
];

const CLASSIFY_OPTIONS = [
  { value: "none", label: "Select..." },
  { value: "dimension", label: "Dimension" },
  { value: "measure", label: "Measure" },
];

export default function DataDictionaryForm({
  open,
  onClose,
  sourceId,
  sourceName,
  userId,
  selectedTableId,
  onTableChange,
}: DataDictionaryFormProps) {
  const [sources, setSources] = useState<DataDictionarySource[]>([]);
  const [loading, setLoading] = useState(false);
  const [saving, setSaving] = useState(false);
  const [activeTableId, setActiveTableId] = useState<string | undefined>(selectedTableId);
  const [editedColumns, setEditedColumns] = useState<Record<string, Partial<DataDictionaryColumn>>>({});

  useEffect(() => {
    if (open) {
      loadCatalog();
    }
  }, [open, sourceId, userId]);

  useEffect(() => {
    if (selectedTableId) {
      setActiveTableId(selectedTableId);
    }
  }, [selectedTableId]);

  const loadCatalog = async () => {
    setLoading(true);
    try {
      const source = await getDataCatalogSource(userId, sourceId);
      if (source) {
        // Transform to our type - source now contains full tables/columns
        const transformed: DataDictionarySource = {
          source_id: source.source_id,
          source_type: source.source_type as "schema" | "tabular" | "unstructured",
          name: source.name,
          tables: (source.tables || []).map((t: any) => ({
            table_id: t.table_id || t.name,
            name: t.name,
            row_count: t.row_count,
            columns: (t.columns || []).map((c: any) => ({
              column_id: c.column_id,
              name: c.name,
              data_type: c.data_type,
              nullable: c.nullable,
              pii_flag: c.pii_flag,
              sample_values: c.sample_values,
              stats: c.stats,
              description: c.description,
              business_meaning: c.business_meaning,
              data_classification: c.data_classification,
              is_key: c.is_key,
            })),
          })),
        };
        setSources([transformed]);
        // Auto-select first table if none selected
        if (!activeTableId && transformed.tables.length > 0) {
          setActiveTableId(transformed.tables[0].table_id);
          onTableChange?.(transformed.tables[0].table_id);
        }
      }
    } catch (err) {
      console.error("Failed to load catalog:", err);
      toast.error("Failed to load data dictionary");
    } finally {
      setLoading(false);
    }
  };

  const handleColumnChange = useCallback(
    (columnId: string, field: keyof DataDictionaryColumn, value: any) => {
      setEditedColumns((prev) => ({
        ...prev,
        [columnId]: {
          ...prev[columnId],
          [field]: value,
        },
      }));
    },
    []
  );

  const handleSave = async () => {
    if (saving) return;
    setSaving(true);
    try {
      const source = sources[0];
      if (!source) throw new Error("Source not found");

      // Prepare updated tables - apply edits to ALL tables
      const updatedTables = source.tables.map((table) => {
        const updatedColumns = table.columns.map((col) => {
          const edits = editedColumns[col.column_id] || {};
          return {
            ...col,
            ...edits,
          };
        });
        return { ...table, columns: updatedColumns };
      });

      // Update the source in catalog
      const updatedSource = {
        ...source,
        tables: updatedTables,
      };

      await saveDataDictionary(userId, sourceId, updatedSource);
      toast.success("Data dictionary saved successfully");
      onClose();
    } catch (err) {
      console.error("Failed to save data dictionary:", err);
      toast.error(err instanceof Error ? err.message : "Failed to save data dictionary");
    } finally {
      setSaving(false);
    }
  };

  if (!open) return null;

  const source = sources[0];
  const activeTable = source?.tables.find((t) => t.table_id === activeTableId);

  return (
    <div className="fixed inset-0 z-50 flex items-center justify-center bg-black/40 p-4 backdrop-blur-sm">
      <div className="flex max-h-[90vh] w-full max-w-5xl flex-col rounded-2xl bg-white shadow-xl">
        {/* Header */}
        <div className="flex items-center justify-between px-5 py-4 border-b border-slate-100">
          <div className="flex items-center gap-3">
            <button onClick={onClose} className="p-1.5 -ml-1 rounded-lg text-slate-400 hover:text-slate-700 hover:bg-slate-100 transition">
              <X className="w-4 h-4" />
            </button>
            <div className="w-7 h-7 rounded-lg bg-orange-100 flex items-center justify-center">
              <HelpCircle className="w-4 h-4 text-[#FF8F00]" />
            </div>
            <div>
              <h2 className="text-sm font-semibold text-slate-900">Data Dictionary</h2>
              <p className="text-xs text-slate-400">{sourceName}</p>
            </div>
          </div>
        </div>

        {/* Content */}
        <div className="flex-1 overflow-y-auto px-5 py-4">
          {loading ? (
            <div className="flex justify-center py-10">
              <Loader2 className="w-5 h-5 animate-spin text-slate-300" />
            </div>
          ) : source ? (
            <div className="space-y-4">
              {/* Table selector for multi-table sources */}
              {source.tables.length > 1 && (
                <div className="flex items-center gap-2 px-3 py-2.5 rounded-xl border border-slate-100">
                  <Label htmlFor="table-select" className="text-xs font-medium text-slate-500 mr-2">
                    Select Table/Worksheet:
                  </Label>
                  <Select value={activeTableId} onValueChange={setActiveTableId}>
                    <SelectTrigger id="table-select" className="w-64">
                      <SelectValue placeholder="Choose a table..." />
                    </SelectTrigger>
                    <SelectContent>
                      {source.tables.map((table) => (
                        <SelectItem key={table.table_id} value={table.table_id}>
                          {table.name} ({table.columns.length} columns)
                        </SelectItem>
                      ))}
                    </SelectContent>
                  </Select>
                </div>
              )}

              {/* Columns Table */}
              {activeTable && (
                <Card>
                  <CardHeader className="pb-2">
                    <CardTitle className="text-base">
                      {activeTable.name} ({activeTable.columns.length} columns)
                      {activeTable.row_count && (
                        <span className="ml-2 text-xs text-slate-400 font-normal">
                          {activeTable.row_count.toLocaleString()} rows
                        </span>
                      )}
                    </CardTitle>
                  </CardHeader>
                  <CardContent className="p-0">
                    <div className="overflow-x-auto">
                      <table className="w-full">
                        <thead>
                          <tr className="border-b border-slate-100 bg-slate-50">
                            <th className="px-3 py-2 text-left text-xs font-semibold text-slate-500 uppercase tracking-wider w-8">
                              #
                            </th>
                            <th className="px-3 py-2 text-left text-xs font-semibold text-slate-500 uppercase tracking-wider">
                              Column
                            </th>
                            <th className="px-3 py-2 text-left text-xs font-semibold text-slate-500 uppercase tracking-wider w-24">
                              Type
                            </th>
                            <th className="px-3 py-2 text-left text-xs font-semibold text-slate-500 uppercase tracking-wider w-16">
                              Nullable
                            </th>
                            <th className="px-3 py-2 text-left text-xs font-semibold text-slate-500 uppercase tracking-wider">
                              Description
                            </th>
                            <th className="px-3 py-2 text-left text-xs font-semibold text-slate-500 uppercase tracking-wider w-32">
                              Classify
                            </th>
                            <th className="px-3 py-2 text-left text-xs font-semibold text-slate-500 uppercase tracking-wider w-32">
                              Aggregation
                            </th>
                          </tr>
                        </thead>
                        <tbody className="divide-y divide-slate-100">
                          {activeTable.columns.map((col, index) => {
                            const edits = editedColumns[col.column_id] || {};
                            const classify = (edits.classify ?? col.classify ?? "") as string;
                            const isMeasure = classify === "measure";
                            return (
                              <tr key={col.column_id} className="hover:bg-slate-50/50">
                                <td className="px-3 py-2 text-xs text-slate-400 font-mono">{index + 1}</td>
                                <td className="px-3 py-2">
                                  <div className="flex items-center gap-2">
                                    <span className="text-sm font-medium text-slate-800 truncate max-w-[150px]">
                                      {col.name}
                                    </span>
                                  </div>
                                </td>
                                <td className="px-3 py-2 text-xs text-slate-500 font-mono">{col.data_type}</td>
                                <td className="px-3 py-2">
                                  <Checkbox
                                    checked={col.nullable}
                                    disabled
                                    className="text-slate-400"
                                  />
                                </td>
                                <td className="px-3 py-2">
                                  <Input
                                    value={edits.description ?? col.description ?? ""}
                                    onChange={(e) => handleColumnChange(col.column_id, "description", e.target.value)}
                                    placeholder="Add description..."
                                    className="text-xs h-7 px-2"
                                  />
                                </td>
                                <td className="px-3 py-2">
                                  <Select
                                    value={classify || undefined}
                                    onValueChange={(v) => handleColumnChange(col.column_id, "classify", v)}
                                  >
                                    <SelectTrigger className="text-xs h-7 px-2">
                                      <SelectValue placeholder="Select..." />
                                    </SelectTrigger>
                                    <SelectContent>
                                      {CLASSIFY_OPTIONS.map((opt) => (
                                        <SelectItem key={opt.value} value={opt.value}>
                                          {opt.label}
                                        </SelectItem>
                                      ))}
                                    </SelectContent>
                                  </Select>
                                </td>
                                <td className="px-3 py-2">
                                  <Select
                                    value={isMeasure ? ((edits.aggregation ?? col.aggregation) === "none" ? undefined : (edits.aggregation ?? col.aggregation)) : undefined}
                                    onValueChange={(v) => handleColumnChange(col.column_id, "aggregation", v === "none" ? "" : v)}
                                    disabled={!isMeasure}
                                  >
                                    <SelectTrigger className="text-xs h-7 px-2">
                                      <SelectValue placeholder="None" />
                                    </SelectTrigger>
                                    <SelectContent>
                                      {AGGREGATION_OPTIONS.map((opt) => (
                                        <SelectItem key={opt.value} value={opt.value}>
                                          {opt.label}
                                        </SelectItem>
                                      ))}
                                    </SelectContent>
                                  </Select>
                                </td>
                              </tr>
                            );
                          })}
                        </tbody>
                      </table>
                    </div>
                  </CardContent>
                </Card>
              )}

              {!source && (
                <div className="text-center py-8">
                  <p className="text-sm text-slate-400">No catalog data found for this source.</p>
                </div>
              )}
            </div>
          ) : (
            <div className="text-center py-8">
              <p className="text-sm text-slate-400">No catalog data found for this source.</p>
            </div>
          )}
        </div>

        {/* Footer */}
        <div className="px-5 py-3 border-t border-slate-100 flex justify-end gap-2">
          <Button variant="ghost" onClick={onClose} disabled={saving}>
            Cancel
          </Button>
          <Button onClick={handleSave} disabled={saving}>
            {saving ? (
              <>
                <Loader2 className="w-4 h-4 animate-spin mr-2" />
                Saving...
              </>
            ) : (
              <>
                <Save className="w-4 h-4 mr-2" />
                Save Data Dictionary
              </>
            )}
          </Button>
        </div>
      </div>
    </div>
  );
}

export { DataDictionaryForm };