src/lib/components/MonsterGenerator/MonsterGenerator.svelte

<script lang="ts">
  import type { MonsterGeneratorProps, MonsterWorkflowState, CaptionType, CaptionLength, MonsterStats } from '$lib/types';
  import UploadStep from './UploadStep.svelte';
  import WorkflowProgress from './WorkflowProgress.svelte';
  import MonsterResult from './MonsterResult.svelte';
  import { makeWhiteTransparent } from '$lib/utils/imageProcessing';
  import { saveMonster } from '$lib/db/monsters';
  
  interface Props extends MonsterGeneratorProps {}
  
  let { joyCaptionClient, rwkvClient, fluxClient }: Props = $props();
  
  let state: MonsterWorkflowState = $state({
    currentStep: 'upload',
    userImage: null,
    imageCaption: null,
    monsterConcept: null,
    imagePrompt: null,
    monsterImage: null,
    error: null,
    isProcessing: false
  });
  
  // Prompt templates
  const MONSTER_CONCEPT_PROMPT = (caption: string) => `Based on this image caption: "${caption}"

Create a Pokémon-style monster that transforms the object into an imaginative creature. The monster should clearly be inspired by the object's appearance but reimagined as a living monster.

Guidelines:
- Take the object's key visual elements (colors, shapes, materials) incorporating all of them into a single creature design
- Add eyes (can be glowing, mechanical, multiple, etc.) positioned where they make sense
- Include limbs (legs, arms, wings, tentacles) that grow from or replace parts of the object
- Add a mouth, beak, or feeding apparatus if appropriate
- Add creature elements like tail, fins, claws, or horns where fitting

Rarity assessment: Common objects = weaker monsters. Unique/rare objects = stronger monsters.

Include:
- A creative name that hints at the original object
- Physical description showing how the object becomes a creature
- Special abilities derived from the object's function
- Personality traits based on the object's purpose`;

  const IMAGE_GENERATION_PROMPT = (concept: string) => `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 MONSTER_STATS_PROMPT = (concept: string) => `Convert the following monster concept into a JSON object with stats:
"${concept}"

The output should be formatted as a JSON instance that conforms to the JSON schema below.

