src/lib/services/trainerScanService.ts

import type { GradioClient } from '$lib/types';
import { 
  initializeTrainerScanProgress, 
  getNextPendingImage, 
  markImageProcessingStarted,
  markImageProcessingCompleted,
  markImageProcessingFailed,
  getScanningStats,
  getCurrentProcessingImage
} from '$lib/db/trainerScanning';
import { savePicletInstance, generatedDataToPicletInstance } from '$lib/db/piclets';
import { extractPicletMetadata } from './picletMetadata';
import { removeBackground } from '$lib/utils/professionalImageProcessing';
import type { PicletStats } from '$lib/types';

export interface TrainerScanState {
  isScanning: boolean;
  currentImage: string | null;
  currentTrainer: string | null;
  progress: {
    total: number;
    completed: number;
    failed: number;
    pending: number;
  };
  error: string | null;
}

export class TrainerScanService {
  private joyCaptionClient: GradioClient;
  private zephyrClient: GradioClient;
  private fluxClient: GradioClient;
  
  private isScanning = false;
  private shouldStop = false;
  private stateCallbacks: ((state: TrainerScanState) => void)[] = [];
  
  constructor(
    joyCaptionClient: GradioClient,
    zephyrClient: GradioClient,
    fluxClient: GradioClient
  ) {
    this.joyCaptionClient = joyCaptionClient;
    this.zephyrClient = zephyrClient;
    this.fluxClient = fluxClient;
  }
  
  // Subscribe to state changes
  onStateChange(callback: (state: TrainerScanState) => void) {
    this.stateCallbacks.push(callback);
  }
  
  // Notify all subscribers of state changes
  private notifyStateChange(state: Partial<TrainerScanState>) {
    const fullState = { ...this.getCurrentState(), ...state };
    this.stateCallbacks.forEach(callback => callback(fullState));
  }
  
  // Get current scanning state
  private async getCurrentState(): Promise<TrainerScanState> {
    const stats = await getScanningStats();
    return {
      isScanning: this.isScanning,
      currentImage: null,
      currentTrainer: null,
      progress: {
        total: stats.total,
        completed: stats.completed,
        failed: stats.failed,
        pending: stats.pending
      },
      error: null
    };
  }
  
  // Initialize scanning database with image paths from file
  async initializeFromFile(): Promise<void> {
    try {
      const response = await fetch('/trainer_image_paths.txt');
      if (!response.ok) {
        throw new Error(`Failed to fetch trainer_image_paths.txt: ${response.statusText}`);
      }
      
      const content = await response.text();
      if (!content) {
        throw new Error('trainer_image_paths.txt is empty');
      }
      
      const imagePaths = content.trim().split('\n')
        .map(path => typeof path === 'string' ? path.trim() : '')
        .filter(path => path.length > 0);
      console.log(`Loaded ${imagePaths.length} trainer image paths`);
      
      if (imagePaths.length === 0) {
        throw new Error('No valid image paths found in trainer_image_paths.txt');
      }
      
      await initializeTrainerScanProgress(imagePaths);
      this.notifyStateChange(await this.getCurrentState());
    } catch (error) {
      console.error('Failed to initialize trainer scan progress:', error);
      throw new Error('Failed to load trainer image paths');
    }
  }
  
  // Start automated scanning
  async startScanning(): Promise<void> {
    if (this.isScanning) {
      throw new Error('Scanning is already in progress');
    }
    
    // Initialize database if needed
    const stats = await getScanningStats();
    if (stats.total === 0) {
      await this.initializeFromFile();
    }
    
    // Check for interrupted processing
    const currentProcessing = await getCurrentProcessingImage();
    if (currentProcessing) {
      // Reset interrupted processing back to pending
      await markImageProcessingFailed(currentProcessing.imagePath, 'Process interrupted');
    }
    
    this.isScanning = true;
    this.shouldStop = false;
    this.notifyStateChange({ isScanning: true, error: null });
    
    try {
      await this.processingLoop();
    } catch (error) {
      console.error('Scanning error:', error);
      this.notifyStateChange({ error: error instanceof Error ? error.message : 'Unknown error' });
    } finally {
      this.isScanning = false;
      this.notifyStateChange({ isScanning: false, currentImage: null, currentTrainer: null });
      
      // Log final summary
      const finalStats = await getScanningStats();
      console.log(`🏁 Scanning session complete:`, {
        total: finalStats.total,
        completed: finalStats.completed,
        failed: finalStats.failed,
        pending: finalStats.pending,
        successRate: finalStats.total > 0 ? Math.round((finalStats.completed / finalStats.total) * 100) + '%' : '0%'
      });
    }
  }
  
