import * as React from "react";
import { Check } from "lucide-react";
import { cn } from "@/lib/utils";

/**
 * Stepper — WealthX Design System
 *
 * Built from scratch (no Base UI primitive). Supports horizontal and vertical
 * orientations with active, completed, and error step states.
 *
 * Figma: https://www.figma.com/design/9V9F0NGVsif8LGmEhVjOcT/Design-System---shadcn
 *
 * Composition:
 *   <Stepper activeStep={1}>
 *     <Step><StepItem label="..." description="..." /></Step>
 *     ...
 *   </Stepper>
 *
 * Vertical with expandable content:
 *   <Stepper activeStep={1} orientation="vertical">
 *     <Step>
 *       <div className="flex gap-4">
 *         <StepIndicator />
 *         <StepLabel description="...">Label</StepLabel>
 *       </div>
 *       <StepContent>...</StepContent>
 *     </Step>
 *   </Stepper>
 */

// ---------------------------------------------------------------------------
// Context
// ---------------------------------------------------------------------------

interface StepperContextValue {
  activeStep: number;
  orientation: "horizontal" | "vertical";
  totalSteps: number;
}

const StepperContext = React.createContext<StepperContextValue>({
  activeStep: 0,
  orientation: "horizontal",
  totalSteps: 0,
});

interface StepContextValue {
  index: number;
  isActive: boolean;
  isCompleted: boolean;
  isFirst: boolean;
  isLast: boolean;
}

const StepContext = React.createContext<StepContextValue>({
  index: 0,
  isActive: false,
  isCompleted: false,
  isFirst: true,
  isLast: false,
});

// ---------------------------------------------------------------------------
// Stepper
// ---------------------------------------------------------------------------

export interface StepperProps {
  /** 0-based index of the currently active step. */
  activeStep?: number;
  orientation?: "horizontal" | "vertical";
  className?: string;
  children: React.ReactNode;
}

function Stepper({
  activeStep = 0,
  orientation = "horizontal",
  className,
  children,
}: StepperProps) {
  const totalSteps = React.Children.count(children);

  return (
    <StepperContext.Provider value={{ activeStep, orientation, totalSteps }}>
      <div
        data-slot="stepper"
        className={cn(
          "font-sans",
          orientation === "horizontal"
            ? "flex w-full items-start"
            : "flex flex-col",
          className,
        )}
      >
        {React.Children.map(children, (step, index) =>
          React.isValidElement(step)
            ? React.cloneElement(
                step as React.ReactElement<StepProps>,
                { _index: index } as Partial<StepProps>,
              )
            : step,
        )}
      </div>
    </StepperContext.Provider>
  );
}

// ---------------------------------------------------------------------------
// Step
// ---------------------------------------------------------------------------

export interface StepProps {
  /**
   * Injected by `<Stepper>` via `React.cloneElement` — do not pass manually.
   * @internal
   */
  _index?: number;
  /** Mark this step as errored regardless of active state. */
  error?: boolean;
  className?: string;
  children: React.ReactNode;
}

function Step({ _index = 0, error, className, children }: StepProps) {
  const { activeStep, orientation, totalSteps } =
    React.useContext(StepperContext);

  const isActive = _index === activeStep;
  const isCompleted = _index < activeStep;
  const isFirst = _index === 0;
  const isLast = _index === totalSteps - 1;

  return (
    <StepContext.Provider
      value={{ index: _index, isActive, isCompleted, isFirst, isLast }}
    >
      <div
        data-slot="step"
        aria-current={isActive ? "step" : undefined}
        data-active={isActive || undefined}
        data-completed={isCompleted || undefined}
        data-error={error || undefined}
        className={cn(
          orientation === "horizontal"
            ? "flex flex-1 flex-col items-center"
            : "flex flex-row gap-3",
          className,
        )}
      >
        {children}
      </div>
    </StepContext.Provider>
  );
}

// ---------------------------------------------------------------------------
// StepIndicator — circle + connector lines
// ---------------------------------------------------------------------------

export interface StepIndicatorProps {
  /** Pass `true` to render the circle in the error/destructive state. */
  error?: boolean;
  className?: string;
}

