semblance-dev/src/components/persona/PersonaModificationModal.tsx

import React, { useState } from 'react';
import { zodResolver } from "@hookform/resolvers/zod";
import { useForm } from "react-hook-form";
import { z } from "zod";
import { Loader2, Bot, Wand2 } from 'lucide-react';
import { Button } from "@/components/ui/button";
import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogHeader,
  DialogTitle,
  DialogFooter,
} from "@/components/ui/dialog";
import {
  Form,
  FormControl,
  FormDescription,
  FormField,
  FormItem,
  FormLabel,
  FormMessage,
} from "@/components/ui/form";
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
import { Textarea } from "@/components/ui/textarea";
import { personasApi } from '@/lib/api';
import { toastService } from '@/lib/toast';
import { useCancellableGeneration } from '@/hooks/useCancellableGeneration';
import { getSocket } from '@/services/websocketServiceNew';
import GenerationProgressBar from '@/components/ui/GenerationProgressBar';
import { Persona } from '@/types/persona';

const modificationFormSchema = z.object({
  modificationPrompt: z.string().min(10, {
    message: "Modification prompt must be at least 10 characters long.",
  }),
  llm_model: z.string().min(1, {
    message: "Please select an AI model.",
  }),
  reasoning_effort: z.string().optional(),
  verbosity: z.string().optional(),
});

type ModificationFormData = z.infer<typeof modificationFormSchema>;

interface PersonaModificationModalProps {
  persona: Persona;
  isOpen: boolean;
  onClose: () => void;
  onPersonaPreview: (modifiedPersona: Persona) => void;
}