\`\`\`json
{
  "properties": {
    "name": {"type": "string", "description": "The monster's unique name"},
    "description": {"type": "string", "description": "A brief physical description of the monster's appearance"},
    "rarity": {"type": "integer", "minimum": 0, "maximum": 100, "description": "How rare/unique the monster is (0=very common, 100=legendary)"},
    "HP": {"type": "integer", "minimum": 0, "maximum": 100, "description": "Health/vitality stat (0=fragile, 100=incredibly tanky)"},
    "defence": {"type": "integer", "minimum": 0, "maximum": 100, "description": "Defensive/armor stat (0=paper thin, 100=impenetrable fortress)"},
    "attack": {"type": "integer", "minimum": 0, "maximum": 100, "description": "Physical attack power (0=harmless, 100=devastating force)"},
    "speed": {"type": "integer", "minimum": 0, "maximum": 100, "description": "Movement and reaction speed (0=immobile, 100=lightning fast)"},
    "specialPassiveTraitDescription": {"type": "string", "description": "Describe a passive ability that gives this monster a unique advantage in battle"},
    "attackActionName": {"type": "string", "description": "Name of the monster's primary damage-dealing attack (e.g., 'Flame Burst', 'Toxic Bite')"},
    "attackActionDescription": {"type": "string", "description": "Describe how this attack damages the opponent and any special effects"},
    "buffActionName": {"type": "string", "description": "Name of the monster's self-enhancement ability (e.g., 'Iron Defense', 'Speed Boost')"},
    "buffActionDescription": {"type": "string", "description": "Describe which stats are boosted and how this improves the monster's battle performance"},
    "debuffActionName": {"type": "string", "description": "Name of the monster's enemy-weakening ability (e.g., 'Intimidate', 'Slow Poison')"},
    "debuffActionDescription": {"type": "string", "description": "Describe which enemy stats are lowered and how this weakens the opponent"},
    "specialActionName": {"type": "string", "description": "Name of the monster's ultimate move (one use per battle)"},
    "specialActionDescription": {"type": "string", "description": "Describe this powerful finishing move and its dramatic effects in battle"}
  },
  "required": ["name", "description", "rarity", "HP", "defence", "attack", "speed", "specialPassiveTraitDescription", "attackActionName", "attackActionDescription", "buffActionName", "buffActionDescription", "debuffActionName", "debuffActionDescription", "specialActionName", "specialActionDescription"]
}
\`\`\`

Remember to base the stats on how unique/powerful the original object was. Common objects should have lower stats, unique objects should have higher stats.

Write your response within \`\`\`json\`\`\``;

  async function handleImageSelected(file: File) {
    if (!joyCaptionClient || !rwkvClient || !fluxClient) {
      state.error = "Services not connected. Please wait...";
      return;
    }
    
    state.userImage = file;
    state.error = null;
    startWorkflow();
  }
  
  async function startWorkflow() {
    state.isProcessing = true;
    
    try {
      // Step 1: Caption the image
      await captionImage();
      await new Promise(resolve => setTimeout(resolve, 100)); // Small delay for state update
      
      // Step 2: Generate monster concept
      await generateMonsterConcept();
      await new Promise(resolve => setTimeout(resolve, 100)); // Small delay for state update
      
      // Step 3: Generate monster stats
      await generateStats();
      await new Promise(resolve => setTimeout(resolve, 100)); // Small delay for state update
      
      // Step 4: Generate image prompt
      await generateImagePrompt();
      await new Promise(resolve => setTimeout(resolve, 100)); // Small delay for state update
      
      // Step 5: Generate monster image
      await generateMonsterImage();
      
      // Step 6: Auto-save the monster
      await autoSaveMonster();
      
      state.currentStep = 'complete';
    } catch (err) {
      console.error('Workflow error:', err);
      
      // Check for GPU quota error
      if (err && typeof err === 'object' && 'message' in err) {
        const errorMessage = String(err.message);
        if (errorMessage.includes('exceeded your GPU quota') || errorMessage.includes('GPU quota')) {
          state.error = 'GPU quota exceeded! You need to sign in with Hugging Face for free GPU time, or upgrade to Hugging Face Pro for more quota.';
        } else {
          state.error = errorMessage;
        }
      } else if (err instanceof Error) {
        state.error = err.message;
      } else {
        state.error = 'An unknown error occurred';
      }
    } finally {
      state.isProcessing = false;
    }
  }
  
  function handleAPIError(error: any): never {
    console.error('API Error:', error);
    
    // Check if it's a GPU quota error
    if (error && typeof error === 'object' && 'message' in error) {
      const errorMessage = String(error.message);
      if (errorMessage.includes('exceeded your GPU quota') || errorMessage.includes('GPU quota')) {
        throw new Error('GPU quota exceeded! You need to sign in with Hugging Face for free GPU time, or upgrade to Hugging Face Pro for more quota.');
      }
      throw new Error(errorMessage);
    }
    
    // Check if error has a different structure (like the status object from the logs)
    if (error && typeof error === 'object' && 'type' in error && error.type === 'status') {
      const statusError = error as any;
      if (statusError.message && statusError.message.includes('GPU quota')) {
        throw new Error('GPU quota exceeded! You need to sign in with Hugging Face for free GPU time, or upgrade to Hugging Face Pro for more quota.');
      }
      throw new Error(statusError.message || 'API request failed');
    }
    
    throw error;
  }
  
  async function captionImage() {
    state.currentStep = 'captioning';
    
    if (!joyCaptionClient || !state.userImage) {
      throw new Error('Caption service not available or no image provided');
    }
    
    try {
      const output = await joyCaptionClient.predict("/stream_chat", [
        state.userImage,
        "Descriptive" as CaptionType,
        "very long" as CaptionLength,
        [], // extra_options
        "", // name_input
        ""  // custom_prompt
      ]);
      
      const [prompt, caption] = output.data;
      state.imageCaption = caption;
      console.log('Caption generated:', caption);
    } catch (error) {
      handleAPIError(error);
    }
  }
  
  async function generateMonsterConcept() {
    state.currentStep = 'conceptualizing';
    
    if (!rwkvClient || !state.imageCaption) {
      throw new Error('Text generation service not available or no caption');
    }
    
    const conceptPrompt = MONSTER_CONCEPT_PROMPT(state.imageCaption);
    const systemPrompt = "You are a creative Pokémon designer. Transform everyday objects into imaginative monsters by blending their characteristics with creature features. Think like the designers who turned a Pokéball into Voltorb, a keyring into Klefki, or an ice cream cone into Vanillite. Be creative and whimsical while keeping the object recognizable.";
    
    console.log('Generating monster concept with prompt:', conceptPrompt);
    
    try {
      const output = await rwkvClient.predict("/chat", [
        conceptPrompt,          // message
        [],                     // chat_history
        systemPrompt,          // system_prompt
        2048,                  // max_new_tokens
        0.7,                   // temperature
        0.95,                  // top_p
        50,                    // top_k
        1.0                    // repetition_penalty
      ]);
      
      console.log('Zephyr output:', output.data[0]);
      let conceptText = output.data[0];
      
      state.monsterConcept = conceptText;
      console.log('Monster concept generated:', state.monsterConcept);
      
      if (!state.monsterConcept || state.monsterConcept.trim() === '') {
        throw new Error('Failed to generate monster concept - received empty response');
      }
    } catch (error) {
      handleAPIError(error);
    }
  }
  
  async function generateImagePrompt() {
    state.currentStep = 'promptCrafting';
    
    if (!rwkvClient || !state.monsterConcept) {
      throw new Error('Text generation service not available or no concept');
    }
    
    const promptGenerationPrompt = IMAGE_GENERATION_PROMPT(state.monsterConcept);
    const systemPrompt = "You are an expert at creating concise visual descriptions for image generation. Extract ONLY visual appearance details and describe them in ONE sentence (max 50 words). Focus on colors, shape, eyes, limbs, and distinctive features. Omit all non-visual information like abilities, personality, or backstory.";
    
    console.log('Generating image prompt from concept');
    
    try {
      const output = await rwkvClient.predict("/chat", [
        promptGenerationPrompt, // message
        [],                     // chat_history
        systemPrompt,          // system_prompt
        1024,                   // max_new_tokens
        0.7,                   // temperature
        0.95,                  // top_p
        50,                    // top_k
        1.0                    // repetition_penalty
      ]);
      
      console.log('Image prompt output:', output);
      let promptText = output.data[0];
      
      state.imagePrompt = promptText;
      console.log('Image prompt generated:', state.imagePrompt);
      
      if (!state.imagePrompt || state.imagePrompt.trim() === '') {
        throw new Error('Failed to generate image prompt - received empty response');
      }
    } catch (error) {
      handleAPIError(error);
    }
  }
  
  async function generateMonsterImage() {
    state.currentStep = 'generating';
    
    if (!fluxClient || !state.imagePrompt) {
      throw new Error('Image generation service not available or no prompt');
    }
    
    try {
      const output = await fluxClient.predict("/infer", [
        `${state.imagePrompt}\nNow generate a Pokémon-Anime-style image of the monster in an idle pose with a white background. The monster should not be attacking or in motion. The full monster must be visible within the frame.`,
        0,      // seed
        true,   // randomizeSeed
        1024,   // width
        1024,   // height
        4       // steps
      ]);
      
      const [image, usedSeed] = output.data;
      let url: string | undefined;
      
      if (typeof image === "string") url = image;
      else if (image && image.url) url = image.url;
      else if (image && image.path) url = image.path;
      
      if (url) {
        // Process the image to make white background transparent
        console.log('Processing image for transparency...');
        try {
          const transparentBase64 = await makeWhiteTransparent(url);
          state.monsterImage = {
            imageUrl: url,
            imageData: transparentBase64,
            seed: usedSeed,
            prompt: state.imagePrompt
          };
        } catch (processError) {
          console.error('Failed to process image for transparency:', processError);
          // Fallback to original image
          state.monsterImage = {
            imageUrl: url,
            seed: usedSeed,
            prompt: state.imagePrompt
          };
        }
      } else {
        throw new Error('Failed to generate monster image');
      }
    } catch (error) {
      handleAPIError(error);
    }
  }
  
  async function generateStats() {
    state.currentStep = 'statsGenerating';
    
    if (!rwkvClient || !state.monsterConcept) {
      throw new Error('Text generation service not available or no concept');
    }
    
    const statsPrompt = MONSTER_STATS_PROMPT(state.monsterConcept);
    const systemPrompt = "You are a game designer specializing in monster stats and abilities. You must ONLY output valid JSON that matches the provided schema exactly. Do not include any text before or after the JSON. Do not include null values in your JSON response. Your entire response should be wrapped in a ```json``` code block.";
    
    console.log('Generating monster stats from concept');
    
    try {
      const output = await rwkvClient.predict("/chat", [
        statsPrompt,          // message
        [],                   // chat_history
        systemPrompt,         // system_prompt
        2048,                 // max_new_tokens
        0.3,                  // temperature
        0.95,                 // top_p
        50,                   // top_k
        1.0                   // repetition_penalty
      ]);
      
      console.log('Stats output:', output);
      let jsonString = output.data[0];
      
      // Extract JSON from the response (remove markdown if present)
      let cleanJson = jsonString;
      if (jsonString.includes('```')) {
        const matches = jsonString.match(/```(?:json)?\s*([\s\S]*?)```/);
        if (matches) {
          cleanJson = matches[1];
        } else {
          // If no closing ```, just remove the opening ```json
          cleanJson = jsonString.replace(/^```(?:json)?\s*/, '').replace(/```\s*$/, '');
        }
      }
      
      try {
        // Remove any trailing text after the JSON object
        const jsonMatch = cleanJson.match(/^\s*\{[\s\S]*?\}\s*/);
        if (jsonMatch) {
          cleanJson = jsonMatch[0];
        }
        
        const parsedStats = JSON.parse(cleanJson.trim());
        
        // Remove any extra fields not in our schema
        const allowedFields = ['name', 'description', 'rarity', 'HP', 'defence', 'attack', 'speed',
          'specialPassiveTraitDescription', 'attackActionName', 'attackActionDescription',
          'buffActionName', 'buffActionDescription', 'debuffActionName', 'debuffActionDescription',
          'specialActionName', 'specialActionDescription', 'boostActionName', 'boostActionDescription',
          'disparageActionName', 'disparageActionDescription'];
        
        for (const key in parsedStats) {
          if (!allowedFields.includes(key)) {
            delete parsedStats[key];
          }
        }
        
        // Ensure numeric fields are actually numbers
        const numericFields = ['rarity', 'HP', 'defence', 'attack', 'speed'];
        
        for (const field of numericFields) {
          if (parsedStats[field] !== undefined) {
            // Convert string numbers to actual numbers
            parsedStats[field] = parseInt(parsedStats[field]);
            
            // Clamp to 0-100 range
            parsedStats[field] = Math.max(0, Math.min(100, parsedStats[field]));
          }
        }
        
        // Map field names from schema to interface
        if (parsedStats.specialPassiveTraitDescription) {
          parsedStats.specialPassiveTrait = parsedStats.specialPassiveTraitDescription;
          delete parsedStats.specialPassiveTraitDescription;
        }
        
        // Handle potential old field names from LLM
        if (parsedStats.boostActionName) {
          parsedStats.buffActionName = parsedStats.boostActionName;
          delete parsedStats.boostActionName;
        }
        if (parsedStats.boostActionDescription) {
          parsedStats.buffActionDescription = parsedStats.boostActionDescription;
          delete parsedStats.boostActionDescription;
        }
        if (parsedStats.disparageActionName) {
          parsedStats.debuffActionName = parsedStats.disparageActionName;
          delete parsedStats.disparageActionName;
        }
        if (parsedStats.disparageActionDescription) {
          parsedStats.debuffActionDescription = parsedStats.disparageActionDescription;
          delete parsedStats.disparageActionDescription;
        }
        
        const stats: MonsterStats = parsedStats;
        state.monsterStats = stats;
        console.log('Monster stats generated:', stats);
        console.log('Monster stats JSON:', JSON.stringify(stats, null, 2));
      } catch (parseError) {
        console.error('Failed to parse JSON:', parseError, 'Raw output:', cleanJson);
        throw new Error('Failed to parse monster stats JSON');
      }
    } catch (error) {
      handleAPIError(error);
    }
  }
  
  async function autoSaveMonster() {
    if (!state.monsterImage || !state.imageCaption || !state.monsterConcept || !state.imagePrompt || !state.monsterStats) {
      console.error('Cannot auto-save: missing required data');
      return;
    }
    
    try {
      // Create a clean copy of stats to ensure it's serializable
      const cleanStats = JSON.parse(JSON.stringify(state.monsterStats));
      
      const monsterData = {
        name: state.monsterStats.name,
        imageUrl: state.monsterImage.imageUrl,
        imageData: state.monsterImage.imageData,
        imageCaption: state.imageCaption,
        concept: state.monsterConcept,
        imagePrompt: state.imagePrompt,
        stats: cleanStats
      };
      
      // Check for any non-serializable data
      console.log('Checking monster data for serializability:');
      console.log('- name type:', typeof monsterData.name);
      console.log('- imageUrl type:', typeof monsterData.imageUrl);
      console.log('- imageData type:', typeof monsterData.imageData, monsterData.imageData ? `length: ${monsterData.imageData.length}` : 'null/undefined');
      console.log('- imageCaption type:', typeof monsterData.imageCaption);
      console.log('- concept type:', typeof monsterData.concept);
      console.log('- imagePrompt type:', typeof monsterData.imagePrompt);
      console.log('- stats:', cleanStats);
      
      const id = await saveMonster(monsterData);
      console.log('Monster auto-saved with ID:', id);
    } catch (err) {
      console.error('Failed to auto-save monster:', err);
      console.error('Monster data that failed to save:', {
        name: state.monsterStats?.name,
        hasImageUrl: !!state.monsterImage?.imageUrl,
        hasImageData: !!state.monsterImage?.imageData,
        hasStats: !!state.monsterStats
      });
      // Don't throw - we don't want to interrupt the workflow
    }
  }
  
  function reset() {
    state = {
      currentStep: 'upload',
      userImage: null,
      imageCaption: null,
      monsterConcept: null,
      monsterStats: null,
      imagePrompt: null,
      monsterImage: null,
      error: null,
      isProcessing: false
    };
  }
