script.js

const promptInput = document.getElementById('promptInput');
const negativePromptInput = document.getElementById('negativePromptInput');
const seedInput = document.getElementById('seedInput');
const styleSelect = document.getElementById('styleSelect');
const renderQualitySelect = document.getElementById('renderQualitySelect');
const lightingSelect = document.getElementById('lightingSelect');
const controlNetSelect = document.getElementById('controlNetSelect');
const anatomyControlSelect = document.getElementById('anatomyControlSelect');
const generateButton = document.getElementById('generateButton');
const generatedImage = document.getElementById('generatedImage');
const loadingIndicator = document.getElementById('loading');
const technicalExplanation = document.querySelector('.technical-explanation');

// Prompt Enhancer Module
const promptEnhancerModule = {
  enhancements: {
    "a cat": "a photorealistic fluffy white cat sitting on a windowsill, soft light, shallow depth of field, 4k resolution",
    "a fox": "a majestic red fox standing in a misty forest, golden hour lighting, detailed fur, realistic anatomy, 8k resolution",
    "a landscape": "a breathtaking mountain landscape with vibrant colors, sunrise lighting, crisp details, epic panoramic view, high resolution"
  },

  enhance: async function(originalPrompt) {
    // First, use the new advanced prompt engineering module
    const semanticAnalysis = await promptEngineeringModule.analyzeAndCorrectPrompt(originalPrompt);
  
    // Then enhance with style and quality details
    const enhancedPrompt = promptEngineeringModule.enhancePrompt(semanticAnalysis, {
      style: styleSelect.value,
      renderQuality: renderQualitySelect.value,
      lighting: lightingSelect.value
    });

    return enhancedPrompt;
  }
};

// Advanced Prompt Engineering Module
const promptEngineeringModule = {
  // Enhanced prompt generation with more specific details
  enhancePrompt: function(originalPrompt, options = {}) {
    const {
      style = 'default', 
      renderQuality = 'standard', 
      lighting = 'default',
      detailLevel = 'high'
    } = options;

    const styleMap = {
      'default': '',
      'photorealistic': 'photorealistic, high detail, sharp focus',
      'anime': 'anime style, clean lines, vibrant colors',
      'digital-art': 'digital art, crisp edges, stylized',
      'studio-ghibli': 'Studio Ghibli style, soft colors, whimsical',
      'oil-painting': 'oil painting texture, rich brush strokes',
      'realistic': 'hyper-realistic, ultra-detailed',
      'dreamscape': 'surreal, dreamy, soft focus'
    };

    const lightingMap = {
      'default': '',
      'cinematic': 'cinematic lighting, dramatic shadows',
      'soft': 'soft, diffused light, gentle shadows',
      'dramatic': 'high contrast, strong directional light',
      'sunset': 'golden hour, warm tones, long shadows'
    };

    const qualityMap = {
      'standard': '512x512 resolution',
      'high': '768x768 resolution, enhanced details',
      'ultra': '1024x1024 resolution, maximum detail',
      '4k': '2048x2048 resolution, ultra high quality'
    };

    const detailLevelMap = {
      'low': 'basic shapes, minimal detail',
      'medium': 'clear definition, moderate detail',
      'high': 'intricate details, sharp edges, textural complexity'
    };

    // Combine enhanced prompt with style and quality specifications
    const enhancedPrompt = [
      originalPrompt,
      styleMap[style],
      lightingMap[lighting],
      qualityMap[renderQuality],
      `${detailLevelMap[detailLevel]}, precise composition`
    ].filter(Boolean).join(', ');

    return enhancedPrompt;
  },

  // Advanced semantic analysis and prompt correction
  analyzeAndCorrectPrompt: async function(prompt) {
    try {
      const completion = await websim.chat.completions.create({
        messages: [
          {
            role: "system",
            content: `You are an expert prompt engineer for AI image generation.
            Analyze the given prompt and:
            1. Improve clarity and specificity
            2. Add missing descriptive details
            3. Ensure semantic consistency
            4. Optimize for better image generation`
          },
          {
            role: "user",
            content: prompt
          }
        ]
      });

      return completion.content || prompt;
    } catch (error) {
      console.warn("Prompt analysis fallback triggered:", error);
      return prompt;
    }
  }
};

