import { useState, useMemo } from "react";
import { apiFetch } from '@/utils/api';
import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import { Button } from "@/components/ui/button";
import {
  Table,
  TableBody,
  TableCell,
  TableHead,
  TableHeader,
  TableRow,
} from "@/components/ui/table";
import { Input } from "@/components/ui/input";
import { Textarea } from "@/components/ui/textarea";
import { Select, SelectContent, SelectGroup, SelectItem, SelectLabel, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Label } from "@/components/ui/label";
import { Badge } from "@/components/ui/badge";
import { ArrowUpDown, Pencil, Trash2, PlusCircle } from "lucide-react";
import config from "@/config";
import {
  useReactTable,
  getCoreRowModel,
  getSortedRowModel,
  SortingState,
  flexRender,
  type ColumnDef,
} from "@tanstack/react-table";
import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogFooter,
  DialogHeader,
  DialogTitle,
} from "@/components/ui/dialog";
import {
  AlertDialog,
  AlertDialogAction,
  AlertDialogCancel,
  AlertDialogContent,
  AlertDialogDescription,
  AlertDialogFooter,
  AlertDialogHeader,
  AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { toast } from "sonner";

interface FieldOption {
  label: string;
  value: string;
}

// Form uses task + role, which gets composed into prompt_type on submit
interface PromptFormData {
  id?: number;
  prompt_text: string;
  task: string;
  role: "system" | "general" | "company_specific";
  company: string | null;
}

interface AiPrompt {
  id: number;
  prompt_text: string;
  prompt_type: string;
  company: string | null;
  is_singleton: boolean;
  created_at: string;
  updated_at: string;
}

interface FieldOptions {
  companies: FieldOption[];
}

// Predefined tasks (can also enter custom)
const PREDEFINED_TASKS = [
  { value: "name_validation", label: "Name Validation", description: "Inline validation of product names (Groq)" },
  { value: "description_validation", label: "Description Validation", description: "Inline validation of product descriptions (Groq)" },
  { value: "sanity_check", label: "Sanity Check", description: "Batch product consistency review (Groq)" },
  { value: "bulk_validation", label: "Bulk Validation", description: "Full product validation during import (GPT-5)" },
];

// Role options
const ROLES = [
  { value: "system", label: "System", description: "Initial instructions that set the AI's behavior" },
  { value: "general", label: "General", description: "Rules that apply to all products" },
  { value: "company_specific", label: "Company-Specific", description: "Rules unique to a specific company" },
];

// Parse prompt_type into task and role
function parsePromptType(promptType: string): { task: string; role: "system" | "general" | "company_specific" } {
  if (promptType.endsWith("_company_specific")) {
    return { task: promptType.replace("_company_specific", ""), role: "company_specific" };
  }
  if (promptType.endsWith("_general")) {
    return { task: promptType.replace("_general", ""), role: "general" };
  }
  if (promptType.endsWith("_system")) {
    return { task: promptType.replace("_system", ""), role: "system" };
  }
  // Fallback - assume it's a system prompt
  return { task: promptType, role: "system" };
}

// Compose task + role into prompt_type
function composePromptType(task: string, role: string): string {
  return `${task}_${role}`;
}

// Get human-readable task name
function getTaskDisplayName(task: string): string {
  const predefined = PREDEFINED_TASKS.find(t => t.value === task);
  if (predefined) return predefined.label;
  // Format custom task nicely
  return task
    .split('_')
    .map(word => word.charAt(0).toUpperCase() + word.slice(1))
    .join(' ');
}

// Get human-readable role name
function getRoleDisplayName(role: string): string {
  const roleInfo = ROLES.find(r => r.value === role);
  return roleInfo?.label || role;
}

export function PromptManagement() {
  const [isFormOpen, setIsFormOpen] = useState(false);
  const [isDeleteOpen, setIsDeleteOpen] = useState(false);
  const [promptToDelete, setPromptToDelete] = useState<AiPrompt | null>(null);
  const [editingPrompt, setEditingPrompt] = useState<AiPrompt | null>(null);
  const [sorting, setSorting] = useState<SortingState>([
    { id: "prompt_type", desc: false },
    { id: "company", desc: false }
  ]);
  const [searchQuery, setSearchQuery] = useState("");
  const [useCustomTask, setUseCustomTask] = useState(false);
  const [customTask, setCustomTask] = useState("");
  const [formData, setFormData] = useState<PromptFormData>({
    prompt_text: "",
    task: "",
    role: "system",
    company: null,
  });

  const queryClient = useQueryClient();

  const { data: prompts, isLoading } = useQuery<AiPrompt[]>({
    queryKey: ["ai-prompts"],
    queryFn: async () => {
      const response = await apiFetch(`${config.apiUrl}/ai-prompts`);
      if (!response.ok) {
        throw new Error("Failed to fetch AI prompts");
      }
      return response.json();
    },
  });

  const { data: fieldOptions } = useQuery<FieldOptions>({
    queryKey: ["fieldOptions"],
    queryFn: async () => {
      const response = await apiFetch(`${config.apiUrl}/import/field-options`);
      if (!response.ok) {
        throw new Error("Failed to fetch field options");
      }
      return response.json();
    },
  });

  // Track which prompts exist for disabling options
  const existingPrompts = useMemo(() => {
    if (!prompts) return { system: new Set<string>(), general: new Set<string>(), companySpecific: new Map<string, Set<string>>() };

    const system = new Set<string>();
    const general = new Set<string>();
    const companySpecific = new Map<string, Set<string>>();

    prompts.forEach(p => {
      const { task, role } = parsePromptType(p.prompt_type);
      if (role === "system") {
        system.add(task);
      } else if (role === "general") {
        general.add(task);
      } else if (role === "company_specific" && p.company) {
        if (!companySpecific.has(task)) {
          companySpecific.set(task, new Set());
        }
        companySpecific.get(task)!.add(p.company);
      }
    });

    return { system, general, companySpecific };
  }, [prompts]);

  const createMutation = useMutation({
    mutationFn: async (data: { prompt_text: string; prompt_type: string; company: string | null; is_singleton: boolean }) => {
      const response = await apiFetch(`${config.apiUrl}/ai-prompts`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(data),
      });
      if (!response.ok) {
        const error = await response.json();
        throw new Error(error.message || error.error || "Failed to create prompt");
      }
      return response.json();
    },
    onSuccess: (newPrompt) => {
      queryClient.setQueryData<AiPrompt[]>(["ai-prompts"], (old) => {
        if (!old) return [newPrompt];
        return [...old, newPrompt];
      });
      toast.success("Prompt created successfully");
      resetForm();
    },
    onError: (error) => {
      toast.error(error instanceof Error ? error.message : "Failed to create prompt");
    },
  });

  const updateMutation = useMutation({
    mutationFn: async (data: { id: number; prompt_text: string; prompt_type: string; company: string | null; is_singleton: boolean }) => {
      const response = await apiFetch(`${config.apiUrl}/ai-prompts/${data.id}`, {
        method: "PUT",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(data),
      });
      if (!response.ok) {
        const error = await response.json();
        throw new Error(error.message || error.error || "Failed to update prompt");
      }
      return response.json();
    },
    onSuccess: (updatedPrompt) => {
      queryClient.setQueryData<AiPrompt[]>(["ai-prompts"], (old) => {
        if (!old) return [updatedPrompt];
        return old.map((prompt) =>
          prompt.id === updatedPrompt.id ? updatedPrompt : prompt
        );
      });
      toast.success("Prompt updated successfully");
      resetForm();
    },
    onError: (error) => {
      toast.error(error instanceof Error ? error.message : "Failed to update prompt");
    },
  });

  const deleteMutation = useMutation({
    mutationFn: async (id: number) => {
      const response = await apiFetch(`${config.apiUrl}/ai-prompts/${id}`, {
        method: "DELETE",
      });
      if (!response.ok) {
        throw new Error("Failed to delete prompt");
      }
      return id;
    },
    onSuccess: (deletedId) => {
      queryClient.setQueryData<AiPrompt[]>(["ai-prompts"], (old) => {
        if (!old) return [];
        return old.filter((prompt) => prompt.id !== deletedId);
      });
      toast.success("Prompt deleted successfully");
    },
    onError: (error) => {
      toast.error(error instanceof Error ? error.message : "Failed to delete prompt");
    },
  });

  const handleEdit = (prompt: AiPrompt) => {
    setEditingPrompt(prompt);
    const { task, role } = parsePromptType(prompt.prompt_type);
    const isPredefinedTask = PREDEFINED_TASKS.some(t => t.value === task);
    setUseCustomTask(!isPredefinedTask);
    setCustomTask(isPredefinedTask ? "" : task);
    setFormData({
      id: prompt.id,
      prompt_text: prompt.prompt_text,
      task: task,
      role: role,
      company: prompt.company,
    });
    setIsFormOpen(true);
  };

  const handleDeleteClick = (prompt: AiPrompt) => {
    setPromptToDelete(prompt);
    setIsDeleteOpen(true);
  };

  const handleDeleteConfirm = () => {
    if (promptToDelete) {
      deleteMutation.mutate(promptToDelete.id);
      setIsDeleteOpen(false);
      setPromptToDelete(null);
    }
  };

  const handleSubmit = (e: React.FormEvent) => {
    e.preventDefault();

    const actualTask = useCustomTask ? customTask.trim() : formData.task;

    if (!actualTask) {
      toast.error("Please select or enter a task");
      return;
    }

    if (!formData.role) {
      toast.error("Please select a role");
      return;
    }

    if (formData.role === "company_specific" && !formData.company) {
      toast.error("Please select a company");
      return;
    }

    const promptType = composePromptType(actualTask, formData.role);
    const submitData = {
      id: formData.id,
      prompt_text: formData.prompt_text,
      prompt_type: promptType,
      company: formData.role === "company_specific" ? formData.company : null,
      is_singleton: true, // Always singleton
    };

    if (editingPrompt) {
      updateMutation.mutate(submitData as { id: number; prompt_text: string; prompt_type: string; company: string | null; is_singleton: boolean });
    } else {
      createMutation.mutate(submitData);
    }
  };

  const resetForm = () => {
    setFormData({
      prompt_text: "",
      task: "",
      role: "system",
      company: null,
    });
    setEditingPrompt(null);
    setUseCustomTask(false);
    setCustomTask("");
    setIsFormOpen(false);
  };

  const handleCreateClick = () => {
    resetForm();
    // Find first task + role combo that doesn't exist
    for (const task of PREDEFINED_TASKS) {
      if (!existingPrompts.system.has(task.value)) {
        setFormData(prev => ({ ...prev, task: task.value, role: "system" }));
        setIsFormOpen(true);
        return;
      }
      if (!existingPrompts.general.has(task.value)) {
        setFormData(prev => ({ ...prev, task: task.value, role: "general" }));
        setIsFormOpen(true);
        return;
      }
    }
    // All base prompts exist, default to first task with company-specific
    setFormData(prev => ({ ...prev, task: PREDEFINED_TASKS[0].value, role: "company_specific" }));
    setIsFormOpen(true);
  };

  const handleTaskChange = (value: string) => {
    if (value === "__custom__") {
      setUseCustomTask(true);
      setFormData(prev => ({ ...prev, task: "" }));
    } else {
      setUseCustomTask(false);
      setCustomTask("");
      setFormData(prev => ({ ...prev, task: value }));
    }
  };

  // Get the effective task for checking what exists
  const effectiveTask = useCustomTask ? customTask.trim() : formData.task;

  // Check if current selection would conflict
  const wouldConflict = useMemo(() => {
    if (!effectiveTask) return false;

    // If editing the same prompt, no conflict
    if (editingPrompt) {
      const { task: editTask, role: editRole } = parsePromptType(editingPrompt.prompt_type);
      if (editTask === effectiveTask && editRole === formData.role) {
        if (formData.role !== "company_specific") return false;
        if (editingPrompt.company === formData.company) return false;
      }
    }

    if (formData.role === "system") {
      return existingPrompts.system.has(effectiveTask);
    }
    if (formData.role === "general") {
      return existingPrompts.general.has(effectiveTask);
    }
    if (formData.role === "company_specific" && formData.company) {
      const taskCompanies = existingPrompts.companySpecific.get(effectiveTask);
      return taskCompanies?.has(formData.company) || false;
    }
    return false;
  }, [effectiveTask, formData.role, formData.company, existingPrompts, editingPrompt]);

  const columns = useMemo<ColumnDef<AiPrompt>[]>(() => [
    {
      accessorKey: "prompt_type",
      header: ({ column }) => (
        <Button
          variant="ghost"
          onClick={() => column.toggleSorting(column.getIsSorted() === "asc")}
        >
          Task
          <ArrowUpDown className="ml-2 h-4 w-4" />
        </Button>
      ),
      cell: ({ row }) => {
        const { task } = parsePromptType(row.getValue("prompt_type") as string);
        return (
          <span className="font-medium">{getTaskDisplayName(task)}</span>
        );
      },
    },
    {
      id: "role",
      header: "Role",
      cell: ({ row }) => {
        const { role } = parsePromptType(row.getValue("prompt_type") as string);
        const variant = role === "system" ? "default" : role === "general" ? "secondary" : "outline";
        return <Badge variant={variant}>{getRoleDisplayName(role)}</Badge>;
      },
    },
    {
      accessorFn: (row) => row.prompt_text.length,
      id: "length",
      header: ({ column }) => (
        <Button
          variant="ghost"
          onClick={() => column.toggleSorting(column.getIsSorted() === "asc")}
        >
          Length
          <ArrowUpDown className="ml-2 h-4 w-4" />
        </Button>
      ),
      cell: ({ getValue }) => {
        const length = getValue() as number;
        return <span>{length.toLocaleString()}</span>;
      },
    },
    {
      accessorKey: "company",
      header: ({ column }) => (
        <Button
          variant="ghost"
          onClick={() => column.toggleSorting(column.getIsSorted() === "asc")}
        >
          Company
          <ArrowUpDown className="ml-2 h-4 w-4" />
        </Button>
      ),
      cell: ({ row }) => {
        const companyId = row.getValue("company");
        if (!companyId) return <span className="text-muted-foreground">—</span>;
        return fieldOptions?.companies.find(c => c.value === companyId)?.label || companyId;
      },
    },
    {
      accessorKey: "updated_at",
      header: ({ column }) => (
        <Button
          variant="ghost"
          onClick={() => column.toggleSorting(column.getIsSorted() === "asc")}
        >
          Last Updated
          <ArrowUpDown className="ml-2 h-4 w-4" />
        </Button>
      ),
      cell: ({ row }) => new Date(row.getValue("updated_at")).toLocaleDateString(),
    },
    {
      id: "actions",
      cell: ({ row }) => (
        <div className="flex gap-0 justify-end">
          <Button variant="ghost" onClick={() => handleEdit(row.original)}>
            <Pencil className="h-4 w-4" />
            Edit
          </Button>
          <Button
            variant="ghost"
            className="text-destructive hover:text-destructive"
            onClick={() => handleDeleteClick(row.original)}
          >
            <Trash2 className="h-4 w-4" />
            Delete
          </Button>
        </div>
      ),
    },
  ], [fieldOptions]);

  const filteredData = useMemo(() => {
    if (!prompts) return [];
    return prompts.filter((prompt) => {
      const searchString = searchQuery.toLowerCase();
      const { task, role } = parsePromptType(prompt.prompt_type);
      const taskName = getTaskDisplayName(task).toLowerCase();
      const roleName = getRoleDisplayName(role).toLowerCase();
      return (
        taskName.includes(searchString) ||
        roleName.includes(searchString) ||
        prompt.prompt_type.toLowerCase().includes(searchString) ||
        (prompt.company && prompt.company.toLowerCase().includes(searchString))
      );
    });
  }, [prompts, searchQuery]);

  const table = useReactTable({
    data: filteredData,
    columns,
    state: { sorting },
    onSortingChange: setSorting,
    getSortedRowModel: getSortedRowModel(),
    getCoreRowModel: getCoreRowModel(),
  });

  return (
    <div className="space-y-6">
      <div className="flex items-center justify-between">
        <h2 className="text-2xl font-bold">AI Validation Prompts</h2>
        <Button onClick={handleCreateClick}>
          <PlusCircle className="mr-2 h-4 w-4" />
          Create New Prompt
        </Button>
      </div>

      <div className="flex items-center gap-4">
        <Input
          placeholder="Search prompts..."
          value={searchQuery}
          onChange={(e) => setSearchQuery(e.target.value)}
          className="max-w-sm"
        />
      </div>

      {isLoading ? (
        <div>Loading prompts...</div>
      ) : (
        <div className="border rounded-lg">
          <Table>
            <TableHeader className="bg-muted">
              {table.getHeaderGroups().map((headerGroup) => (
                <TableRow key={headerGroup.id} className="">
                  {headerGroup.headers.map((header) => (
                    <TableHead key={header.id} className="whitespace-nowrap">
                      {header.isPlaceholder
                        ? null
                        : flexRender(header.column.columnDef.header, header.getContext())}
                    </TableHead>
                  ))}
                </TableRow>
              ))}
            </TableHeader>
            <TableBody>
              {table.getRowModel().rows?.length ? (
                table.getRowModel().rows.map((row) => (
                  <TableRow key={row.id} className="hover:bg-gray-100">
                    {row.getVisibleCells().map((cell) => (
                      <TableCell key={cell.id} className="pl-3 whitespace-nowrap">
                        {flexRender(cell.column.columnDef.cell, cell.getContext())}
                      </TableCell>
                    ))}
                  </TableRow>
                ))
              ) : (
                <TableRow>
                  <TableCell colSpan={columns.length} className="text-center">
                    No prompts found
                  </TableCell>
                </TableRow>
              )}
            </TableBody>
          </Table>
        </div>
      )}

      {/* Prompt Form Dialog */}
      <Dialog open={isFormOpen} onOpenChange={setIsFormOpen}>
        <DialogContent className="max-w-3xl">
          <DialogHeader>
            <DialogTitle>{editingPrompt ? "Edit Prompt" : "Create New Prompt"}</DialogTitle>
            <DialogDescription>
              {editingPrompt
                ? "Update this AI validation prompt."
                : "Create a new AI validation prompt. Select a task and role, then enter the prompt text."}
            </DialogDescription>
          </DialogHeader>

          <form onSubmit={handleSubmit}>
            <div className="grid gap-4 py-4">
              {/* Task Selector */}
              <div className="grid gap-2">
                <Label>Task</Label>
                {!useCustomTask ? (
                  <Select value={formData.task} onValueChange={handleTaskChange}>
                    <SelectTrigger>
                      <SelectValue placeholder="Select task" />
                    </SelectTrigger>
                    <SelectContent>
                      <SelectGroup>
                        <SelectLabel className="text-xs text-muted-foreground">Predefined Tasks</SelectLabel>
                        {PREDEFINED_TASKS.map((task) => (
                          <SelectItem key={task.value} value={task.value}>
                            <span className="flex items-center gap-2">
                              <span>{task.label}</span>
                              <span className="text-xs text-muted-foreground">{task.description}</span>
                            </span>
                          </SelectItem>
                        ))}
                      </SelectGroup>
                      <SelectGroup>
                        <SelectLabel className="text-xs text-muted-foreground">Custom</SelectLabel>
                        <SelectItem value="__custom__">Custom Task...</SelectItem>
                      </SelectGroup>
                    </SelectContent>
                  </Select>
                ) : (
                  <div className="flex gap-2">
                    <Input
                      value={customTask}
                      onChange={(e) => setCustomTask(e.target.value.toLowerCase().replace(/\s+/g, '_'))}
                      placeholder="e.g., my_custom_task"
                      className="font-mono"
                    />
                    <Button
                      type="button"
                      variant="outline"
                      onClick={() => {
                        setUseCustomTask(false);
                        setCustomTask("");
                      }}
                    >
                      Cancel
                    </Button>
                  </div>
                )}
              </div>

              {/* Role Selector */}
              <div className="grid gap-2">
                <Label>Role</Label>
                <Select
                  value={formData.role}
                  onValueChange={(value: "system" | "general" | "company_specific") =>
                    setFormData(prev => ({ ...prev, role: value, company: value !== "company_specific" ? null : prev.company }))
                  }
                >
                  <SelectTrigger>
                    <SelectValue placeholder="Select role" />
                  </SelectTrigger>
                  <SelectContent>
                    {ROLES.map((role) => {
                      // Check if this role is already taken for the current task
                      let isDisabled = false;
                      if (effectiveTask) {
                        if (role.value === "system") {
                          isDisabled = existingPrompts.system.has(effectiveTask);
                        } else if (role.value === "general") {
                          isDisabled = existingPrompts.general.has(effectiveTask);
                        }
                        // Company-specific is never disabled at the role level
                      }
                      // Allow if editing the same prompt
                      if (editingPrompt && effectiveTask) {
                        const { task: editTask, role: editRole } = parsePromptType(editingPrompt.prompt_type);
                        if (editTask === effectiveTask && editRole === role.value) {
                          isDisabled = false;
                        }
                      }
                      return (
                        <SelectItem key={role.value} value={role.value} disabled={isDisabled}>
                          <span className="flex items-center gap-2">
                            {role.label}
                            {isDisabled && <Badge variant="secondary" className="text-xs">exists</Badge>}
                          </span>
                        </SelectItem>
                      );
                    })}
                  </SelectContent>
                </Select>
                <p className="text-xs text-muted-foreground">
                  {ROLES.find(r => r.value === formData.role)?.description}
                </p>
              </div>

              {/* Company Selector (only for company_specific role) */}
              {formData.role === "company_specific" && (
                <div className="grid gap-2">
                  <Label>Company</Label>
                  <Select
                    value={formData.company || ''}
                    onValueChange={(value) => setFormData(prev => ({ ...prev, company: value }))}
                  >
                    <SelectTrigger>
                      <SelectValue placeholder="Select company" />
                    </SelectTrigger>
                    <SelectContent>
                      {fieldOptions?.companies.map((company) => {
                        const taskCompanies = existingPrompts.companySpecific.get(effectiveTask);
                        const isExisting = taskCompanies?.has(company.value);
                        const isCurrentEdit = editingPrompt?.company === company.value;
                        return (
                          <SelectItem
                            key={company.value}
                            value={company.value}
                            disabled={isExisting && !isCurrentEdit}
                          >
                            <span className="flex items-center gap-2">
                              {company.label}
                              {isExisting && !isCurrentEdit && (
                                <Badge variant="secondary" className="text-xs">exists</Badge>
                              )}
                            </span>
                          </SelectItem>
                        );
                      })}
                    </SelectContent>
                  </Select>
                </div>
              )}

              {/* Conflict Warning */}
              {wouldConflict && (
                <p className="text-sm text-destructive">
                  A prompt for this task + role combination already exists.
                </p>
              )}

              {/* Prompt Text */}
              <div className="grid gap-2">
                <Label htmlFor="prompt_text">Prompt Text</Label>
                <Textarea
                  id="prompt_text"
                  value={formData.prompt_text}
                  onChange={(e) => setFormData(prev => ({ ...prev, prompt_text: e.target.value }))}
                  placeholder="Enter your prompt text..."
                  className="h-80 font-mono text-sm"
                  required
                />
              </div>
            </div>

            <DialogFooter>
              <Button type="button" variant="outline" onClick={() => { resetForm(); setIsFormOpen(false); }}>
                Cancel
              </Button>
              <Button
                type="submit"
                disabled={wouldConflict || !effectiveTask || !formData.role || (formData.role === "company_specific" && !formData.company)}
              >
                {editingPrompt ? "Update" : "Create"} Prompt
              </Button>
            </DialogFooter>
          </form>
        </DialogContent>
      </Dialog>

      {/* Delete Confirmation Dialog */}
      <AlertDialog open={isDeleteOpen} onOpenChange={setIsDeleteOpen}>
        <AlertDialogContent>
          <AlertDialogHeader>
            <AlertDialogTitle>Delete Prompt</AlertDialogTitle>
            <AlertDialogDescription>
              Are you sure you want to delete this prompt? This action cannot be undone.
            </AlertDialogDescription>
          </AlertDialogHeader>
          <AlertDialogFooter>
            <AlertDialogCancel onClick={() => { setIsDeleteOpen(false); setPromptToDelete(null); }}>
              Cancel
            </AlertDialogCancel>
            <AlertDialogAction onClick={handleDeleteConfirm}>
              Delete
            </AlertDialogAction>
          </AlertDialogFooter>
        </AlertDialogContent>
      </AlertDialog>
    </div>
  );
}
