MyMidas/frontend/src/pages/auth/TwoFactorSetup.tsx

import { useState, useEffect } from "react";
import { useNavigate } from "react-router-dom";
import { useForm } from "react-hook-form";
import { zodResolver } from "@hookform/resolvers/zod";
import { z } from "zod";
import { getTotpSetup, enableTotp } from "@/api/auth";
import { useAuthStore } from "@/store/authStore";
import { useQuery, useMutation } from "@tanstack/react-query";
import { ShieldCheck, Copy, CheckCircle, Loader2 } from "lucide-react";
import { useDemoMode } from "@/hooks/useDemoMode";

const schema = z.object({ code: z.string().length(6, "6-digit code required") });
type Form = z.infer<typeof schema>;

export default function TwoFactorSetupPage() {
  const navigate = useNavigate();
  const { setTotpEnabled } = useAuthStore();
  const isDemo = useDemoMode();

  useEffect(() => {
    if (isDemo) navigate("/settings", { replace: true });
  }, [isDemo, navigate]);
  const [copied, setCopied] = useState(false);
  const [secret, setSecret] = useState<string | null>(null);
  const [error, setError] = useState<string | null>(null);

  const { data, isLoading } = useQuery({
    queryKey: ["totp-setup"],
    queryFn: async () => {
      const res = await getTotpSetup();
      setSecret(res.secret);
      return res;
    },
    enabled: !isDemo,
  });

  const { register, handleSubmit, formState } = useForm<Form>({
    resolver: zodResolver(schema),
  });

  const enableMutation = useMutation({
    mutationFn: ({ code }: { code: string }) => enableTotp(secret!, code),
    onSuccess: () => {
      setTotpEnabled(true);
      navigate("/settings");
    },
    onError: () => setError("Invalid code — try again"),
  });

  function copySecret() {
    if (data?.secret) {
      navigator.clipboard.writeText(data.secret);
      setCopied(true);
      setTimeout(() => setCopied(false), 2000);
    }
  }

  if (isLoading) {
    return (
      <div className="flex items-center justify-center h-64">
        <Loader2 className="w-6 h-6 animate-spin text-muted-foreground" />
      </div>
    );
  }

  return (
    <div className="max-w-md mx-auto mt-8">
      <div className="bg-card border border-border rounded-xl p-8">
        <div className="flex items-center gap-2 mb-6">
          <ShieldCheck className="w-6 h-6 text-primary" />
          <h1 className="text-xl font-semibold">Set up two-factor authentication</h1>
        </div>

        {/* QR code */}
        <div className="flex justify-center mb-6">
          <div className="p-3 bg-white rounded-lg">
            {data?.qr_code_png_b64 && (
              <img
                src={`data:image/png;base64,${data.qr_code_png_b64}`}
                alt="TOTP QR code"
                className="w-48 h-48"
              />
            )}
          </div>
        </div>

        <p className="text-sm text-muted-foreground mb-2 text-center">
          Scan with your authenticator app (Authy, Google Authenticator, etc.)
        </p>

        {/* Manual secret */}
        <div className="mb-6">
          <p className="text-xs text-muted-foreground mb-1">Or enter the secret manually:</p>
          <div className="flex items-center gap-2">
            <code className="flex-1 text-xs bg-secondary px-3 py-2 rounded font-mono break-all">
              {data?.secret}
            </code>
            <button onClick={copySecret} className="text-muted-foreground hover:text-foreground">
              {copied ? <CheckCircle className="w-4 h-4 text-success" /> : <Copy className="w-4 h-4" />}
            </button>
          </div>
        </div>

        {/* Backup codes */}
        {data?.backup_codes && (
          <div className="mb-6">
            <p className="text-xs font-medium text-warning mb-2">
              Save these backup codes — you can only see them once:
            </p>
            <div className="grid grid-cols-2 gap-1">
              {data.backup_codes.map((code) => (
                <code key={code} className="text-xs bg-secondary px-2 py-1 rounded font-mono text-center">
                  {code}
                </code>
              ))}
            </div>
          </div>
        )}

        {/* Verify */}
        <form
          onSubmit={handleSubmit(({ code }) => enableMutation.mutate({ code }))}
          className="space-y-3"
        >
          <div>
            <label className="text-sm font-medium block mb-1.5">
              Enter code to confirm setup
            </label>
            <input
              {...register("code")}
              type="text"
              inputMode="numeric"
              maxLength={6}
              className="w-full rounded-md border border-input bg-background px-3 py-2 text-sm text-center tracking-widest font-mono focus:outline-none focus:ring-2 focus:ring-ring"
              placeholder="000000"
            />
            {formState.errors.code && (
              <p className="text-destructive text-xs mt-1">{formState.errors.code.message}</p>
            )}
          </div>

          {error && (
            <p className="text-destructive text-sm bg-destructive/10 rounded px-3 py-2">{error}</p>
          )}

          <button
            type="submit"
            disabled={enableMutation.isPending}
            className="w-full flex items-center justify-center gap-2 bg-primary text-primary-foreground rounded-md py-2.5 text-sm font-medium hover:bg-primary/90 disabled:opacity-50 transition-colors"
          >
            {enableMutation.isPending && <Loader2 className="w-4 h-4 animate-spin" />}
            Enable 2FA
          </button>

          <button
            type="button"
            onClick={() => navigate("/")}
            className="w-full text-sm text-muted-foreground hover:text-foreground transition-colors"
          >
            Skip for now
          </button>
        </form>
      </div>
    </div>
  );
}