import type { Express } from "express";
import { createServer, type Server } from "http";
import Stripe from "stripe";
import { storage } from "./storage";
import { setupAuth, isAuthenticated } from "./replitAuth";
import multer from "multer";
import Replicate from "replicate";
import { insertPhotoshootSchema, type Photoshoot } from "@shared/schema";
import { z } from "zod";
import JSZip from "jszip";

if (!process.env.STRIPE_SECRET_KEY) {
  throw new Error('Missing required Stripe secret: STRIPE_SECRET_KEY');
}

const stripe = new Stripe(process.env.STRIPE_SECRET_KEY);

if (!process.env.REPLICATE_API_TOKEN) {
  throw new Error('Missing required Replicate API token: REPLICATE_API_TOKEN');
}
const replicate = new Replicate({
  auth: process.env.REPLICATE_API_TOKEN,
});

// Configure multer for file uploads
const upload = multer({
  storage: multer.memoryStorage(),
  limits: {
    fileSize: 10 * 1024 * 1024, // 10MB limit
  },
  fileFilter: (req, file, cb) => {
    if (file.mimetype.startsWith('image/')) {
      cb(null, true);
    } else {
      cb(new Error('Only image files are allowed'));
    }
  },
});

// Helper function to get processing delay based on package type
function getProcessingDelay(packageType: string): number {
  switch (packageType?.toLowerCase()) {
    case 'executive': return 30; // 30 minutes
    case 'professional': return 120; // 2 hours
    case 'novice':
    default: return 180; // 3 hours
  }
}

// Helper function to get number of photos based on package type
function getPackagePhotoCount(packageType: string): number {
  switch (packageType?.toLowerCase()) {
    case 'executive': return 120;
    case 'professional': return 80;
    case 'novice':
    default: return 40;
  }
}