  // Stop scanning
  stopScanning(): void {
    this.shouldStop = true;
  }
  
  // Main processing loop
  private async processingLoop(): Promise<void> {
    while (!this.shouldStop) {
      const nextImage = await getNextPendingImage();
      
      if (!nextImage) {
        // No more pending images
        break;
      }
      
      this.notifyStateChange({ 
        currentImage: nextImage.imagePath,
        currentTrainer: typeof nextImage.trainerName === 'string' ? nextImage.trainerName : null
      });
      
      try {
        await this.processImage(nextImage.imagePath, nextImage.remoteUrl);
        console.log(`✅ Successfully processed: ${nextImage.imagePath} (${nextImage.trainerName})`);
        // Add small delay between images to prevent overwhelming the system
        await new Promise(resolve => setTimeout(resolve, 1000));
      } catch (error) {
        const errorMessage = error instanceof Error ? error.message : 'Unknown error';
        console.error(`❌ Failed to process ${nextImage.imagePath} (${nextImage.trainerName}):`, {
          imagePath: nextImage.imagePath,
          trainerName: nextImage.trainerName,
          remoteUrl: nextImage.remoteUrl,
          error: errorMessage,
          fullError: error
        });
        
        await markImageProcessingFailed(nextImage.imagePath, errorMessage);
        
        // Continue to next image - don't let individual failures stop the whole process
        console.log(`🔄 Continuing to next image despite failure...`);
      }
      
      // Update progress
      this.notifyStateChange(await this.getCurrentState());
    }
  }
  
  // Process a single remote image
  private async processImage(imagePath: string, remoteUrl: string): Promise<void> {
    await markImageProcessingStarted(imagePath);
    
    try {
      console.log(`🔄 Processing ${imagePath}: Fetching remote image...`);
      // Fetch remote image
      const imageFile = await this.fetchRemoteImage(remoteUrl, imagePath);
      
      console.log(`🔄 Processing ${imagePath}: Captioning image...`);
      // Caption the image
      const imageCaption = await this.captionImage(imageFile);
      
      console.log(`🔄 Processing ${imagePath}: Generating concept...`);
      // Generate monster concept
      const picletConcept = await this.generatePicletConcept(imageCaption);
      
      console.log(`🔄 Processing ${imagePath}: Generating stats...`);
      // Generate stats
      const picletStats = await this.generatePicletStats(picletConcept);
      
      console.log(`🔄 Processing ${imagePath}: Generating image prompt...`);
      // Generate image prompt
      const imagePrompt = await this.generateImagePrompt(picletConcept);
      
      console.log(`🔄 Processing ${imagePath}: Generating monster image...`);
      // Generate monster image
      const picletImageUrl = await this.generateMonsterImage(imagePrompt);
      
      console.log(`🔄 Processing ${imagePath}: Processing generated image...`);
      // Process generated image (remove background)
      const imageData = await this.processGeneratedImage(picletImageUrl);
      
      console.log(`🔄 Processing ${imagePath}: Creating piclet instance...`);
      // Create piclet instance
      const generatedData = {
        name: this.extractNameFromConcept(picletConcept),
        imageUrl: picletImageUrl,
        imageData,
        imageCaption,
        concept: picletConcept,
        imagePrompt,
        stats: picletStats,
        createdAt: new Date()
      };
      
      const picletInstance = await generatedDataToPicletInstance(generatedData, 5);
      const savedId = await savePicletInstance(picletInstance);
      
      await markImageProcessingCompleted(imagePath, savedId);
      
    } catch (error) {
      // Add context about which step failed
      const enhancedError = new Error(`Failed during processing of ${imagePath}: ${error instanceof Error ? error.message : 'Unknown error'}`);
      enhancedError.cause = error;
      throw enhancedError;
    }
  }
  
