// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
import { Spinner } from "@/components/ui/spinner";
import { useT } from "@/i18n";
import { cn } from "@/lib/utils";
import {
  AlertCircleIcon,
  CheckmarkCircle02Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { Sparkles } from "lucide-react";

const ROLE_LABELS: Record<string, string> = {
  system: "System",
  user: "User",
  assistant: "Assistant",
  human: "Human",
  gpt: "GPT",
  instruction: "Instruction",
  input: "Input",
  output: "Output",
  image: "Image",
  text: "Text",
  audio: "Audio",
  speaker_id: "Speaker ID",
};

export function HeaderRolePicker({
  currentRole,
  onRoleChange,
  availableRoles,
}: {
  currentRole: string | undefined;
  onRoleChange: (role: string | undefined) => void;
  availableRoles: readonly string[];
}) {
  return (
    <Select
      value={currentRole ?? "_none"}
      onValueChange={(v) => onRoleChange(v === "_none" ? undefined : v)}
    >
      <SelectTrigger className="h-6 w-[90px] text-ui-10 px-2 py-0 border-dashed cursor-pointer">
        <SelectValue placeholder="Role..." />
      </SelectTrigger>
      <SelectContent>
        <SelectItem value="_none" className="text-ui-11">
          None
        </SelectItem>
        {availableRoles.map((role) => (
          <SelectItem key={role} value={role} className="text-ui-11">
            {ROLE_LABELS[role] ?? role}
          </SelectItem>
        ))}
      </SelectContent>
    </Select>
  );
}

export function DatasetMappingCard({
  mapping,
  mappingOk,
  autoDetected = false,
  isVlm = false,
  isAudio = false,
  format,
  onAiAssist,
  isAiLoading = false,
  aiError,
  advisorNotification,
  advisorSystemPrompt,
}: {
  mapping: Record<string, string>;
  mappingOk: boolean;
  autoDetected?: boolean;
  isVlm?: boolean;
  isAudio?: boolean;
  format?: string;
  onAiAssist?: () => void;
  isAiLoading?: boolean;
  aiError?: string | null;
  advisorNotification?: string | null;
  advisorSystemPrompt?: string;
}) {
  const t = useT();
  const entries = Object.entries(mapping);
  const requiredLabel = isAudio
    ? t("studio.dataset.mappingRequirements.audioAndText")
    : isVlm
      ? t("studio.dataset.mappingRequirements.imageAndText")
      : format === "alpaca"
        ? t("studio.dataset.mappingRequirements.instructionAndOutput")
        : format === "sharegpt"
          ? t("studio.dataset.mappingRequirements.humanAndGpt")
          : t("studio.dataset.mappingRequirements.userAndAssistant");

  return (
    <div
      className={cn(
        "rounded-xl corner-squircle ring-1 px-5 py-4 mb-4",
        mappingOk
          ? "ring-emerald-200/70 bg-emerald-50/70 text-emerald-950 dark:ring-emerald-900/50 dark:bg-emerald-950/30 dark:text-emerald-50"
          : "ring-amber-200/70 bg-amber-50/70 text-amber-950 dark:ring-amber-900/50 dark:bg-amber-950/30 dark:text-amber-50",
      )}
    >
      <div className="flex items-start gap-3">
        <div
          className={cn(
            "rounded-xl corner-squircle p-2 shrink-0",
            mappingOk ? "bg-emerald-500/15" : "bg-amber-500/15",
          )}
        >
          <HugeiconsIcon
            icon={mappingOk ? CheckmarkCircle02Icon : AlertCircleIcon}
            className={cn(
              "size-4",
              mappingOk
                ? "text-emerald-700 dark:text-emerald-300"
                : "text-amber-700 dark:text-amber-300",
            )}
          />
        </div>
        <div className="min-w-0">
          <p className="text-sm font-semibold tracking-tight">
            {mappingOk
              ? autoDetected
                ? t("studio.dataset.mappingStatus.heuristicTitle")
                : t("studio.dataset.mappingStatus.readyTitle")
              : t("studio.dataset.mappingStatus.requiredTitle")}
          </p>
          <p
            className={cn(
              "text-xs mt-0.5",
              mappingOk
                ? "text-emerald-800/80 dark:text-emerald-200/80"
                : "text-amber-800/80 dark:text-amber-200/80",
            )}
          >
            {mappingOk
              ? autoDetected
                ? t("studio.dataset.mappingStatus.heuristicDescription")
                : t("studio.dataset.mappingStatus.readyDescription")
              : t("studio.dataset.mappingStatus.requiredDescription", {
                  required: requiredLabel,
                })}
          </p>
          {entries.length > 0 && (
            <div className="mt-3 flex flex-wrap items-center gap-2">
              {entries.map(([col, role]) => (
                <Badge
                  key={col}
                  variant="outline"
                  className="h-6 text-ui-11 bg-white/60 dark:bg-transparent"
                >
                  <span className="font-mono">{col}</span>
                  <span className="mx-1 text-muted-foreground/60">&rarr;</span>
                  <span>{ROLE_LABELS[role] ?? role}</span>
                </Badge>
              ))}
            </div>
          )}
          {!mappingOk && entries.length === 0 && (
            <p className="mt-2 text-xs text-amber-800/80 dark:text-amber-200/80">
              Use the dropdowns in the column headers to assign roles.
            </p>
          )}
          {onAiAssist && (
            <div className="mt-3 flex items-center gap-2">
              <Button
                variant="outline"
                size="sm"
                onClick={onAiAssist}
                disabled={isAiLoading}
                className="cursor-pointer bg-white/60 dark:bg-transparent"
              >
                {isAiLoading ? (
                  <>
                    <Spinner className="mr-1.5 size-3.5" />
                    Analyzing dataset...
                  </>
                ) : (
                  <>
                    <Sparkles className="mr-1.5 h-3.5 w-3.5" />
                    AI Assist
                    <Badge
                      variant="outline"
                      className="ml-1.5 text-ui-9 px-1 py-0 h-4 font-medium"
                    >
                      Beta
                    </Badge>
                  </>
                )}
              </Button>
              {aiError && (
                <p className="text-xs text-amber-700 dark:text-amber-300">
                  {aiError}
                </p>
              )}
            </div>
          )}
          {advisorNotification && (
            <div className="mt-3 rounded-lg border border-indigo-200 bg-indigo-50 px-3 py-2.5 text-xs text-indigo-700 dark:border-indigo-800 dark:bg-indigo-950 dark:text-indigo-300 space-y-2">
              <div className="flex items-start gap-2">
                <Sparkles className="size-3.5 shrink-0 mt-0.5" />
                <span>{advisorNotification}</span>
              </div>
              {advisorSystemPrompt && (
                <div className="pl-5.5 text-ui-11 font-mono text-indigo-600/80 dark:text-indigo-400/80">
                  <span className="font-sans font-medium text-indigo-500 dark:text-indigo-400">
                    System:
                  </span>{" "}
                  <span className="break-words">{advisorSystemPrompt}</span>
                </div>
              )}
            </div>
          )}
        </div>
      </div>
    </div>
  );
}

export function DatasetMappingFooter({
  mappingOk,
  startBlocked,
  stopRequested,
  startError,
  onCancel,
  onStartTraining,
}: {
  mappingOk: boolean;
  startBlocked: boolean;
  stopRequested: boolean;
  startError: string | null;
  onCancel: () => void;
  onStartTraining: () => Promise<void>;
}) {
  const t = useT();

  return (
    <div className="mt-3 flex flex-col gap-2">
      <div className="flex items-center justify-between gap-3">
        <p className="text-ui-11 text-muted-foreground/70 leading-relaxed">
          Tip: use the role dropdowns in the column headers to assign roles.
        </p>
        <div className="flex items-center gap-2">
          <Button
            variant="outline"
            size="sm"
            className="cursor-pointer"
            onClick={onCancel}
          >
            Cancel
          </Button>
          <Button
            size="sm"
            className="cursor-pointer"
            disabled={!mappingOk || startBlocked}
            onClick={() => void onStartTraining()}
          >
            {stopRequested
              ? t("studio.training.stopping")
              : startBlocked
                ? "Starting..."
                : "Continue"}
          </Button>
        </div>
      </div>

      {startError && (
        <p className="text-xs text-red-500 leading-relaxed text-center">
          {startError}
        </p>
      )}
    </div>
  );
}