// Process training completion and generate headshots
async function processTrainingCompletion(photoshootId: number, trainingId: string, destinationModel: string) {
  try {
    console.log(`Processing training completion for photoshoot ${photoshootId}`);

    // Get the photoshoot
    const photoshoot = await storage.getPhotoshoot(photoshootId);
    if (!photoshoot) {
      console.error(`Photoshoot ${photoshootId} not found`);
      return;
    }

    // Check training status first
    const training = await replicate.trainings.get(trainingId);

    if (training.status === 'failed') {
      console.error(`Training ${trainingId} failed:`, training.error);
      await storage.updatePhotoshoot(photoshootId, {
        status: 'failed',
        error: `Training failed: ${training.error || 'Unknown error'}`
      });
      return;
    }

    if (training.status !== 'succeeded') {
      console.log(`Training ${trainingId} not yet completed. Status: ${training.status}`);
      // Reschedule check in 10 minutes
      setTimeout(() => processTrainingCompletion(photoshootId, trainingId, destinationModel), 10 * 60 * 1000);
      return;
    }

    // Get the trained model information from the training result
    console.log('Training completed successfully. Training details:', {
      id: training.id,
      status: training.status,
      version: training.version,
      output: training.output
    });

    // When training completes, Replicate provides the trained model reference
    // We need to get the actual trained model path from the training result
    let trainedModelPath;

    console.log(`Training output:`, JSON.stringify(training.output, null, 2));
    console.log(`Training urls:`, training.urls);

    // Check if training.output contains the model reference
    if (training.output) {
      // Sometimes the output contains the model path directly
      if (typeof training.output === 'string' && training.output.includes(':')) {
        trainedModelPath = training.output;
        console.log(`Found trained model path in output string: ${trainedModelPath}`);
      } else if (training.output.model) {
        trainedModelPath = training.output.model;
        console.log(`Found trained model path in output.model: ${trainedModelPath}`);
      } else if (training.output.version) {
        // Construct the model path using the destination and version
        trainedModelPath = `${destinationModel}:${training.output.version}`;
        console.log(`Constructed model path from output.version: ${trainedModelPath}`);
      }
    }

    // Fallback: check training.version
    if (!trainedModelPath && training.version) {
      trainedModelPath = `${destinationModel}:${training.version}`;
      console.log(`Constructed model path from training.version: ${trainedModelPath}`);
    }

    // Final fallback: get the latest version from the destination model
    if (!trainedModelPath) {
      try {
        const [owner, modelName] = destinationModel.split('/');
        const model = await replicate.models.get(owner, modelName);

        if (model.latest_version && model.latest_version.id) {
          trainedModelPath = `${destinationModel}:${model.latest_version.id}`;
          console.log(`Constructed model path from latest version: ${trainedModelPath}`);
        }
      } catch (modelError) {
        console.error('Error getting model from API:', modelError);
      }
    }

    if (!trainedModelPath) {
      console.error(`Could not determine trained model path for training ${trainingId}`);
      console.error('Training object:', JSON.stringify(training, null, 2));
      await storage.updatePhotoshoot(photoshootId, {
        status: 'failed',
        error: 'Could not determine trained model path after training completion'
      });
      return;
    }

    console.log(`Using trained model for generation: ${trainedModelPath}`);

    // Update status to processing
    await storage.updatePhotoshoot(photoshootId, {
      status: 'processing'
    });

    // Generate professional outfit prompts based on gender and selections
    const prompts = generateHeadshotPrompts(
      photoshoot.gender || 'male',
      photoshoot.selectedOutfits || [],
      photoshoot.selectedBackdrops || [],
      getPackagePhotoCount(photoshoot.packageType)
    );

    console.log(`Generating ${prompts.length} headshots for photoshoot ${photoshootId}`);

    // Generate all headshots using the trained model
    console.log(`Starting generation of ${prompts.length} headshots using model: ${trainedModelPath}`);
    const generatedPhotos: string[] = [];
    const batchSize = 3; // Smaller batches for better reliability

    for (let i = 0; i < prompts.length; i += batchSize) {
      const batch = prompts.slice(i, i + batchSize);
      console.log(`Processing batch ${Math.floor(i/batchSize) + 1} of ${Math.ceil(prompts.length/batchSize)} (${batch.length} images)`);

      const batchPromises = batch.map(async (prompt, batchIndex) => {
        const imageIndex = i + batchIndex + 1;
        try {
          console.log(`Generating image ${imageIndex}/${prompts.length} using trained model: ${trainedModelPath}`);
          console.log(`Prompt: "${prompt.substring(0, 100)}..."`);

          // Generate image using trained model with ONLY prompt parameter (as specified)
          console.log(`Calling replicate.run with model: ${trainedModelPath}`);
          const prediction = await replicate.run(trainedModelPath as any, {
            input: {
              prompt: prompt
            }
          });

          console.log(`Raw prediction response for image ${imageIndex}:`, JSON.stringify(prediction, null, 2));

          console.log(`Generated image ${imageIndex}, prediction result:`, prediction);

          // Handle different possible return formats
          let imageUrl = null;
          if (Array.isArray(prediction) && prediction.length > 0) {
            imageUrl = prediction[0];
          } else if (typeof prediction === 'string') {
            imageUrl = prediction;
          } else if (prediction && (prediction as any).output) {
            const output = (prediction as any).output;
            imageUrl = Array.isArray(output) ? output[0] : output;
          }

          if (!imageUrl) {
            console.error(`No image URL returned for image ${imageIndex}`);
            return null;
          }

          console.log(`Image ${imageIndex} generated successfully: ${imageUrl}`);

          // Upscale the generated image
          try {
            console.log(`Upscaling image ${imageIndex}...`);
            const upscaled = await replicate.run("philz1337x/clarity-upscaler", {
              input: {
                image: imageUrl,
                scale: 2,
                dynamic: 6,
                creativity: 0.35,
                resemblance: 0.6,
                hdr: 0,
                clarity: 0
              }
            });

            const upscaledUrl = Array.isArray(upscaled) ? upscaled[0] : upscaled;
            console.log(`Image ${imageIndex} upscaled successfully: ${upscaledUrl}`);
            return upscaledUrl;
          } catch (upscaleError) {
            console.error(`Error upscaling image ${imageIndex}:`, upscaleError);
            // Return original image if upscaling fails
            return imageUrl;
          }

        } catch (error) {
          console.error(`Error generating image ${imageIndex}:`, error);
          return null;
        }
      });

      const batchResults = await Promise.all(batchPromises);
      const validResults = batchResults.filter(Boolean);
      generatedPhotos.push(...validResults);

      console.log(`Batch ${Math.floor(i/batchSize) + 1} completed. Generated ${validResults.length}/${batch.length} images successfully.`);

      // Delay between batches to avoid rate limits
      if (i + batchSize < prompts.length) {
        console.log('Waiting 3 seconds before next batch...');
        await new Promise(resolve => setTimeout(resolve, 3000));
      }
    }

    console.log(`Generated ${generatedPhotos.length} headshots for photoshoot ${photoshootId}`);

    // Update photoshoot with generated photos
    await storage.updatePhotoshoot(photoshootId, {
      status: 'completed',
      generatedPhotos: { urls: generatedPhotos },
      completedAt: new Date(),
    });

    console.log(`Photoshoot ${photoshootId} completed successfully`);
  } catch (error) {
    console.error(`Error processing training completion for photoshoot ${photoshootId}:`, error);
    await storage.updatePhotoshoot(photoshootId, { status: 'failed' });
  }
}