// Semantic Consistency Check Module
const semanticConsistencyModule = {
  checks: [
    {
      name: "Anatomical Consistency",
      check: async (prompt, imageUrl) => {
        try {
          const result = await websim.chat.completions.create({
            messages: [
              {
                role: "system",
                content: "Analyze the image for anatomical correctness based on the prompt. Report any inconsistencies concisely."
              },
              {
                role: "user",
                content: [
                  { type: "text", text: `Check anatomical correctness for prompt: ${prompt}` },
                  { type: "image_url", image_url: { url: imageUrl } }
                ]
              }
            ]
          });
          return `Anatomical Check: ${result.content}`;
        } catch (error) {
          console.error("Anatomical consistency check failed:", error);
          return `Anatomical Check: Error performing check - ${error.message}`;
        }
      }
    },
    {
      name: "Prompt Adherence Check",
      check: async (prompt, imageUrl) => {
        try {
          const result = await websim.chat.completions.create({
            messages: [
              {
                role: "system",
                content: "Compare the image against the provided prompt. Summarize how well the image matches the key elements of the prompt."
              },
              {
                role: "user",
                content: [
                  { type: "text", text: `Check adherence to prompt: ${prompt}` },
                  { type: "image_url", image_url: { url: imageUrl } }
                ]
              }
            ]
          });
          return `Prompt Adherence: ${result.content}`;
        } catch (error) {
          console.error("Prompt adherence check failed:", error);
          return `Prompt Adherence: Error performing check - ${error.message}`;
        }
      }
    },
    {
      name: "Negative Prompt Check",
      check: async (prompt, negativePrompt, imageUrl) => {
        if (!negativePrompt || negativePrompt.trim() === '') {
          return "Negative Prompt Check: No negative prompt provided.";
        }
        try {
          const result = await websim.chat.completions.create({
            messages: [
              {
                role: "system",
                content: "Analyze the image to see if it contains elements explicitly mentioned in the negative prompt."
              },
              {
                role: "user",
                content: [
                  { type: "text", text: `Check negative prompt: ${prompt} with negative prompt: ${negativePrompt}` },
                  { type: "image_url", image_url: { url: imageUrl } }
                ]
              }
            ]
          });
          return `Negative Prompt Check: ${result.content}`;
        } catch (error) {
          console.error("Negative prompt check failed:", error);
          return `Negative Prompt Check: Error performing check - ${error.message}`;
        }
      }
    }
  ],

  performChecks: async function(prompt, negativePrompt, imageUrl) {
    const results = [];
    // Clear previous results while checks are running
    const semanticCheckList = document.getElementById('semanticCheckList');
    if (semanticCheckList) {
      semanticCheckList.innerHTML = '<li>Running semantic checks...</li>';
    }

    for (let check of this.checks) {
      let checkResult;
      if (check.name === "Negative Prompt Check") {
        checkResult = await check.check(prompt, negativePrompt, imageUrl);
      } else {
        checkResult = await check.check(prompt, imageUrl);
      }
      results.push(checkResult);
    }
    return results;
  }
};

// Layered Generation Module (basic placeholder)
const layeredGenerationModule = {
  generateLayered: async function(prompt) {
    // Placeholder for future implementation
    console.log("Layered generation for:", prompt);
    return null;
  }
};

