encoach_frontend_new_v2/frontend/src/components/ai/AiGeneratorModal.tsx

import { useState } from "react";
import { useMutation } from "@tanstack/react-query";
import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogTrigger } from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Sparkles, Loader2, Trash2 } from "lucide-react";
import type { ExamModule } from "@/types";
import { generationService } from "@/services/generation.service";
import { useToast } from "@/hooks/use-toast";

type ExerciseRow = { title?: string; description?: string; marks?: number };

function exerciseLabel(ex: unknown, index: number): ExerciseRow {
  if (ex && typeof ex === "object") {
    const o = ex as Record<string, unknown>;
    return {
      title: typeof o.title === "string" ? o.title : `Exercise ${index + 1}`,
      description: typeof o.description === "string" ? o.description : undefined,
      marks: typeof o.marks === "number" ? o.marks : typeof o.marks === "string" ? Number(o.marks) : undefined,
    };
  }
  return { title: String(ex) };
}

export default function AiGeneratorModal() {
  const [open, setOpen] = useState(false);
  const [moduleType, setModuleType] = useState<ExamModule>("reading");
  const [difficulty, setDifficulty] = useState("b2");
  const [count, setCount] = useState(5);
  const [topic, setTopic] = useState("");
  const [localExercises, setLocalExercises] = useState<unknown[] | null>(null);
  const { toast } = useToast();

  const generateMutation = useMutation({
    mutationFn: () =>
      generationService.generate(moduleType, {
        title: topic.trim() || `${moduleType} practice set`,
        difficulty,
        count,
      }),
    onSuccess: (res: Record<string, unknown>) => {
      const items = Array.isArray(res.questions) ? res.questions : Array.isArray(res.exercises) ? res.exercises : [];
      setLocalExercises(items);
    },
    onError: (err: Error) => {
      toast({
        variant: "destructive",
        title: "Generation failed",
        description: err.message || "Could not generate exercises.",
      });
    },
  });

  const handleGenerate = () => {
    setLocalExercises(null);
    generateMutation.mutate();
  };

  const saveMutation = useMutation({
    mutationFn: () => generationService.saveGenerated(moduleType, localExercises ?? []),
    onSuccess: (res) => {
      toast({ title: "Saved", description: `${res.saved} assignments saved successfully.` });
      setOpen(false);
    },
    onError: (err: Error) => {
      toast({
        variant: "destructive",
        title: "Save failed",
        description: err.message || "Could not save generated assignments.",
      });
    },
  });

  const generated = localExercises;

  const handleRemove = (index: number) => {
    setLocalExercises((prev) => (prev ? prev.filter((_, i) => i !== index) : null));
  };

  return (
    <Dialog
      open={open}
      onOpenChange={(v) => {
        setOpen(v);
        if (!v) {
          setLocalExercises(null);
          generateMutation.reset();
        }
      }}
    >
      <DialogTrigger asChild>
        <Button variant="outline" size="sm">
          <Sparkles className="h-3.5 w-3.5 mr-1 text-primary" /> Generate with AI
        </Button>
      </DialogTrigger>
      <DialogContent className="max-w-2xl max-h-[80vh] overflow-y-auto">
        <DialogHeader>
          <DialogTitle className="flex items-center gap-2">
            <Sparkles className="h-4 w-4 text-primary" /> AI Assignment Generator
          </DialogTitle>
        </DialogHeader>
        
        {!generated ? (
          <div className="space-y-4">
            <div className="space-y-2">
              <Label>Module Type</Label>
              <Select
                value={moduleType}
                onValueChange={(v) => setModuleType(v as ExamModule)}
              >
                <SelectTrigger>
                  <SelectValue />
                </SelectTrigger>
                <SelectContent>
                  <SelectItem value="reading">Reading</SelectItem>
                  <SelectItem value="listening">Listening</SelectItem>
                  <SelectItem value="writing">Writing</SelectItem>
                  <SelectItem value="speaking">Speaking</SelectItem>
                </SelectContent>
              </Select>
            </div>
            <div className="grid grid-cols-2 gap-4">
              <div className="space-y-2">
                <Label>Difficulty</Label>
                <Select value={difficulty} onValueChange={setDifficulty}>
                  <SelectTrigger>
                    <SelectValue />
                  </SelectTrigger>
                  <SelectContent>
                    {["A1", "A2", "B1", "B2", "C1", "C2"].map((l) => (
                      <SelectItem key={l} value={l.toLowerCase()}>
                        {l}
                      </SelectItem>
                    ))}
                  </SelectContent>
                </Select>
              </div>
              <div className="space-y-2">
                <Label>Count</Label>
                <Input
                  type="number"
                  value={count}
                  min={1}
                  max={10}
                  onChange={(e) => setCount(Number(e.target.value) || 1)}
                />
              </div>
            </div>
            <div className="space-y-2">
              <Label>Topic / Focus Area</Label>
              <Input
                placeholder="e.g. Academic reading comprehension"
                value={topic}
                onChange={(e) => setTopic(e.target.value)}
              />
            </div>
            <Button
              className="w-full"
              onClick={handleGenerate}
              disabled={generateMutation.isPending}
            >
              {generateMutation.isPending ? (
                <>
                  <Loader2 className="h-4 w-4 mr-2 animate-spin" /> Generating...
                </>
              ) : (
                <>
                  <Sparkles className="h-4 w-4 mr-2" /> Generate Assignments
                </>
              )}
            </Button>
          </div>
        ) : (
          <div className="space-y-3">
            <p className="text-sm text-muted-foreground">
              AI generated {generated.length} assignments. Edit or remove before saving.
            </p>
            {generated.map((item, i) => {
              const row = exerciseLabel(item, i);
              return (
                <div key={i} className="rounded-lg border p-3 space-y-1">
                  <div className="flex items-start justify-between">
                    <div>
                      <p className="text-sm font-medium">{row.title}</p>
                      {row.description && (
                        <p className="text-xs text-muted-foreground mt-1">{row.description}</p>
                      )}
                      {row.marks !== undefined && !Number.isNaN(row.marks) && (
                        <p className="text-xs font-semibold mt-1">{row.marks} marks</p>
                      )}
                    </div>
                    <Button
                      variant="ghost"
                      size="icon"
                      className="h-7 w-7 shrink-0"
                      onClick={() => handleRemove(i)}
                    >
                      <Trash2 className="h-3.5 w-3.5 text-destructive" />
                    </Button>
                  </div>
                </div>
              );
            })}
            <div className="flex gap-2">
              <Button
                className="flex-1"
                onClick={() => saveMutation.mutate()}
                disabled={saveMutation.isPending || !generated?.length}
              >
                {saveMutation.isPending ? (
                  <><Loader2 className="h-4 w-4 mr-2 animate-spin" /> Saving...</>
                ) : (
                  "Save All"
                )}
              </Button>
              <Button
                variant="outline"
                onClick={() => {
                  setLocalExercises(null);
                  generateMutation.reset();
                }}
              >
                Regenerate
              </Button>
            </div>
          </div>
        )}
      </DialogContent>
    </Dialog>
  );
}