// Generate professional headshot prompts
function generateHeadshotPrompts(gender: string, outfits: string[], backdrops: string[], count: number): string[] {
  const basePrompt = "OHWX person, professional business headshot, ultra sharp, high resolution, professional studio lighting, clean composition";

  const genderSpecific = gender.toLowerCase() === 'female'
    ? "professional businesswoman"
    : "professional businessman";

  // Default options if none selected
  const defaultOutfits = [
    "navy blue business suit",
    "charcoal gray suit",
    "black blazer with white shirt",
    "professional dress shirt",
    "business casual blazer"
  ];

  const defaultBackdrops = [
    "clean white studio background",
    "soft gray backdrop",
    "neutral office environment",
    "minimalist corporate setting",
    "professional studio backdrop"
  ];

  const finalOutfits = outfits.length > 0 ? outfits : defaultOutfits;
  const finalBackdrops = backdrops.length > 0 ? backdrops : defaultBackdrops;

  const prompts: string[] = [];

  // Create variations combining outfits and backdrops
  for (let i = 0; i < count; i++) {
    const outfit = finalOutfits[i % finalOutfits.length];
    const backdrop = finalBackdrops[i % finalBackdrops.length];

    const variations = [
      "looking directly at camera",
      "slight smile, confident expression",
      "serious professional expression",
      "friendly business demeanor",
      "confident executive pose"
    ];

    const variation = variations[i % variations.length];

    const prompt = `${basePrompt}, ${genderSpecific}, wearing ${outfit}, ${backdrop}, ${variation}, corporate headshot photography, executive portrait style, professional business photo`;
    prompts.push(prompt);
  }

  return prompts;
}