const styleMap = {
  'default': {
    prompt_prefix: '',
    model: 'stable-diffusion-xl',
    sampler: 'Euler a'
  },
  'photorealistic': {
    prompt_prefix: 'highly detailed photorealistic, sharp focus, natural lighting',
    model: 'realistic-vision-v5',
    sampler: 'DPM++ 2M Karras'
  },
  'anime': {
    prompt_prefix: 'anime style, clean lines, vibrant colors, sharp details, kyoto animation aesthetic',
    model: 'deliberate-v2',
    sampler: 'Euler a'
  },
  'digital-art': {
    prompt_prefix: 'digital illustration, crisp edges, stylized, graphic novel aesthetic, high contrast',
    model: 'dreamshaper-v8',
    sampler: 'DDIM'
  },
  'studio-ghibli': {
    prompt_prefix: 'Studio Ghibli style, soft color palette, whimsical, hand-drawn texture, Miyazaki-inspired',
    model: 'ghibli-diffusion',
    sampler: 'Euler'
  },
  'oil-painting': {
    prompt_prefix: 'oil painting texture, rich brush strokes, impasto technique, canvas-like rendering',
    model: 'anything-v5',
    sampler: 'DPM++ SDE Karras'
  },
  'realistic': {
    prompt_prefix: 'ultra-realistic, hyper-detailed, natural textures, precise anatomical accuracy',
    model: 'realistic-vision-v5',
    sampler: 'DPM++ 2M Karras'
  },
  'dreamscape': {
    prompt_prefix: 'surreal, dreamy atmosphere, soft focus, ethereal lighting, blended reality',
    model: 'stable-diffusion-xl',
    sampler: 'Euler a'
  }
};

function mapRenderQuality(quality) {
  const qualityMap = {
    "standard": {
      width: 512,
      height: 512,
      description: "Standard Quality (512x512)"
    },
    "high": {
      width: 768,
      height: 768,
      description: "High Quality (768x768)"
    },
    "ultra": {
      width: 1024,
      height: 1024,
      description: "Ultra HD (1024x1024)"
    },
    "4k": {
      width: 2048,
      height: 2048,
      description: "4K Ultra-Detailed (2048x2048)"
    }
  };
  return qualityMap[quality] || qualityMap["standard"];
}

// Advanced Face Generation Module
const faceGenerationModule = {
  // Advanced prompt engineering for high-quality face generation
  enhanceFacePrompt: function(basePrompt, options = {}) {
    const {
      faceQuality = 'high',
      symmetry = true,
      style = 'photorealistic'
    } = options;

    const qualityPrefixes = {
      'low': 'basic face details',
      'medium': 'clear facial features, moderate detail',
      'high': 'perfect symmetrical face, high resolution, soft studio lighting, crisp details',
      'ultra': 'hyper-realistic face, professional DSLR photo quality, cinematic lighting, intricate skin texture'
    };

    const stylePrefixes = {
      'photorealistic': 'DSLR photo style, natural skin tones, precise facial structure',
      'anime': 'anime style, clean lines, large expressive eyes',
      'artistic': 'stylized face, painterly details, soft color palette',
      'sketch': 'pencil sketch style, delicate line work'
    };

    const symmetryPrefix = symmetry ? 'perfectly symmetrical, ' : '';

    const enhancedPrompt = [
      basePrompt,
      symmetryPrefix,
      qualityPrefixes[faceQuality],
      stylePrefixes[style]
    ].filter(Boolean).join(', ');

    return enhancedPrompt;
  },

  // Advanced face correction using AI techniques
  correctFace: async function(imageUrl) {
    try {
      // Simulate using GFPGAN or CodeFormer for face enhancement
      const result = await websim.chat.completions.create({
        messages: [
          {
            role: "system",
            content: "Analyze and correct facial details in the image. Improve symmetry, skin texture, and overall facial structure. Provide a URL to the corrected image or indicate if no correction was needed."
          },
          {
            role: "user",
            content: [
              { type: "text", text: "Enhance facial details" },
              { type: "image_url", image_url: { url: imageUrl } }
            ]
          }
        ]
      });

      try {
          const response = JSON.parse(result.content);
          if (response && response.url) {
              console.log("Face correction applied. New image URL:", response.url);
              return response.url; 
          } else {
              console.log("Face correction analysis:", result.content);
              return imageUrl; 
          }
      } catch (e) {
          console.log("Face correction analysis (non-JSON response):", result.content);
          return imageUrl; 
      }

    } catch (error) {
      console.error("Face correction failed:", error);
      return null;
    }
  },

  // Advanced semantic face analysis
  analyzeFace: async function(imageUrl) {
    try {
      const result = await websim.chat.completions.create({
        messages: [
          {
            role: "system",
            content: `Perform comprehensive face analysis:
            1. Check facial symmetry
            2. Verify anatomical correctness
            3. Assess style adherence
            4. Detect any artifacts or inconsistencies
            5. Provide a concise summary.`
          },
          {
            role: "user",
            content: [
              { type: "text", text: "Analyze face in image" },
              { type: "image_url", image_url: { url: imageUrl } }
            ]
          }
        ]
      });

      return result.content;
    } catch (error) {
      console.error("Face analysis failed:", error);
      return null;
    }
  }
};

