App.tsx

import React, { useState, useEffect, useRef } from 'react';
import { APP_NAME } from './constants';
import { JobData, Model, EnhancementData } from './types';
import * as api from './services/api';
import * as db from './services/db';
import * as creditApi from './services/credit';

import CustomVoice from './components/CustomVoice';
import ModelGrid from './components/ModelGrid';
import ModelModal from './components/ModelModal';
import CreateModelModal from './components/CreateModelModal'; 
import History from './components/History';
import BottomNav from './components/BottomNav';
import DeleteModal from './components/DeleteModal';
import UpgradeModal from './components/UpgradeModal';
import FlyToHistory from './components/FlyToHistory';

type TabType = 'home' | 'custom' | 'history';
type SubscriptionStatus = 'free' | 'paid';
type DeleteMode = 'job' | 'model';

interface JobPayload {
    sourceFile?: File;
    text?: string;
    refFile?: File | Blob;
    refFileUrl?: string;
    type: 'custom' | 'model';
    mode: 'vc' | 'tts';
    model?: Model | null;
    pitch: number;
    retryCount: number;
}

export default function App() {
  const [jobs, setJobs] = useState<Record<string, JobData>>({});
  const [customModels, setCustomModels] = useState<Model[]>([]);
  const [activeJobId, setActiveJobId] = useState<string | null>(null);
  const [isLoading, setIsLoading] = useState(false);
  const [activeTab, setActiveTab] = useState<TabType>('home');
  
  const [isModalOpen, setIsModalOpen] = useState(false);
  const [isCreateModelOpen, setIsCreateModelOpen] = useState(false); 
  const [selectedModel, setSelectedModel] = useState<Model | null>(null);
  const [editingModel, setEditingModel] = useState<Model | null>(null); 

  const [deleteModalOpen, setDeleteModalOpen] = useState(false);
  const [deleteTarget, setDeleteTarget] = useState<'all' | string | null>(null);
  const [deleteMode, setDeleteMode] = useState<DeleteMode>('job');

  const [subscriptionStatus, setSubscriptionStatus] = useState<SubscriptionStatus>('free');
  const [upgradeModalOpen, setUpgradeModalOpen] = useState(false);
  const [upgradeModalMode, setUpgradeModalMode] = useState<'limit' | 'feature'>('limit');
  const [resetDate, setResetDate] = useState<number | null>(null);
  const [userFingerprint, setUserFingerprint] = useState<string | null>(null);

  const [showFlyAnimation, setShowFlyAnimation] = useState(false);
  const [lastJobMode, setLastJobMode] = useState<'vc' | 'tts'>('vc');

  const activeJobPayloads = useRef<Record<string, JobPayload>>({});
  const historyBtnRef = useRef<HTMLButtonElement>(null);
  const PREMIUM_PAGE_ID = '1149636';
  
  const FULL_UPGRADE_URL = "#/nav/online/news/getSingle/1149636/eyJpdiI6IkVWcVZ2RFRERVhyaXBxaGtVSHpFK2c9PSIsInZhbHVlIjoiMmJWVllYNnhDNDJjbU96Q0Z5RHdzeE9HWjkrWTVNVzVFTzl3bFlUTzVsWHVTU1pCMWRteUJFeUZ6RVBYUks2VSIsIm1hYyI6IjU3Y2ZjM2IyZDEyNjY4ZTQ3NDgyM2JkMzE0Zjg4Y2ZmZDc5OWQyNTRmYzM4MzBjY2M5ZTNmMmViZDE3N2U5YmQiLCJ0YWciOiIifQ==/20934991";

  useEffect(() => {
    const initUser = async () => {
        const fp = await creditApi.getBrowserFingerprint();
        setUserFingerprint(fp);
        window.parent.postMessage({ type: 'REQUEST_USER_STATUS' }, '*');
    };
    initUser();

    const handleMessage = (event: MessageEvent) => {
        if (event.data?.type === 'USER_STATUS_RESPONSE') {
            try {
                if (event.data.error || !event.data.payload) {
                    setSubscriptionStatus('free');
                    return;
                }
                const userObject = JSON.parse(event.data.payload);
                const isPaid = userObject && userObject.isLogin && userObject.accessible_pages &&
                    (userObject.accessible_pages.includes(PREMIUM_PAGE_ID) || userObject.accessible_pages.includes(parseInt(PREMIUM_PAGE_ID)));
                setSubscriptionStatus(isPaid ? 'paid' : 'free');
            } catch (e) {
                setSubscriptionStatus('free');
            }
        }
    };
    window.addEventListener('message', handleMessage);
    return () => window.removeEventListener('message', handleMessage);
  }, []);

  useEffect(() => {
    const loadData = async () => {
      try {
        const [storedJobs, storedModels] = await Promise.all([
            db.getAllJobs(),
            db.getAllCustomModels()
        ]);
        const jobsMap: Record<string, JobData> = {};
        storedJobs.forEach(job => { jobsMap[job.job_id] = job; });
        setJobs(jobsMap);
        setCustomModels(storedModels);
      } catch (err) {
        console.error("Failed to load data from DB", err);
      }
    };
    loadData();
  }, []);

  const handleJobFailure = async (jobId: string, errorMessage: string) => {
    const payload = activeJobPayloads.current[jobId];
    if (payload && payload.retryCount < 3) {
        payload.retryCount += 1;
        setJobs(prev => {
            const job = prev[jobId]; if (!job) return prev;
            const updated = { ...job, status: 'started' as const, statusMessage: `تلاش مجدد (${payload.retryCount}/3)...`, retryCount: payload.retryCount };
            db.saveJob(updated);
            return { ...prev, [jobId]: updated };
        });
        setTimeout(() => executeJobProcess(jobId, payload), 2000);
    } else {
        delete activeJobPayloads.current[jobId];
        setJobs(prev => {
            const job = prev[jobId]; if (!job) return prev;
            const failedJob = { ...job, status: 'failed' as const, statusMessage: 'خطا: ' + errorMessage };
            db.saveJob(failedJob);
            return { ...prev, [jobId]: failedJob };
        });
    }
  };

  const executeJobProcess = async (tempId: string, data: JobPayload) => {
    try {
        let actualRefFile = data.refFile;
        const isLegacy = data.model?.type === 'legacy_rvc';

        if (!isLegacy && !actualRefFile && data.refFileUrl) {
            actualRefFile = await api.fetchRefAudioAsFile(data.refFileUrl, 'ref.wav');
        }
        
        const finalRef = actualRefFile || new Blob([''], { type: 'audio/wav' });

        const resultJob = await api.processUnifiedJob(
            data.sourceFile, data.text, finalRef, data.model || null, data.mode, data.pitch,
            (msg, progress) => {
                setJobs(prev => {
                    const job = prev[tempId]; if (!job) return prev;
                    return { ...prev, [tempId]: { ...job, statusMessage: msg, progress: progress !== undefined ? progress : job.progress } };
                });
            },
            data.type
        );

        const existingJob = jobs[tempId] || {};
        
        // CRITICAL FIX: Ensure metadata (type, modelName, image) is preserved from existing/payload 
        // if missing or undefined in resultJob to prevent UI reverting to "Custom Project".
        const finalJobData: JobData = { 
            ...existingJob, 
            ...resultJob,
            type: existingJob.type || resultJob.type || data.type,
            modelName: existingJob.modelName || resultJob.modelName || data.model?.name,
            modelImage: existingJob.modelImage || resultJob.modelImage || (typeof data.model?.image === 'string' ? data.model.image : undefined),
            mode: existingJob.mode || resultJob.mode || data.mode
        };
        
        await db.saveJob(finalJobData);
        await db.deleteJob(tempId);

        setJobs(prev => {
            const nextJobs = { ...prev };
            delete nextJobs[tempId];
            nextJobs[finalJobData.job_id] = finalJobData;
            if (finalJobData.status === 'completed' || finalJobData.status === 'failed') {
                delete activeJobPayloads.current[tempId];
            }
            return nextJobs;
        });

    } catch (error: any) {
        handleJobFailure(tempId, error.message);
    }
  };

  useEffect(() => {
    const pollActiveJobs = async () => {
      const activeJobsList = Object.values(jobs) as JobData[];
      const jobsToPoll = activeJobsList.filter(job => 
        (job.status === 'started' || job.status === 'processing') && 
        !job.downloadUrl && !job.job_id.startsWith('temp_')    
      );
      if (jobsToPoll.length === 0) return;

      const updates: Record<string, JobData> = {};
      let hasChanges = false;

      await Promise.all(jobsToPoll.map(async (job) => {
        try {
          const statusRes = await api.checkJobStatus(job);
          
          if (statusRes.status === 'completed' || statusRes.status === 'failed' || statusRes.status === 'error') {
            const isSuccess = statusRes.status === 'completed';
            const finalStatus: JobData['status'] = isSuccess ? 'completed' : 'failed';
            const updatedJob: JobData = { 
                ...job, 
                status: finalStatus, 
                filename: statusRes.filename, 
                downloadUrl: statusRes.downloadUrl, 
                progress: isSuccess ? 100 : job.progress,
                statusMessage: statusRes.detail || (isSuccess ? 'تکمیل شد' : 'خطا در پردازش')
            };
            updates[job.job_id] = updatedJob;
            hasChanges = true;
            db.saveJob(updatedJob);
          } else {
             // Check for progress, detail, OR chunks updates
             let shouldUpdate = false;
             let newChunks = job.chunks;

             // Critical for distributed processing: update chunks if server returns new task IDs
             if (statusRes.chunks && JSON.stringify(statusRes.chunks) !== JSON.stringify(job.chunks)) {
                 newChunks = statusRes.chunks;
                 shouldUpdate = true;
             }
             
             if ((statusRes.detail && statusRes.detail !== job.statusMessage) || (statusRes.progress && statusRes.progress > (job.progress || 0))) {
                 shouldUpdate = true;
             }

             if (shouldUpdate) {
                updates[job.job_id] = { 
                    ...job, 
                    progress: statusRes.progress || job.progress, 
                    statusMessage: statusRes.detail || job.statusMessage,
                    chunks: newChunks
                };
                hasChanges = true;
             }
          }
        } catch (err) { console.error(err); }
      }));
      if (hasChanges) setJobs(prev => ({ ...prev, ...updates }));
    };
    const interval = setInterval(pollActiveJobs, 3000);
    return () => clearInterval(interval);
  }, [jobs]);

  useEffect(() => {
    const pollEnhancementJobs = async () => {
      const enhanceJobsList = (Object.values(jobs) as JobData[]).filter(job => job.enhancement && job.enhancement.status === 'processing');
      if (enhanceJobsList.length === 0) return;
      const updates: Record<string, JobData> = {};
      let hasChanges = false;
      await Promise.all(enhanceJobsList.map(async (job) => {
        const enh = job.enhancement!;
        try {
          const statusRes = await api.checkEnhancementStatus(enh.metadata);
          let nextProgress = enh.progress || 5; 
          if (statusRes.status === 'completed') nextProgress = 100;
          else nextProgress = Math.min(95, nextProgress + 1);

          if (statusRes.status !== enh.status || nextProgress !== enh.progress || statusRes.downloadUrl !== enh.downloadUrl) {
                 const updatedJob = { 
                    ...job, 
                    enhancement: { 
                        ...enh, 
                        status: statusRes.status, 
                        progress: nextProgress, 
                        filename: statusRes.filename || enh.filename,
                        downloadUrl: statusRes.downloadUrl || enh.downloadUrl // ذخیره آدرس مستقیم بله در مرورگر کاربر برای مقاومت در برابر ری‌استارت
                    } 
                 };
                 updates[job.job_id] = updatedJob;
                 hasChanges = true;
                 db.saveJob(updatedJob);
              }
        } catch (err) { console.error(err); }
      }));
      if (hasChanges) setJobs(prev => ({ ...prev, ...updates }));
    };
    const interval = setInterval(pollEnhancementJobs, 3000);
    return () => clearInterval(interval);
  }, [jobs]);

  const startJob = async (data: any) => {
    const isCustomModel = data.type === 'custom' || data.model?.isCustom;
    if (subscriptionStatus === 'free') {
        if (isCustomModel) { setUpgradeModalMode('feature'); setUpgradeModalOpen(true); return; }
        if (userFingerprint) {
            const check = await creditApi.checkCredit(userFingerprint, data.model?.id || 'standard_vc');
            if (check.limit_reached) { setResetDate(check.reset_timestamp); setUpgradeModalMode('limit'); setUpgradeModalOpen(true); return; }
            creditApi.useCredit(userFingerprint, data.model?.id || 'standard_vc');
        }
    }
    const tempJobId = `temp_${Date.now()}`;
    const initialJob: JobData = {
        job_id: tempJobId, status: 'started' as const, progress: 0, statusMessage: 'ارسال درخواست...', total_chunks: 1, chunks: [],
        date: new Date().toLocaleString('fa-IR'), timestamp: Date.now(), 
        type: data.type, mode: data.mode, modelName: data.model?.name, modelImage: typeof data.model?.image === 'string' ? data.model.image : undefined
    };
    setJobs(prev => ({ ...prev, [tempJobId]: initialJob }));
    await db.saveJob(initialJob); 
    setIsModalOpen(false);
    setLastJobMode(data.mode);
    setShowFlyAnimation(true);
    activeJobPayloads.current[tempJobId] = { ...data, retryCount: 0 };
    executeJobProcess(tempJobId, data);
  };

  const handleModelConfirm = (data: { file?: File; text?: string; mode: 'vc' | 'tts'; pitch: number }) => {
    if (!selectedModel) return;
    startJob({ 
      sourceFile: data.file, 
      text: data.text, 
      refFile: typeof selectedModel.refFile !== 'string' ? selectedModel.refFile : undefined, 
      refFileUrl: typeof selectedModel.refFile === 'string' ? selectedModel.refFile : undefined, 
      type: 'model', 
      mode: data.mode, 
      model: selectedModel, 
      pitch: data.pitch 
    });
  };

  const handleCustomStart = (src: File, ref: File) => startJob({ sourceFile: src, refFile: ref, type: 'custom', mode: 'vc', pitch: 0 });

  const handleEnhance = async (job: JobData) => {
    setJobs(prev => ({ ...prev, [job.job_id]: { ...job, enhancement: { status: 'processing' as const, progress: 5 } } }));
    try {
        const audioFile = await api.fetchRefAudioAsFile(job.downloadUrl || `https://ezmarynoori-sada.hf.space/download/${job.filename}`, 'orig.wav');
        const response = await api.uploadEnhancementJob(audioFile);
        const processingJob: JobData = { ...job, enhancement: { status: 'processing' as const, job_id: response.job_id, progress: 5, metadata: response } };
        setJobs(prev => ({ ...prev, [job.job_id]: processingJob }));
        db.saveJob(processingJob);
    } catch (err) { setJobs(prev => ({ ...prev, [job.job_id]: { ...job, enhancement: { status: 'failed' as const, progress: 0 } } })); }
  };

  const handleSaveModel = async (model: Model) => {
    await db.saveCustomModel(model);
    setCustomModels(await db.getAllCustomModels());
    setIsCreateModelOpen(false);
    setEditingModel(null);
  };

  const handleDeleteCustomModel = (id: string) => { 
    setDeleteMode('model'); 
    setDeleteTarget(id); 
    setDeleteModalOpen(true); 
    setIsCreateModelOpen(false); 
  };

  const confirmDelete = async () => {
    if (deleteMode === 'job') {
      if (deleteTarget === 'all') { setJobs({}); await db.clearAllJobs(); }
      else if (typeof deleteTarget === 'string') { setJobs(prev => { const next = { ...prev }; delete next[deleteTarget]; return next; }); await db.deleteJob(deleteTarget); }
    } else if (deleteMode === 'model' && typeof deleteTarget === 'string') { 
      await db.deleteCustomModel(deleteTarget); 
      setCustomModels(await db.getAllCustomModels()); 
      setEditingModel(null); 
    }
    setDeleteModalOpen(false);
  };

  return (
    <div className="min-h-screen bg-[#f8fafc] text-gray-800 font-vazir pb-24 overflow-x-hidden">
      {showFlyAnimation && <FlyToHistory targetRef={historyBtnRef} onComplete={() => setShowFlyAnimation(false)} mode={lastJobMode} />}
      <div className="max-w-lg mx-auto p-4 pt-8">
        <div className="text-center mb-8 relative animate-[fadeIn_0.8s_ease-out]">
            <div className="w-20 h-20 mx-auto mb-5 bg-white rounded-3xl flex items-center justify-center shadow-xl shadow-indigo-500/10 rotate-3 border border-indigo-50 relative group">
                <div className="relative bg-gradient-to-tr from-primary to-secondary w-12 h-12 rounded-xl flex items-center justify-center text-white shadow-lg shadow-primary/40"><i className="fas fa-wave-square text-2xl animate-[pulse_3s_infinite]"></i></div>
            </div>
            <h1 className="text-2xl font-black text-gray-800 mb-3">تغییر صدا با <span className="text-transparent bg-clip-text bg-gradient-to-r from-blue-600 to-purple-600">هوش مصنوعی آلفا</span></h1>
            
            {/* Added subtext from the provided image */}
            <p className="text-gray-500 text-[13px] font-medium mb-6 max-w-[280px] mx-auto leading-relaxed">
                استودیو پیشرفته تبدیل متن به گفتار و تغییر صدای خوانندگان و مشاهیر
            </p>

            <div className={`inline-flex items-center gap-2 px-4 py-1.5 rounded-full text-xs font-bold transition-all ${subscriptionStatus === 'paid' ? 'bg-amber-300 shadow-yellow-500/20' : 'bg-white border border-gray-200'}`}>
                {subscriptionStatus === 'paid' ? <><i className="fas fa-crown"></i><span>نسخه نامحدود</span></> : <><i className="fas fa-leaf text-green-500"></i><span>نسخه رایگان</span></>}
            </div>
        </div>
        {activeTab === 'home' && <ModelGrid customModels={customModels} onSelectModel={(m) => { setSelectedModel(m); setIsModalOpen(true); }} onOpenCustom={() => setActiveTab('custom')} onAddNewModel={() => { if(subscriptionStatus==='paid') setIsCreateModelOpen(true); else { setUpgradeModalMode('feature'); setUpgradeModalOpen(true); } }} onEditModel={(m) => { setEditingModel(m); setIsCreateModelOpen(true); }} />}
        {activeTab === 'custom' && <CustomVoice onStart={handleCustomStart} isLoading={isLoading} />}
        {activeTab === 'history' && <History jobs={Object.values(jobs)} onClearAll={() => { setDeleteMode('job'); setDeleteTarget('all'); setDeleteModalOpen(true); }} onDeleteSingle={(id) => { setDeleteMode('job'); setDeleteTarget(id); setDeleteModalOpen(true); }} onSelect={(j) => setActiveJobId(j.job_id)} onEnhance={handleEnhance} />}
      </div>
      <ModelModal 
        isOpen={isModalOpen} 
        model={selectedModel} 
        onClose={() => setIsModalOpen(false)} 
        onConfirm={handleModelConfirm} 
        isLoading={isLoading} 
        subscriptionStatus={subscriptionStatus}
        onUpgrade={() => { 
          setIsModalOpen(false); 
          // ارسال مستقیم پیام هدایت به بخش خرید اشتراک به برنامه اصلی
          window.parent.postMessage({ type: 'NAVIGATE_TO_PREMIUM', payload: { url: FULL_UPGRADE_URL } }, '*'); 
        }}
      />
      <CreateModelModal isOpen={isCreateModelOpen} editingModel={editingModel} onClose={() => setIsCreateModelOpen(false)} onSave={handleSaveModel} onDelete={handleDeleteCustomModel} />
      <DeleteModal isOpen={deleteModalOpen} type={deleteTarget === 'all' ? 'all' : 'single'} title={deleteMode === 'model' ? 'حذف فایل' : undefined} message={deleteMode === 'model' ? 'آیا از حذف این مدل اختصاصی اطمینان دارید؟' : undefined} onClose={() => setDeleteModalOpen(false)} onConfirm={confirmDelete} />
      <UpgradeModal isOpen={upgradeModalOpen} mode={upgradeModalMode} onClose={() => setUpgradeModalOpen(false)} onUpgrade={() => window.parent.postMessage({ type: 'NAVIGATE_TO_PREMIUM', payload: { url: FULL_UPGRADE_URL } }, '*')} resetDate={resetDate || undefined} />
      <BottomNav activeTab={activeTab} onChange={setActiveTab} badgeCount={(Object.values(jobs) as JobData[]).filter(j=>j.status==='processing'||j.status==='started').length} historyRef={historyBtnRef} />
    </div>
  );
}