export default function PersonaModificationModal({ 
  persona, 
  isOpen, 
  onClose, 
  onPersonaPreview 
}: PersonaModificationModalProps) {
  // Cancellable generation for persona modification
  const socket = getSocket();
  const [modificationState, modificationControls] = useCancellableGeneration('persona modification', socket);

  const form = useForm<ModificationFormData>({
    resolver: zodResolver(modificationFormSchema),
    defaultValues: {
      modificationPrompt: "",
      llm_model: "gemini-3-pro-preview",
      reasoning_effort: "medium",
      verbosity: "medium",
    },
  });

  const handleClose = () => {
    if (modificationState.isGenerating) return; // Prevent closing while processing
    form.reset();
    modificationControls.resetGeneration();
    onClose();
  };

  const onSubmit = async (values: ModificationFormData) => {
    modificationControls.startGeneration();
    
    try {
      toastService.info("Generating persona preview...", {
        description: `Using ${values.llm_model} to create a preview of your modifications`
      });

      // Use the persona's MongoDB _id or fallback to id
      const personaId = persona._id || persona.id;
      
      const response = await personasApi.modifyWithAI(personaId, {
        modification_prompt: values.modificationPrompt,
        llm_model: values.llm_model,
        reasoning_effort: values.reasoning_effort || 'medium',
        verbosity: values.verbosity || 'medium',
        preview_only: true
      });

      // Set task ID if available
      if (response.data?.task_id) {
        modificationControls.setTaskId(response.data.task_id);
      }

      if (response.data && response.data.persona) {
        modificationControls.completeGeneration();
        
        toastService.success("Preview generated successfully!", {
          description: `Ready to review proposed changes to ${persona.name}`
        });
        
        onPersonaPreview(response.data.persona);
        handleClose();
      } else {
        throw new Error("Invalid response from server");
      }
      
    } catch (error: any) {
      // Check if this was a cancellation
      if (error.response?.status === 499) {
        return;
      }
      
      console.error("Error generating persona preview:", error);
      
      modificationControls.failGeneration(error.message || 'Failed to generate preview');
      
      if (error.response) {
        const errorMessage = error.response.data?.error || "Server error occurred";
        toastService.error("Failed to generate preview", {
          description: errorMessage
        });
      } else if (error.request) {
        toastService.error("Network error", {
          description: "Unable to connect to the server"
        });
      } else {
        toastService.error("Preview generation failed", {
          description: error.message || "An unexpected error occurred"
        });
      }
    }
  };

  return (
    <Dialog open={isOpen} onOpenChange={handleClose}>
      <DialogContent className="max-w-2xl">
        <DialogHeader>
          <DialogTitle className="flex items-center gap-3">
            <Bot className="h-5 w-5 text-primary" />
            Modify Persona with Natural Language
          </DialogTitle>
          <DialogDescription>
            Describe the changes you'd like to make to <strong>{persona.name}</strong>, 
            and AI will modify the persona data accordingly.
          </DialogDescription>
        </DialogHeader>

        {/* Progress Bar for persona modification */}
        {modificationState.isGenerating && (
          <div className="mb-4">
            <GenerationProgressBar
              isActive={modificationState.isGenerating}
              isComplete={modificationState.isComplete}
              hasError={modificationState.hasError}
              isCancelling={modificationState.isCancelling}
              label="Modifying persona with AI..."
              onCancel={modificationControls.cancelGeneration}
              taskId={modificationState.taskId}
              showCancel={true}
              className="w-full"
            />
          </div>
        )}

        <Form {...form}>
          <form onSubmit={form.handleSubmit(onSubmit)} className="space-y-6">
            {/* Modification Prompt */}
            <FormField
              control={form.control}
              name="modificationPrompt"
              render={({ field }) => (
                <FormItem>
                  <FormLabel>Modification Instructions</FormLabel>
                  <FormControl>
                    <Textarea
                      {...field}
                      placeholder="E.g., 'Make this person more tech-savvy and interested in sustainable products' or 'Increase their income level and add marketing experience'"
                      className="min-h-[120px]"
                      disabled={modificationState.isGenerating}
                    />
                  </FormControl>
                  <FormDescription>
                    Describe what aspects of the persona you'd like to modify or enhance
                  </FormDescription>
                  <FormMessage />
                </FormItem>
              )}
            />

            {/* AI Model Selection */}
            <FormField
              control={form.control}
              name="llm_model"
              render={({ field }) => (
                <FormItem>
                  <FormLabel>AI Model</FormLabel>
                  <Select 
                    onValueChange={field.onChange} 
                    value={field.value}
                    disabled={modificationState.isGenerating}
                  >
                    <FormControl>
                      <SelectTrigger>
                        <SelectValue placeholder="Select AI model" />
                      </SelectTrigger>
                    </FormControl>
                    <SelectContent>
                      <SelectItem value="gemini-3-pro-preview">Gemini 3 Pro</SelectItem>
                      <SelectItem value="gpt-4.1">GPT-4.1</SelectItem>
                      <SelectItem value="gpt-5">GPT-5</SelectItem>
                    </SelectContent>
                  </Select>
                  <FormDescription>
                    Choose which AI model to use for modifying the persona
                  </FormDescription>
                  <FormMessage />
                </FormItem>
              )}
            />

            {/* GPT-5 specific parameters */}
            {form.watch("llm_model") === "gpt-5" && (
              <>
                {/* Reasoning Effort Parameter */}
                <FormField
                  control={form.control}
                  name="reasoning_effort"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>Reasoning Effort</FormLabel>
                      <Select 
                        onValueChange={field.onChange} 
                        value={field.value}
                        disabled={modificationState.isGenerating}
                      >
                        <FormControl>
                          <SelectTrigger>
                            <SelectValue placeholder="Select reasoning effort" />
                          </SelectTrigger>
                        </FormControl>
                        <SelectContent>
                          <SelectItem value="minimal">Minimal - Fast responses</SelectItem>
                          <SelectItem value="low">Low - Quick thinking</SelectItem>
                          <SelectItem value="medium">Medium - Balanced (default)</SelectItem>
                          <SelectItem value="high">High - Deep reasoning</SelectItem>
                        </SelectContent>
                      </Select>
                      <FormDescription>
                        Controls how much the AI thinks through the modification
                      </FormDescription>
                      <FormMessage />
                    </FormItem>
                  )}
                />

                {/* Verbosity Parameter */}
                <FormField
                  control={form.control}
                  name="verbosity"
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>Response Detail</FormLabel>
                      <Select 
                        onValueChange={field.onChange} 
                        value={field.value}
                        disabled={modificationState.isGenerating}
                      >
                        <FormControl>
                          <SelectTrigger>
                            <SelectValue placeholder="Select verbosity level" />
                          </SelectTrigger>
                        </FormControl>
                        <SelectContent>
                          <SelectItem value="low">Concise</SelectItem>
                          <SelectItem value="medium">Balanced (default)</SelectItem>
                          <SelectItem value="high">Detailed</SelectItem>
                        </SelectContent>
                      </Select>
                      <FormDescription>
                        Controls how detailed the AI's modifications will be
                      </FormDescription>
                      <FormMessage />
                    </FormItem>
                  )}
                />
              </>
            )}

            <DialogFooter>
              <Button 
                type="button" 
                variant="outline" 
                onClick={handleClose}
                disabled={modificationState.isGenerating}
              >
                Cancel
              </Button>
              <Button 
                type="submit" 
                disabled={modificationState.isGenerating}
                className="flex items-center gap-2"
              >
                {modificationState.isGenerating ? (
                  <>
                    <Loader2 className="h-4 w-4 animate-spin" />
                    Generating Preview...
                  </>
                ) : (
                  <>
                    <Wand2 className="h-4 w-4" />
                    Generate Preview
                  </>
                )}
              </Button>
            </DialogFooter>
          </form>
        </Form>
      </DialogContent>
    </Dialog>
  );
}