async function generateImage(prompt, options = {}) {
  const {
    negativePrompt = '',
    seed,
    style = 'default',
    renderQuality = 'standard',
    lighting,
    controlNetMode,
    faceQuality,
    faceStyle,
    symmetry,
    perspective,
    objectPlacement,
    lightSource,
    depthPerception,
    cfgScale,
    steps
  } = options;

  // Get style-specific configuration
  const styleConfig = styleMap[style] || styleMap['default'];

  // Combine style-specific prefix with original prompt
  let enhancedPrompt = `${styleConfig.prompt_prefix}, ${prompt}`;

  // If face generation is detected, enhance the prompt
  const isFacePrompt = prompt.toLowerCase().includes('face') || 
                       prompt.toLowerCase().includes('portrait');
  
  if (isFacePrompt) {
    enhancedPrompt = faceGenerationModule.enhanceFacePrompt(enhancedPrompt, {
      faceQuality: faceQuality || 'high',
      symmetry: symmetry !== false,
      style: faceStyle || 'photorealistic'
    });
  }

  // Enhanced Perspective Mapping
  const perspectiveMap = {
    'front-view': 'viewed straight on, front perspective',
    'side-view': 'side angle view, profile perspective',
    'overhead': 'viewed from above, overhead perspective',
    'diagonal': 'diagonal angle, dynamic perspective',
    'tilted-view': 'slightly tilted view, subtle angle'
  };

  const objectPlacementMap = {
    'balanced': 'balanced composition, symmetric arrangement',
    'asymmetric': 'asymmetric layout, dynamic positioning',
    'rule-of-thirds': 'following rule of thirds, strategic placement',
    'centered': 'centered focal point, symmetrical design'
  };

  const lightSourceMap = {
    'left-side': 'light source from left side, creating defined shadows',
    'right-side': 'light source from right side, dramatic shadowing',
    'top-down': 'overhead lighting, creating vertical shadows',
    'diagonal': 'diagonal light, creating depth and dimension',
    'backlighting': 'backlight effect, silhouette and rim lighting',
    'soft-ambient': 'soft, diffused ambient lighting'
  };

  const depthPerceptionMap = {
    'shallow-depth': 'shallow depth of field, soft background blur',
    'deep-depth': 'deep depth of field, entire scene in focus',
    'layered-focus': 'layered focus, multiple depth levels',
    'foreground-emphasis': 'foreground elements sharply focused'
  };

  // Enhance prompt with perspective and composition details
  if (perspective) {
    enhancedPrompt += `, ${perspectiveMap[perspective]}`;
  }

  if (objectPlacement) {
    enhancedPrompt += `, ${objectPlacementMap[objectPlacement]}`;
  }

  if (lightSource) {
    enhancedPrompt += `, ${lightSourceMap[lightSource]}`;
  }

  if (depthPerception) {
    enhancedPrompt += `, ${depthPerceptionMap[depthPerception]}`;
  }

  // Use the enhanced mapRenderQuality function
  const qualitySettings = mapRenderQuality(renderQuality);

  try {
    const result = await websim.imageGen({
      prompt: enhancedPrompt,
      negative_prompt: getCombinedNegativePrompt(negativePrompt),
      seed: seed,
      style: style === 'default' ? undefined : style,
      width: qualitySettings.width,
      height: qualitySettings.height,
      lighting: lighting === 'default' ? undefined : lighting,
      controlnet_mode: controlNetMode === 'none' ? undefined : controlNetMode,
      sampler: styleConfig.sampler,
      model: styleConfig.model,
      cfg_scale: cfgScale,
      steps: steps,
      perspective: perspective,
      object_placement: objectPlacement,
      light_source: lightSource,
      depth_perception: depthPerception
    });

    console.log(`Generating image with ${qualitySettings.description}`);

    // Post-generation face correction
    if (isFacePrompt && result.url) {
      const correctedImageUrl = await faceGenerationModule.correctFace(result.url);
      result.url = correctedImageUrl || result.url;

      // Optional face analysis
      const faceAnalysis = await faceGenerationModule.analyzeFace(result.url);
      console.log("Face Analysis:", faceAnalysis);
    }

    return result;
  } catch (error) {
    console.error("Advanced Image Generation Error:", error);
    throw error;
  }
}