export async function registerRoutes(app: Express): Promise<Server> {
  // Auth middleware
  await setupAuth(app);

  // Auth routes
  app.get('/api/auth/user', isAuthenticated, async (req: any, res) => {
    try {
      const userId = req.user.claims.sub;
      const user = await storage.getUser(userId);
      res.json(user);
    } catch (error) {
      console.error("Error fetching user:", error);
      res.status(500).json({ message: "Failed to fetch user" });
    }
  });

  // Get all photoshoots for the authenticated user
  app.get('/api/photoshoots', isAuthenticated, async (req: any, res) => {
    try {
      const userId = req.user.claims.sub;
      const photoshoots = await storage.getUserPhotoshoots(userId);
      res.json(photoshoots);
    } catch (error) {
      console.error("Error fetching photoshoots:", error);
      res.status(500).json({ message: "Failed to fetch photoshoots" });
    }
  });

  // Photoshoot routes
  app.get('/api/photoshoots', isAuthenticated, async (req: any, res) => {
    try {
      const userId = req.user.claims.sub;
      const photoshoots = await storage.getUserPhotoshoots(userId);
      res.json(photoshoots);
    } catch (error) {
      console.error("Error fetching photoshoots:", error);
      res.status(500).json({ message: "Failed to fetch photoshoots" });
    }
  });

  app.get('/api/photoshoots/:id', isAuthenticated, async (req: any, res) => {
    try {
      const userId = req.user.claims.sub;
      const photoshootId = parseInt(req.params.id);
      const photoshoot = await storage.getPhotoshoot(photoshootId);

      if (!photoshoot || photoshoot.userId !== userId) {
        return res.status(404).json({ message: "Photoshoot not found" });
      }

      res.json(photoshoot);
    } catch (error) {
      console.error("Error fetching photoshoot:", error);
      res.status(500).json({ message: "Failed to fetch photoshoot" });
    }
  });

  // Check training status and trigger generation if ready
  app.post('/api/photoshoots/:id/check-training', isAuthenticated, async (req: any, res) => {
    try {
      const userId = req.user.claims.sub;
      const photoshootId = parseInt(req.params.id);
      const photoshoot = await storage.getPhotoshoot(photoshootId);

      if (!photoshoot || photoshoot.userId !== userId) {
        return res.status(404).json({ message: "Photoshoot not found" });
      }

      if (!photoshoot.replicateTrainingId) {
        return res.status(400).json({ message: "No training found for this photoshoot" });
      }

      // Check training status
      const training = await replicate.trainings.get(photoshoot.replicateTrainingId);
      console.log(`Training ${photoshoot.replicateTrainingId} status: ${training.status}`);

      if (training.status === 'succeeded') {
        // Get model name from input photos
        const inputPhotos = photoshoot.inputPhotos as any;
        const modelName = inputPhotos?.modelName;

        if (modelName) {
          const replicateUsername = process.env.REPLICATE_USERNAME;
          const destinationModel = `${replicateUsername}/${modelName}`;
          console.log(`Training completed, starting generation for ${destinationModel}`);

          // Start generation process
          processTrainingCompletion(photoshootId, photoshoot.replicateTrainingId, destinationModel);

          res.json({
            status: 'training_complete',
            message: 'Training completed, starting generation process',
            trainingStatus: training.status
          });
        } else {
          res.status(500).json({ message: "Model name not found in photoshoot data" });
        }
      } else if (training.status === 'failed') {
        await storage.updatePhotoshoot(photoshootId, { status: 'failed' });
        res.json({
          status: 'failed',
          message: 'Training failed',
          trainingStatus: training.status
        });
      } else {
        res.json({
          status: 'training',
          message: 'Training still in progress',
          trainingStatus: training.status
        });
      }
    } catch (error) {
      console.error("Error checking training status:", error);
      res.status(500).json({ message: "Failed to check training status" });
    }
  });

  // Photo upload and training endpoint
  app.post('/api/photoshoots/:id/upload', isAuthenticated, upload.array('photos', 30), async (req: any, res) => {
    console.log('=== Upload endpoint called ===');
    const userId = req.user?.claims?.sub;
    const photoshootId = parseInt(req.params.id);
    const files = req.files as Express.Multer.File[];

    console.log('Initial request data:', {
      userId,
      photoshootId,
      hasUser: !!req.user,
      hasFiles: !!files,
      filesCount: files ? files.length : 0,
      replicateToken: !!process.env.REPLICATE_API_TOKEN,
      stripeKey: !!process.env.STRIPE_SECRET_KEY
    });

    try {
      console.log('Upload request received:', {
        userId,
        photoshootId,
        filesCount: files ? files.length : 0,
        hasFiles: !!files,
        requestFiles: req.files,
        requestBody: Object.keys(req.body || {}),
      });

      if (!files || files.length < 10) {
        console.log(`Upload rejected: ${files ? files.length : 0} files, need at least 10`);
        return res.status(400).json({ message: `Please upload at least 10 photos. Received: ${files ? files.length : 0}` });
      }

      // Get photoshoot and verify ownership
      const photoshoot = await storage.getPhotoshoot(photoshootId);
      if (!photoshoot || photoshoot.userId !== userId) {
        return res.status(404).json({ message: "Photoshoot not found" });
      }

      // Create a zip file containing all the uploaded images
      console.log('Creating zip file from uploaded images...');
      const zip = new JSZip();

      // Add each image file to the zip
      files.forEach((file, index) => {
        const extension = file.originalname.split('.').pop() || 'jpg';
        const filename = `training_image_${index + 1}.${extension}`;
        console.log(`Adding file ${filename}, size: ${file.buffer.length} bytes`);
        zip.file(filename, file.buffer);
      });

      // Generate the zip file as base64
      console.log('Generating zip file...');
      const zipBuffer = await zip.generateAsync({ type: 'nodebuffer' });
      const zipBase64 = zipBuffer.toString('base64');
      const zipDataUrl = `data:application/zip;base64,${zipBase64}`;

      console.log(`Created zip file with ${files.length} images, size: ${zipBuffer.length} bytes`);

      // Create unique destination model name for this user
      const timestamp = Date.now();
      const modelName = `fhotosonic-user-${userId}-${timestamp}`;

      console.log(`Starting training for photoshoot ${photoshootId} with ${files.length} photos`);
      console.log(`Creating destination model: ${modelName}`);

      // Create destination model using the user's Replicate username
      const replicateUsername = process.env.REPLICATE_USERNAME;
      if (!replicateUsername) {
        throw new Error('REPLICATE_USERNAME environment variable is required for training');
      }

      const destinationModel = `${replicateUsername}/${modelName}` as `${string}/${string}`;
      console.log('Training destination model:', destinationModel);

      // First, create the destination model if it doesn't exist
      console.log('Creating destination model...');
      try {
        await replicate.models.create(
          replicateUsername,
          modelName,
          {
            visibility: "private",
            hardware: "gpu-t4",
            description: `Custom FLUX LoRA model for user ${userId} - ${photoshoot.name}`
          }
        );
        console.log('Successfully created destination model:', destinationModel);
        // Small delay to ensure model is properly registered
        await new Promise(resolve => setTimeout(resolve, 2000));
      } catch (modelError: any) {
        console.log('Model creation result:', modelError.message);
        // Model might already exist, which is fine - we'll continue with training
        if (!modelError.message.includes('already exists') && !modelError.message.includes('name is already taken')) {
          console.error('Unexpected model creation error:', modelError);
          // Don't throw here - let's try training anyway in case the model exists
        }
      }

      // Start Replicate training with zip file
      console.log('Starting Replicate training with zip file...');
      console.log('Training parameters:', {
        owner: "ostris",
        model: "flux-dev-lora-trainer",
        version: "4ffd32160efd92e956d39c5338a9b8fbafca58e03f791f6d8011f3e20e8ea6fa",
        destination: destinationModel,
        zipDataLength: zipDataUrl.length
      });

      // Create training with retry logic for network issues
      let training: any;
      let retryCount = 0;
      const maxRetries = 3;

      while (retryCount < maxRetries) {
        try {
          console.log(`Training attempt ${retryCount + 1}/${maxRetries}`);

          training = await replicate.trainings.create(
            "ostris",
            "flux-dev-lora-trainer",
            "4ffd32160efd92e956d39c5338a9b8fbafca58e03f791f6d8011f3e20e8ea6fa",
            {
              destination: destinationModel,
              input: {
                steps: 800, // Reduced steps for stability
                lora_rank: 16,
                optimizer: "adamw8bit",
                batch_size: 1,
                resolution: "1024",
                autocaption: true,
                trigger_word: "OHWX person",
                input_images: zipDataUrl,
                learning_rate: 0.0004,
                cache_latents_to_disk: false, // Avoid disk caching issues
                wandb_project: "", // Disable wandb to avoid extra dependencies
                is_lora: true
              }
            }
          );

          console.log('Training created successfully on attempt', retryCount + 1);
          break; // Success, exit retry loop

        } catch (trainingError: any) {
          console.error(`Training attempt ${retryCount + 1} failed:`, trainingError.message);
          retryCount++;

          if (retryCount >= maxRetries) {
            throw new Error(`Training failed after ${maxRetries} attempts: ${trainingError.message}`);
          }

          // Wait before retry (exponential backoff)
          await new Promise(resolve => setTimeout(resolve, Math.pow(2, retryCount) * 1000));
        }
      }

      if (!training) {
        throw new Error('Training creation failed - no training object returned');
      }

      console.log('Replicate training created successfully:', training.id);

      // Update photoshoot with training info and status
      await storage.updatePhotoshoot(photoshootId, {
        status: 'training',
        replicateTrainingId: training.id,
        inputPhotos: { zipFile: zipDataUrl, count: files.length, modelName },
      });

      // Set up delayed processing based on package type
      const delayMinutes = getProcessingDelay(photoshoot.packageType);
      console.log(`Training started. Scheduled completion in ${delayMinutes} minutes for ${photoshoot.packageType} package`);

      // Schedule the generation process with automatic polling
      setTimeout(async () => {
        console.log(`Scheduled check for training completion - photoshoot ${photoshootId}, training ${training.id}`);
        await processTrainingCompletion(photoshootId, training.id, destinationModel);
      }, delayMinutes * 60 * 1000);

      // Also start immediate polling every 5 minutes to check for early completion
      const pollInterval = setInterval(async () => {
        try {
          const currentPhotoshoot = await storage.getPhotoshoot(photoshootId);
          if (!currentPhotoshoot || currentPhotoshoot.status !== 'training') {
            clearInterval(pollInterval);
            return;
          }

          const currentTraining = await replicate.trainings.get(training.id);
          console.log(`Polling training ${training.id} status: ${currentTraining.status}`);

          if (currentTraining.status === 'succeeded') {
            console.log(`Training ${training.id} completed early! Starting generation immediately.`);
            clearInterval(pollInterval);
            await processTrainingCompletion(photoshootId, training.id, destinationModel);
          } else if (currentTraining.status === 'failed') {
            console.error(`Training ${training.id} failed during polling.`);
            clearInterval(pollInterval);
            await storage.updatePhotoshoot(photoshootId, {
              status: 'failed',
              error: `Training failed: ${currentTraining.error || 'Unknown error'}`
            });
          }
        } catch (pollError) {
          console.error(`Error during training poll for ${training.id}:`, pollError);
        }
      }, 5 * 60 * 1000); // Poll every 5 minutes

      res.json({
        message: "Training started successfully",
        trainingId: training.id,
        estimatedCompletion: delayMinutes
      });
    } catch (error: any) {
      console.error("Error uploading photos:", error);
      console.error('Error details:', {
        message: error.message,
        stack: error.stack,
        photoshootId: parseInt(req.params.id),
        userId: req.user?.claims?.sub,
        filesLength: req.files ? (req.files as Express.Multer.File[]).length : 0
      });

      // Update photoshoot status to failed
      try {
        await storage.updatePhotoshoot(parseInt(req.params.id), { status: 'failed' });
      } catch (updateError) {
        console.error('Failed to update photoshoot status:', updateError);
      }

      res.status(500).json({ message: "Failed to upload photos: " + error.message });
    }
  });

  // Stripe payment route for photoshoot packages
  app.post("/api/create-payment-intent", isAuthenticated, async (req: any, res) => {
    try {
      const userId = req.user.claims.sub;
      const { packageType } = req.body;

      // Package pricing
      const packagePricing = {
        novice: 29,
        professional: 39,
        executive: 59
      };

      const amount = packagePricing[packageType as keyof typeof packagePricing];
      if (!amount) {
        return res.status(400).json({ message: "Invalid package type" });
      }

      const paymentIntent = await stripe.paymentIntents.create({
        amount: Math.round(amount * 100), // Convert to cents
        currency: "usd",
        metadata: {
          userId,
          packageType
        }
      });

      // Create pending photoshoot
      const photoshoot = await storage.createPhotoshoot({
        userId,
        name: `${packageType} Photoshoot`,
        packageType,
        paymentIntentId: paymentIntent.id,
        amount: Math.round(amount * 100), // Store in cents
        status: "pending"
      });

      res.json({
        clientSecret: paymentIntent.client_secret,
        photoshootId: photoshoot.id
      });
    } catch (error: any) {
      res.status(500).json({ message: "Error creating payment intent: " + error.message });
    }
  });

  // Check training status manually
  app.get("/api/photoshoots/:id/check-training", isAuthenticated, async (req: any, res) => {
    try {
      const photoshootId = parseInt(req.params.id);
      const userId = req.user.claims.sub;

      const photoshoot = await storage.getPhotoshoot(photoshootId);
      if (!photoshoot) {
        return res.status(404).json({ message: "Photoshoot not found" });
      }

      if (photoshoot.userId !== userId) {
        return res.status(403).json({ message: "Unauthorized" });
      }

      if (!photoshoot.replicateTrainingId) {
        return res.json({ status: photoshoot.status, message: "No training in progress" });
      }

      // Check training status with Replicate
      const training = await replicate.trainings.get(photoshoot.replicateTrainingId);
      console.log(`Training ${photoshoot.replicateTrainingId} status: ${training.status}`);

      // If training completed and we haven't processed it yet, trigger processing
      if (training.status === 'succeeded' && photoshoot.status === 'training') {
        const inputPhotos = photoshoot.inputPhotos as any;
        const destinationModel = inputPhotos?.modelName || 'unknown';

        console.log(`Training ${photoshoot.replicateTrainingId} completed! Triggering image generation...`);
        console.log('Destination model:', destinationModel);

        // Trigger immediate processing without waiting
        if (photoshoot.replicateTrainingId && destinationModel !== 'unknown') {
          const trainingId = photoshoot.replicateTrainingId;
          setImmediate(() => {
            processTrainingCompletion(photoshootId, trainingId, destinationModel);
          });
        } else {
          console.error('Missing training ID or destination model for processing');
        }

        res.json({
          status: 'processing',
          message: 'Training completed, starting image generation...',
          trainingStatus: training.status
        });
      } else {
        res.json({
          status: photoshoot.status,
          trainingStatus: training.status,
          message: `Training ${training.status}`
        });
      }
    } catch (error: any) {
      console.error("Error checking training status:", error);
      res.status(500).json({ message: "Error checking training status: " + error.message });
    }
  });

  // Handle successful payment confirmation
  app.post("/api/payment-success", isAuthenticated, async (req: any, res) => {
    try {
      const { paymentIntentId, name } = req.body;
      const userId = req.user.claims.sub;

      // Verify payment intent
      const paymentIntent = await stripe.paymentIntents.retrieve(paymentIntentId);

      if (paymentIntent.status !== 'succeeded') {
        return res.status(400).json({ message: "Payment not completed" });
      }

      // Find existing photoshoot by payment intent ID
      const photoshoot = await storage.getPhotoshootByPaymentIntent(paymentIntentId);

      if (!photoshoot) {
        return res.status(404).json({ message: "Photoshoot not found" });
      }

      if (photoshoot.userId !== userId) {
        return res.status(403).json({ message: "Unauthorized" });
      }

      // Update photoshoot with custom name if provided
      if (name && name !== photoshoot.name) {
        await storage.updatePhotoshoot(photoshoot.id, { name });
      }

      res.json({ photoshoot: { ...photoshoot, name: name || photoshoot.name } });
    } catch (error: any) {
      res.status(500).json({ message: "Error confirming payment: " + error.message });
    }
  });

  // Webhook for payment confirmation
  app.post("/api/stripe-webhook", async (req, res) => {
    const sig = req.headers['stripe-signature'];
    let event;

    try {
      event = stripe.webhooks.constructEvent(req.body, sig!, process.env.STRIPE_WEBHOOK_SECRET!);
    } catch (err: any) {
      console.log(`Webhook signature verification failed.`, err.message);
      return res.status(400).send(`Webhook Error: ${err.message}`);
    }

    if (event.type === 'payment_intent.succeeded') {
      const paymentIntent = event.data.object as Stripe.PaymentIntent;
      const { userId, packageType } = paymentIntent.metadata;

      // Update photoshoot status
      await storage.updatePhotoshootStatus(paymentIntent.id, "pending");

      console.log(`Payment confirmed for user ${userId}, package ${packageType}`);
    }

    res.json({ received: true });
  });

  // Update photoshoot with gender, name, outfits, and backdrops
  app.put("/api/photoshoots/:id", isAuthenticated, async (req: any, res) => {
    try {
      const photoshootId = parseInt(req.params.id);
      const userId = req.user.claims.sub;
      const { name, gender, selectedOutfits, selectedBackdrops } = req.body;

      // Verify ownership
      const photoshoot = await storage.getPhotoshoot(photoshootId);
      if (!photoshoot || photoshoot.userId !== userId) {
        return res.status(404).json({ message: "Photoshoot not found" });
      }

      await storage.updatePhotoshoot(photoshootId, {
        name,
        gender,
        selectedOutfits,
        selectedBackdrops,
        status: "uploaded"
      });

      res.json({ message: "Photoshoot updated successfully" });
    } catch (error: any) {
      res.status(500).json({ message: "Error updating photoshoot: " + error.message });
    }
  });

  // File upload endpoint for training photos
  app.post("/api/photoshoots/:id/upload", isAuthenticated, async (req: any, res) => {
    try {
      const photoshootId = parseInt(req.params.id);
      const userId = req.user.claims.sub;

      // Verify ownership
      const photoshoot = await storage.getPhotoshoot(photoshootId);
      if (!photoshoot || photoshoot.userId !== userId) {
        return res.status(404).json({ message: "Photoshoot not found" });
      }

      // TODO: Implement file upload with multer and cloud storage
      // For now, return mock success
      res.json({ message: "Photos uploaded successfully" });
    } catch (error: any) {
      res.status(500).json({ message: "Error uploading photos: " + error.message });
    }
  });

  const httpServer = createServer(app);
  return httpServer;
}