App.tsx

import React, { useState, useCallback, useEffect } from 'react';
import { AppStatus, Provider, AppType } from './types';
import { generateCode } from './services/geminiService';
import ImageUploader from './components/ImageUploader';
import CodeDisplay from './components/CodeDisplay';
import Loader from './components/Loader';
import PythonIcon from './components/icons/PythonIcon';
import CppIcon from './components/icons/CppIcon';
import SparklesIcon from './components/icons/SparklesIcon';
import ApiConfig from './components/ApiConfig';
import AppTypeSelector from './components/AppTypeSelector';

const App: React.FC = () => {
  const [status, setStatus] = useState<AppStatus>(AppStatus.IDLE);
  const [appType, setAppType] = useState<AppType | null>(null);
  const [generatedCode, setGeneratedCode] = useState<string>('');
  const [generatedFiles, setGeneratedFiles] = useState<Record<string, string> | null>(null);
  const [errorMessage, setErrorMessage] = useState<string>('');
  const [imageFiles, setImageFiles] = useState<File[]>([]);
  const [imagePreviewUrls, setImagePreviewUrls] = useState<string[]>([]);
  const [instructions, setInstructions] = useState<string>('');
  const [loadingMessage, setLoadingMessage] = useState<string>('');
  
  const [provider, setProvider] = useState<Provider>('gemini');
  const [geminiApiKey, setGeminiApiKey] = useState<string>('');
  const [openaiApiKey, setOpenaiApiKey] = useState<string>('');
  const [anthropicApiKey, setAnthropicApiKey] = useState<string>('');

  useEffect(() => {
    // Load keys and provider from localStorage on initial render.
    const savedGeminiKey = localStorage.getItem('geminiApiKey');
    const savedOpenaiKey = localStorage.getItem('openaiApiKey');
    const savedAnthropicKey = localStorage.getItem('anthropicApiKey');
    const savedProvider = localStorage.getItem('provider') as Provider;

    if (savedGeminiKey) setGeminiApiKey(savedGeminiKey);
    if (savedOpenaiKey) setOpenaiApiKey(savedOpenaiKey);
    if (savedAnthropicKey) setAnthropicApiKey(savedAnthropicKey);
    if (savedProvider) setProvider(savedProvider);

    // Effect to clean up object URLs
    return () => {
      imagePreviewUrls.forEach(url => URL.revokeObjectURL(url));
    };
  }, []); // Run only once on mount
  
  const resetStateForNewRun = () => {
      setGeneratedCode('');
      setGeneratedFiles(null);
      setErrorMessage('');
      setStatus(AppStatus.IDLE);
  }

  const handleProviderChange = (newProvider: Provider) => {
    setProvider(newProvider);
    localStorage.setItem('provider', newProvider);
  };

  const handleApiKeyChange = (keyType: Provider, key: string) => {
    switch (keyType) {
      case 'gemini':
        setGeminiApiKey(key);
        localStorage.setItem('geminiApiKey', key);
        break;
      case 'openai':
        setOpenaiApiKey(key);
        localStorage.setItem('openaiApiKey', key);
        break;
      case 'anthropic':
        setAnthropicApiKey(key);
        localStorage.setItem('anthropicApiKey', key);
        break;
    }
  };

  const fileToBase64 = (file: File): Promise<{mimeType: string, data: string}> => {
    return new Promise((resolve, reject) => {
      const reader = new FileReader();
      reader.readAsDataURL(file);
      reader.onload = () => {
        const result = reader.result as string;
        const [header, data] = result.split(',');
        const mimeType = header.match(/:(.*?);/)?.[1] || 'application/octet-stream';
        resolve({ mimeType, data });
      };
      reader.onerror = error => reject(error);
    });
  };

  const handleImageUpload = (newFiles: File[]) => {
    const filesToUpload = newFiles.slice(0, 10);
    setImageFiles(filesToUpload);
    // Revoke old URLs before creating new ones to prevent memory leaks
    imagePreviewUrls.forEach(url => URL.revokeObjectURL(url));
    setImagePreviewUrls(filesToUpload.map(file => URL.createObjectURL(file)));
    resetStateForNewRun();
  };

  const handleRemoveImage = (indexToRemove: number) => {
    const updatedFiles = imageFiles.filter((_, index) => index !== indexToRemove);
    URL.revokeObjectURL(imagePreviewUrls[indexToRemove]);
    const updatedUrls = imagePreviewUrls.filter((_, index) => index !== indexToRemove);
    setImageFiles(updatedFiles);
    setImagePreviewUrls(updatedUrls);
  }

  const handleGenerateClick = useCallback(async () => {
    if (!appType) return;
    
    const apiKeys = { gemini: geminiApiKey, openai: openaiApiKey, anthropic: anthropicApiKey };
    const apiKey = apiKeys[provider];

    if (!apiKey) {
      setErrorMessage(`Please provide an API key for ${provider.charAt(0).toUpperCase() + provider.slice(1)}.`);
      setStatus(AppStatus.ERROR);
      return;
    }

    if (imageFiles.length === 0) {
      setErrorMessage('Please upload at least one image first.');
      setStatus(AppStatus.ERROR);
      return;
    }

    setStatus(AppStatus.PROCESSING);
    setErrorMessage('');
    setGeneratedCode('');
    setGeneratedFiles(null);
    setLoadingMessage('Preparing to generate...');

    try {
      const imageParts = await Promise.all(imageFiles.map(file => fileToBase64(file)));
      
      const result = await generateCode(
        appType,
        provider, 
        apiKey, 
        imageParts, 
        instructions,
        (statusMsg) => setLoadingMessage(statusMsg),
        (chunk) => {
             if (appType === 'python') {
                 setGeneratedCode(prev => prev + chunk);
             } else {
                 // For C++, we accumulate the JSON string representation
                 setGeneratedCode(prev => prev + chunk);
             }
        }
      );
      
      if(appType === 'c++') {
          try {
              const finalJsonString = result as string;
              // Clean potential markdown fences
              const fenceRegex = /^```(json)?\s*\n?(.*?)\n?\s*```$/s;
              const match = finalJsonString.match(fenceRegex);
              const cleanedJson = match ? match[2] : finalJsonString;
              const files = JSON.parse(cleanedJson);
              setGeneratedFiles(files);
              setGeneratedCode(''); // Clear the raw string
          } catch(e) {
              console.error("Failed to parse C++ file JSON:", e);
              throw new Error("AI returned an invalid format for C++ files. Please try again.");
          }
      }

      setStatus(AppStatus.SUCCESS);
    } catch (error) {
      console.error(error);
      setErrorMessage(error instanceof Error ? error.message : 'An unknown error occurred.');
      setStatus(AppStatus.ERROR);
    }
  }, [appType, imageFiles, instructions, provider, geminiApiKey, openaiApiKey, anthropicApiKey]);
  
  if (!appType) {
    return <AppTypeSelector onSelect={setAppType} />;
  }
  
  const OutputArea: React.FC = () => {
    switch (status) {
      case AppStatus.PROCESSING:
        return (
          <div className="flex flex-col items-center justify-center h-full text-center p-4">
            <Loader />
            <p className="mt-4 text-lg text-gray-300">{loadingMessage || 'Initializing...'}</p>
            <p className="text-sm text-gray-500 capitalize">Using {provider} to analyze UI and generate code.</p>
            { appType === 'python' && generatedCode && 
              <pre className="mt-4 w-full max-h-64 overflow-y-auto bg-gray-900/50 p-4 rounded-lg text-left">
                <code className="language-python text-xs font-code whitespace-pre-wrap text-gray-400">
                  {generatedCode}
                </code>
              </pre>
            }
          </div>
        );
      case AppStatus.SUCCESS:
        return <CodeDisplay 
            appType={appType}
            pythonCode={generatedCode} 
            cppFiles={generatedFiles}
            status={status} />;
      case AppStatus.ERROR:
        return (
           <div className="flex items-center justify-center h-full">
              <div className="text-center bg-red-900/20 border border-red-600 p-6 rounded-lg">
                  <h3 className="text-2xl font-bold text-red-400">Generation Failed</h3>
                  <p className="mt-2 text-red-300">{errorMessage}</p>
              </div>
          </div>
        );
      case AppStatus.IDLE:
      default:
        return (
          <div className="flex flex-col items-center justify-center h-full text-center text-gray-500 p-8">
            {appType === 'python' ? <PythonIcon className="w-24 h-24 mb-4 opacity-20" /> : <CppIcon className="w-24 h-24 mb-4 opacity-20" />}
            <h3 className="text-2xl font-semibold text-gray-400">Functional Code Output</h3>
            <p className="mt-2 max-w-sm">Configure your API, upload UI images, and provide instructions to generate a functional {appType === 'python' ? 'Python' : 'C++'} script.</p>
          </div>
        );
    }
  };

  const isApiKeyMissing = !( (provider === 'gemini' && geminiApiKey) || (provider === 'openai' && openaiApiKey) || (provider === 'anthropic' && anthropicApiKey) );

  return (
    <div className="min-h-screen bg-gray-900 text-white flex flex-col p-4 sm:p-6 lg:p-8">
      <header className="text-center mb-8">
         <button onClick={() => setAppType(null)} className="text-gray-400 hover:text-white mb-4 text-sm">&larr; Back to selection</button>
        <h1 className="text-4xl sm:text-5xl font-bold bg-gradient-to-r from-green-300 via-blue-400 to-purple-500 text-transparent bg-clip-text">
          Functional {appType === 'python' ? 'PyQt' : 'C++ Qt'} App Generator
        </h1>
        <p className="text-gray-400 mt-2 text-lg">Turn UI Images into Functional {appType === 'python' ? 'Python' : 'C++'} Apps with AI</p>
      </header>

      <main className="flex-grow grid grid-cols-1 lg:grid-cols-2 gap-8">
        <div className="bg-gray-800/50 border border-gray-700 rounded-2xl p-6 flex flex-col shadow-2xl shadow-black/20">
          <div className="flex-grow">
            <ApiConfig 
              provider={provider}
              setProvider={handleProviderChange}
              apiKeys={{ gemini: geminiApiKey, openai: openaiApiKey, anthropic: anthropicApiKey }}
              setApiKey={handleApiKeyChange}
              isProcessing={status === AppStatus.PROCESSING}
            />

            <div>
              <h2 className="text-2xl font-semibold mb-4 text-gray-200">1. Upload UI Images</h2>
              <ImageUploader onImageUpload={handleImageUpload} isProcessing={status === AppStatus.PROCESSING} />
              {imagePreviewUrls.length > 0 && (
                <div className="mt-6 w-full">
                  <p className="text-center mb-3 font-semibold text-gray-300">Image Previews ({imagePreviewUrls.length}/10)</p>
                  <div className="grid grid-cols-2 sm:grid-cols-3 md:grid-cols-4 lg:grid-cols-3 xl:grid-cols-4 gap-4 p-4 border-2 border-dashed border-gray-600 rounded-xl bg-gray-900/50 max-h-80 overflow-y-auto">
                    {imagePreviewUrls.map((url, index) => (
                      <div key={url} className="relative group aspect-square">
                        <img src={url} alt={`UI Preview ${index + 1}`} className="w-full h-full object-contain rounded-lg bg-gray-800" />
                        <button
                          onClick={() => handleRemoveImage(index)}
                          disabled={status === AppStatus.PROCESSING}
                          className="absolute -top-2 -right-2 bg-red-600 text-white rounded-full w-6 h-6 flex items-center justify-center text-sm font-bold opacity-0 group-hover:opacity-100 transition-opacity duration-200 hover:bg-red-700 disabled:opacity-50 disabled:cursor-not-allowed"
                          aria-label={`Remove image ${index + 1}`}
                        >
                          &times;
                        </button>
                      </div>
                    ))}
                  </div>
                </div>
              )}
            </div>
            <div className="mt-8">
              <h2 className="text-2xl font-semibold mb-4 text-gray-200">2. Add Instructions (Optional)</h2>
              <textarea
                value={instructions}
                onChange={(e) => setInstructions(e.target.value)}
                disabled={status === AppStatus.PROCESSING}
                placeholder="e.g., 'Add a dark mode toggle', 'Make the primary button green'"
                className="w-full h-28 p-3 bg-gray-900/70 border-2 border-gray-600 rounded-xl text-gray-300 placeholder-gray-500 focus:ring-2 focus:ring-blue-500 focus:border-blue-500 transition-all duration-200 resize-none"
                aria-label="Refinement instructions"
              />
            </div>
          </div>
          <button
            onClick={handleGenerateClick}
            disabled={imageFiles.length === 0 || status === AppStatus.PROCESSING || isApiKeyMissing}
            className="w-full mt-6 py-4 px-6 rounded-xl bg-gradient-to-r from-blue-500 to-purple-600 text-white font-bold text-lg shadow-lg hover:shadow-xl hover:from-blue-600 hover:to-purple-700 transition-all duration-300 disabled:opacity-50 disabled:cursor-not-allowed disabled:shadow-none flex items-center justify-center gap-3"
          >
            <SparklesIcon className="w-6 h-6"/>
            {status === AppStatus.PROCESSING ? 'Generating...' : '3. Generate Code'}
          </button>
        </div>
        
        <div className="bg-gray-800/50 border border-gray-700 rounded-2xl flex flex-col shadow-2xl shadow-black/20 overflow-hidden">
          <h2 className="text-2xl font-semibold p-6 text-gray-200 border-b border-gray-700">Get Your {appType === 'python' ? 'Python' : 'C++'} Code</h2>
          <div className="flex-grow bg-gray-900/70 overflow-y-auto">
            <OutputArea />
          </div>
        </div>
      </main>
       <footer className="text-center mt-8 text-gray-500 text-sm">
        <p>Powered by AI. Generated code may require adjustments.</p>
      </footer>
    </div>
  );
};

export default App;