function StepIndicator({ error, className }: StepIndicatorProps) {
  const { orientation } = React.useContext(StepperContext);
  const { index, isActive, isCompleted, isFirst, isLast } =
    React.useContext(StepContext);

  const filled = isActive || isCompleted;

  const circle = (
    <div
      className={cn(
        "relative z-10 flex size-8 shrink-0 items-center justify-center rounded-full text-sm font-medium transition-colors",
        !filled &&
          !error &&
          "border-2 border-border bg-background text-muted-foreground",
        isActive &&
          !error &&
          "border-2 border-primary bg-primary text-primary-foreground",
        isCompleted && !error && "bg-primary text-primary-foreground",
        error &&
          "border-2 border-destructive bg-destructive text-destructive-foreground",
      )}
    >
      {isCompleted && !error ? (
        <Check className="size-4" strokeWidth={2.5} />
      ) : (
        <span>{index + 1}</span>
      )}
    </div>
  );

  if (orientation === "horizontal") {
    return (
      <div
        data-slot="step-indicator"
        className={cn("flex w-full items-center", className)}
      >
        <div
          className={cn(
            "h-[2px] flex-1 transition-colors",
            isFirst
              ? "bg-transparent"
              : isActive || isCompleted
                ? "bg-primary"
                : "bg-border",
          )}
        />
        {circle}
        <div
          className={cn(
            "h-[2px] flex-1 transition-colors",
            isLast
              ? "bg-transparent"
              : isCompleted
                ? "bg-primary"
                : "bg-border",
          )}
        />
      </div>
    );
  }

  return (
    <div
      data-slot="step-indicator"
      className={cn("flex w-8 shrink-0 flex-col items-center", className)}
    >
      {circle}
      {!isLast && (
        <div
          data-slot="step-connector"
          className={cn(
            "mt-1 w-[2px] flex-1 transition-colors",
            isActive || isCompleted ? "bg-primary" : "bg-border",
          )}
        />
      )}
    </div>
  );
}

// ---------------------------------------------------------------------------
// StepLabel
// ---------------------------------------------------------------------------

export interface StepLabelProps {
  children: React.ReactNode;
  description?: string;
  /** Pass a string to show a custom error message, or `true` to show description in error colour. */
  error?: string | boolean;
  className?: string;
}

function StepLabel({
  children,
  description,
  error,
  className,
}: StepLabelProps) {
  const { orientation } = React.useContext(StepperContext);
  const { isActive, isCompleted, isLast } = React.useContext(StepContext);

  const subtext = typeof error === "string" ? error : description;

  return (
    <div
      data-slot="step-label"
      className={cn(
        "flex flex-col",
        orientation === "horizontal"
          ? "mt-2 items-center text-center"
          : cn("flex-1", !isLast && "pb-5"),
        className,
      )}
    >
      <span
        className={cn(
          orientation === "vertical" && "flex h-8 items-center",
          "text-label-small transition-colors",
          (isActive || isCompleted) && !error && "text-foreground",
          !isActive && !isCompleted && !error && "text-muted-foreground",
          error && "text-destructive",
        )}
      >
        {children}
      </span>
      {subtext && (
        <span
          className={cn(
            "mt-0.5 text-caption",
            error ? "text-destructive" : "text-muted-foreground",
          )}
        >
          {subtext}
        </span>
      )}
    </div>
  );
}

// ---------------------------------------------------------------------------
// StepContent — visible only on the active step (vertical orientation)
// ---------------------------------------------------------------------------

export interface StepContentProps {
  className?: string;
  children: React.ReactNode;
}

function StepContent({ className, children }: StepContentProps) {
  const { isActive } = React.useContext(StepContext);

  if (!isActive) return null;

  return (
    <div data-slot="step-content" className={cn("pb-4 pt-2", className)}>
      {children}
    </div>
  );
}

// ---------------------------------------------------------------------------
// StepItem — convenience wrapper for horizontal steppers
// Renders StepIndicator + StepLabel in a single element.
// ---------------------------------------------------------------------------

export interface StepItemProps {
  label: string;
  description?: string;
  error?: string | boolean;
}

function StepItem({ label, description, error }: StepItemProps) {
  return (
    <>
      <StepIndicator error={!!error} />
      <StepLabel description={description} error={error}>
        {label}
      </StepLabel>
    </>
  );
}

export { Stepper, Step, StepIndicator, StepLabel, StepContent, StepItem };