function getCombinedNegativePrompt(userNegativePrompt) {
  const baseAnatomyPreventionPrompts = [
    "deformed anatomy",
    "extra limbs",
    "duplicate body parts",
    "unnatural proportions",
    "malformed features",
    "unrealistic anatomy",
    "missing limbs",
    "mutated",
    "twisted",
    "bad eyes",
    "extra fingers",
    "missing fingers"
  ];

  const baseQualityPreventions = [
    "low quality",
    "blurry",
    "watermark",
    "text",
    "disfigured",
    "ugly",
    "nonsensical",
    "bad composition",
    "poorly drawn",
    "sketch",
    "noise",
    "grainy"
  ];

  const textSpecificPreventions = [
    "unreadable text",
    "gibberish",
    "nonsense text",
    "broken letters",
    "random symbols"
  ];

  let combinedPrompts = [...baseAnatomyPreventionPrompts, ...baseQualityPreventions, ...textSpecificPreventions];

  if (userNegativePrompt) {
    combinedPrompts = combinedPrompts.concat(userNegativePrompt.split(',').map(p => p.trim()));
  }

  const uniquePrompts = [...new Set(combinedPrompts)].filter(p => p !== '');

  return uniquePrompts.join(", ");
}

async function performSemanticChecks(prompt, imageUrl) {
  const semanticCheckList = document.getElementById('semanticCheckList');
  semanticCheckList.innerHTML = '<li>Running advanced semantic checks...</li>';

  try {
    const result = await websim.chat.completions.create({
      messages: [
        {
          role: "system",
          content: `Perform a comprehensive semantic analysis of the image:
          1. Check anatomical correctness
          2. Verify prompt adherence
          3. Detect any inconsistencies or artifacts
          4. Suggest potential improvements`
        },
        {
          role: "user",
          content: [
            { type: "text", text: `Analyze image generated from prompt: ${prompt}` },
            { type: "image_url", image_url: { url: imageUrl } }
          ]
        }
      ]
    });

    semanticCheckList.innerHTML = result.content
      .split('\n')
      .map(line => `<li>${line}</li>`)
      .join('');
  } catch (error) {
    console.error("Advanced Semantic Check Failed:", error);
    semanticCheckList.innerHTML = `<li>Advanced semantic check failed: ${error.message}</li>`;
  }
}

const faceQualitySelect = document.createElement('select');
faceQualitySelect.id = 'faceQualitySelect';
faceQualitySelect.innerHTML = `
  <option value="low">Basic Quality</option>
  <option value="medium">Medium Quality</option>
  <option value="high" selected>High Quality</option>
  <option value="ultra">Ultra Quality</option>
`;

const faceStyleSelect = document.createElement('select');
faceStyleSelect.id = 'faceStyleSelect';
faceStyleSelect.innerHTML = `
  <option value="photorealistic" selected>Photorealistic</option>
  <option value="anime">Anime Style</option>
  <option value="artistic">Artistic</option>
  <option value="sketch">Sketch Style</option>
`;

const perspectiveSelect = document.createElement('select');
perspectiveSelect.id = 'perspectiveSelect';
perspectiveSelect.innerHTML = `
  <option value="front-view">Front View</option>
  <option value="side-view">Side View</option>
  <option value="overhead">Overhead View</option>
  <option value="diagonal">Diagonal View</option>
  <option value="tilted-view">Tilted View</option>
`;

const objectPlacementSelect = document.createElement('select');
objectPlacementSelect.id = 'objectPlacementSelect';
objectPlacementSelect.innerHTML = `
  <option value="balanced">Balanced Composition</option>
  <option value="asymmetric">Asymmetric Layout</option>
  <option value="rule-of-thirds">Rule of Thirds</option>
  <option value="centered">Centered Focal Point</option>
`;