</script>

<div class="monster-generator">
  
  {#if state.currentStep !== 'upload'}
    <WorkflowProgress currentStep={state.currentStep} error={state.error} />
  {/if}
  
  {#if state.currentStep === 'upload'}
    <UploadStep 
      onImageSelected={handleImageSelected}
      isProcessing={state.isProcessing}
    />
  {:else if state.currentStep === 'complete'}
    <MonsterResult workflowState={state} onReset={reset} />
  {:else}
    <div class="processing-container">
      <div class="spinner"></div>
      <p class="processing-text">
        {#if state.currentStep === 'captioning'}
          Analyzing your image...
        {:else if state.currentStep === 'conceptualizing'}
          Creating monster concept...
        {:else if state.currentStep === 'statsGenerating'}
          Generating battle stats...
        {:else if state.currentStep === 'promptCrafting'}
          Crafting generation prompt...
        {:else if state.currentStep === 'generating'}
          Generating your monster...
        {/if}
      </p>
    </div>
  {/if}
</div>

<style>
  .monster-generator {
    width: 100%;
    max-width: 1200px;
    margin: 0 auto;
    padding: 2rem;
  }
  
  
  .processing-container {
    display: flex;
    flex-direction: column;
    align-items: center;
    padding: 3rem 1rem;
  }
  
  .spinner {
    width: 60px;
    height: 60px;
    border: 3px solid #f3f3f3;
    border-top: 3px solid #007bff;
    border-radius: 50%;
    animation: spin 1s linear infinite;
    margin-bottom: 2rem;
  }
  
  @keyframes spin {
    0% { transform: rotate(0deg); }
    100% { transform: rotate(360deg); }
  }
  
  .processing-text {
    font-size: 1.2rem;
    color: #333;
    margin-bottom: 2rem;
  }
  
</style>