  // Fetch remote image and convert to File
  private async fetchRemoteImage(remoteUrl: string, originalPath: string): Promise<File> {
    const response = await fetch(remoteUrl);
    if (!response.ok) {
      throw new Error(`Failed to fetch ${remoteUrl}: ${response.statusText}`);
    }
    
    const blob = await response.blob();
    const fileName = originalPath.split('/').pop() || 'trainer_image.jpg';
    
    return new File([blob], fileName, { type: blob.type });
  }
  
  // Caption image using Joy Caption
  private async captionImage(imageFile: File): Promise<string> {
    const result = await this.joyCaptionClient.predict("/caption", [imageFile, "descriptive", "any", false]);
    const captionResult = result.data[0] as string;
    
    if (!captionResult || captionResult.trim() === '') {
      throw new Error('Failed to generate image caption');
    }
    
    return captionResult.trim();
  }
  
  // Generate piclet concept using Zephyr
  private async generatePicletConcept(imageCaption: string): Promise<string> {
    const prompt = `Based on this image description, create a unique creature concept for a Pokemon-style monster collection game called "Pictuary":

"${imageCaption}"

Create a creative, original monster concept that:
1. Is inspired by elements from the image but is clearly a fantastical creature
2. Has a unique name and personality
3. Includes special abilities related to its appearance
4. Is suitable for a family-friendly game

Write a detailed monster concept (2-3 paragraphs).`;

    const result = await this.zephyrClient.predict("/chat", [
      [["user", prompt]],
      512,  // max tokens
      0.7,  // temperature
      0.9,  // top_p
    ]);
    
    const conceptResult = result.data[0][1][1] as string;
    
    if (!conceptResult || conceptResult.trim() === '') {
      throw new Error('Failed to generate piclet concept');
    }
    
    return conceptResult.trim();
  }
  
  // Generate piclet stats
  private async generatePicletStats(concept: string): Promise<PicletStats> {
    return await extractPicletMetadata(concept, this.zephyrClient);
  }
  
  // Generate image prompt for monster creation
  private async generateImagePrompt(concept: string): Promise<string> {
    const prompt = `Extract ONLY the visual appearance from this monster concept and describe it in one concise sentence:
"${concept}"

Focus on: colors, body shape, eyes, limbs, mouth, and key visual features. Omit backstory, abilities, and non-visual details.`;

    const result = await this.zephyrClient.predict("/chat", [
      [["user", prompt]],
      256,  // max tokens
      0.5,  // temperature  
      0.9,  // top_p
    ]);
    
    const promptResult = result.data[0][1][1] as string;
    
    if (!promptResult || promptResult.trim() === '') {
      throw new Error('Failed to generate image prompt');
    }
    
    return promptResult.trim();
  }
  
  // Generate monster image using Flux
  private async generateMonsterImage(imagePrompt: string): Promise<string> {
    const fullPrompt = `${imagePrompt}, digital art, creature design, fantasy monster, clean background, professional illustration`;
    
    const result = await this.fluxClient.predict("/infer", [
      fullPrompt,
      "",  // negative prompt
      832, // width
      1216, // height  
      1,   // num inference steps
      3.0, // guidance scale
      0,   // seed
    ]);
    
    const imageUrl = result.data[0] as string;
    
    if (!imageUrl) {
      throw new Error('Failed to generate monster image');
    }
    
    return imageUrl;
  }
  
  // Process generated image (remove background)
  private async processGeneratedImage(imageUrl: string): Promise<string> {
    try {
      return await removeBackground(imageUrl);
    } catch (error) {
      console.warn('Background removal failed, using original image:', error);
      return imageUrl;
    }
  }
  
  // Extract name from concept text
  private extractNameFromConcept(concept: string): string {
    // Try to find a name in common patterns
    const patterns = [
      /(?:called|named)\s+([A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)/,
      /^([A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)/,
      /this\s+([A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)/,
    ];
    
    for (const pattern of patterns) {
      const match = concept.match(pattern);
      if (match && match[1]) {
        return match[1].trim();
      }
    }
    
    // Fallback to generating a random trainer-inspired name
    const trainerNames = ['Snap', 'Blaze', 'Nimbus', 'Breaker', 'Trinket'];
    const randomName = trainerNames[Math.floor(Math.random() * trainerNames.length)];
    return `Trainer${randomName}`;
  }
}