const lightSourceSelect = document.createElement('select');
lightSourceSelect.id = 'lightSourceSelect';
lightSourceSelect.innerHTML = `
  <option value="left-side">Left Side Lighting</option>
  <option value="right-side">Right Side Lighting</option>
  <option value="top-down">Top-Down Lighting</option>
  <option value="diagonal">Diagonal Lighting</option>
  <option value="backlighting">Backlighting</option>
  <option value="soft-ambient">Soft Ambient Lighting</option>
`;

const depthPerceptionSelect = document.createElement('select');
depthPerceptionSelect.id = 'depthPerceptionSelect';
depthPerceptionSelect.innerHTML = `
  <option value="shallow-depth">Shallow Depth of Field</option>
  <option value="deep-depth">Deep Depth of Field</option>
  <option value="layered-focus">Layered Focus</option>
  <option value="foreground-emphasis">Foreground Emphasis</option>
`;

document.querySelector('.style-grid').appendChild(
  createControlGroup('🖼️ Face Quality:', faceQualitySelect)
);
document.querySelector('.style-grid').appendChild(
  createControlGroup('🎨 Face Style:', faceStyleSelect)
);
document.querySelector('.style-grid').appendChild(
  createControlGroup('📸 Perspective:', perspectiveSelect)
);
document.querySelector('.style-grid').appendChild(
  createControlGroup('🖌️ Object Placement:', objectPlacementSelect)
);
document.querySelector('.style-grid').appendChild(
  createControlGroup('💡 Light Source:', lightSourceSelect)
);
document.querySelector('.style-grid').appendChild(
  createControlGroup('🔍 Depth Perception:', depthPerceptionSelect)
);

function createControlGroup(labelText, selectElement) {
  const container = document.createElement('div');
  container.className = 'style-item';
  
  const label = document.createElement('label');
  label.textContent = labelText;
  
  container.appendChild(label);
  container.appendChild(selectElement);
  
  return container;
}

generateButton.addEventListener('click', async () => {
  const userPrompt = promptInput.value.trim();

  if (!userPrompt) {
    alert('Please enter a prompt to generate an image!');
    return;
  }

  // Show loading indicator
  loadingIndicator.classList.remove('hidden');
  // Hide previous image
  generatedImage.src = "";
  // Clear previous semantic check results
  const semanticCheckList = document.getElementById('semanticCheckList');
  if (semanticCheckList) {
    semanticCheckList.innerHTML = '';
  }
  document.getElementById('semanticCheckResult').classList.add('hidden');


  try {
    // Enhanced image generation with more control
    const imageGenerationOptions = {
      negativePrompt: negativePromptInput.value.trim(),
      seed: seedInput.value ? parseInt(seedInput.value, 10) : undefined,
      style: styleSelect.value,
      renderQuality: renderQualitySelect.value,
      lighting: lightingSelect.value,
      controlNetMode: controlNetSelect.value,
      faceQuality: faceQualitySelect.value,
      faceStyle: faceStyleSelect.value,
      symmetry: true,
      perspective: perspectiveSelect.value,
      objectPlacement: objectPlacementSelect.value,
      lightSource: lightSourceSelect.value,
      depthPerception: depthPerceptionSelect.value,
      cfgScale: parseFloat(document.getElementById('cfgScaleInput').value), 
      steps: parseInt(document.getElementById('stepsInput').value, 10) 
    };

    const result = await generateImage(userPrompt, imageGenerationOptions);

    if (result && result.url) {
      generatedImage.src = result.url;

      // Additional semantic checks and enhancements
      // Show semantic check results panel
      document.getElementById('semanticCheckResult').classList.remove('hidden');
      await performSemanticChecks(userPrompt, result.url);
    }
  } catch (error) {
    console.error("Image Generation Process Failed:", error);
    alert("Image generation failed: " + error.message);
    generatedImage.src = ""; 
  } finally {
    // Hide loading indicator
    loadingIndicator.classList.add('hidden');
  }
});