# app.py (исправленная версия)
import gradio as gr
import os
import time
from pathlib import Path
import zipfile
from music_separator_app import MusicSeparatorApp
from config import AUDIO_DIR, OUTPUT_DIR, SUPPORTED_FORMATS

# Global app instance - автоматически инициализируется при создании
separator_app = MusicSeparatorApp()

def create_zip_archive(output_dir):
    """Create ZIP archive of all separated tracks"""
    zip_path = Path(output_dir).with_suffix('.zip')
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for file_path in Path(output_dir).rglob('*'):
            if file_path.is_file():
                zipf.write(file_path, file_path.relative_to(output_dir))
    return str(zip_path)

def get_initial_status():
    """Get initial system status"""
    system_info = separator_app.get_system_info()
    if system_info['initialized']:
        return "✅ Система готова к работе!"
    elif separator_app.initializing:
        return "🔄 Система инициализируется..."
    elif separator_app.init_error:
        return f"❌ Ошибка инициализации: {separator_app.init_error}"
    else:
        return "⚪ Статус неизвестен"

def separate_audio_interface(audio_file, progress=gr.Progress()):
    """Gradio interface function for audio separation"""
    if not audio_file:
        return "❌ Пожалуйста, загрузите аудиофайл", None
    
    try:
        # Check if system is ready
        if not separator_app.initialized and separator_app.initializing:
            return "🔄 Система все еще инициализируется. Пожалуйста, подождите...", None
        elif not separator_app.initialized and separator_app.init_error:
            return f"❌ Система не готова: {separator_app.init_error}", None
        elif not separator_app.initialized:
            return "❌ Система не инициализирована", None

        # Save uploaded file
        audio_path = AUDIO_DIR / f"input_{int(time.time())}{Path(audio_file).suffix}"
        with open(audio_path, "wb") as f:
            with open(audio_file, "rb") as source_file:
                f.write(source_file.read())
        
        # Process audio
        progress(0.3, desc="Обработка аудио...")
        output_dir, all_outputs = separator_app.process_audio(audio_path)
        
        progress(0.8, desc="Создание ZIP архива...")
        
        # Always create ZIP archive
        zip_output = create_zip_archive(output_dir)
        
        progress(1.0, desc="Завершено!")
        
        # Get track list for message
        track_list = "\n".join([f"• {Path(path).name}" for path in all_outputs.values()])
        message = f"✅ Аудио успешно разделено!\n\nСозданные дорожки:\n{track_list}"
        
        return message, zip_output
        
    except Exception as e:
        return f"❌ Ошибка обработки: {str(e)}", None

def get_example_files():
    """Get example files for demo"""
    examples_dir = Path("examples")
    if examples_dir.exists():
        return [str(f) for f in examples_dir.glob("*") if f.suffix.lower() in SUPPORTED_FORMATS]
    return []

# Create Gradio interface
with gr.Blocks(
    title="🎵 Music Source Separator",
    theme=gr.themes.Soft(),
    css="""
    .status-ready { color: #00cc66; font-weight: bold; }
    .status-warning { color: #ff9900; font-weight: bold; }
    .status-error { color: #ff4444; font-weight: bold; }
    .status-neutral { color: #666666; }
    .track-list { 
        background: #f5f5f5; 
        padding: 10px; 
        border-radius: 5px; 
        margin: 10px 0;
    }
    """
) as demo:
    
    gr.Markdown("""
    # 🎵 Music Source Separator
    ## Разделение музыки на отдельные дорожки
    
    *Автоматически создает ZIP-архив с вокалами, инструментами и аккомпанементом*
    """)
    
    # System status display
    with gr.Row():
        status_display = gr.HTML(
            value=f"<div class='status-warning'>{get_initial_status()}</div>",
            label="Статус системы"
        )
    
    with gr.Row():
        with gr.Column(scale=1):
            # Audio input section
            gr.Markdown("### 🎵 Загрузка аудио")
            audio_input = gr.File(
                label="Выберите музыкальный файл",
                file_types=SUPPORTED_FORMATS,
                type="filepath"
            )
            
            process_btn = gr.Button(
                "🎧 Разделить аудио", 
                variant="primary",
                size="lg"
            )
            
        with gr.Column(scale=1):
            # Results section
            gr.Markdown("### 📊 Результаты")
            status_output = gr.Textbox(
                label="Статус обработки",
                interactive=False,
                max_lines=6
            )
            
            zip_output = gr.File(
                label="📦 Скачать ZIP-архив с дорожками",
                file_types=[".zip"],
                type="filepath"
            )
    
    # Examples section
    gr.Markdown("### 🎼 Примеры для тестирования")
    example_files = get_example_files()
    if example_files:
        gr.Examples(
            examples=example_files,
            inputs=[audio_input],
            outputs=[],
            label="Попробуйте на примерах:"
        )
    
    # Track information
    gr.Markdown("""
    ### 🎹 Содержимое ZIP-архива
    
    Каждый обработка создает ZIP-архив со следующими дорожками:
    
    - **vocals.wav** - выделенные вокалы
    - **bass.wav** - бас-линия
    - **drums.wav** - ударные инструменты  
    - **other.wav** - другие инструменты
    - **accompaniment.wav** - полный аккомпанемент (бас + ударные + другие)
    """)
    
    # Technical info
    gr.Markdown("""
    ### 🔧 Техническая информация
    
    - **Модель**: Demucs (Hybrid Transformer Demucs) - современная архитектура для разделения музыки
    - **Точность**: Fine-tuned версия для улучшенного качества
    - **Форматы входные**: WAV, MP3, FLAC, OGG, M4A, AAC
    - **Форматы выходные**: WAV (16-bit PCM)
    - **Поддержка GPU**: Автоматическое определение и использование GPU при доступности
    - **Авто-инициализация**: Модель загружается автоматически при запуске
    
    ### ⚠️ Рекомендации
    
    - Для лучшего качества используйте файлы с битрейтом 192кбит/с и выше
    - Максимальный размер файла: 100MB
    - Время обработки зависит от длины трека и производительности системы
    - Архив содержит 5 отдельных дорожек для профессиональной работы
    """)
    
    # Event handlers
    process_btn.click(
        fn=separate_audio_interface,
        inputs=[audio_input],
        outputs=[status_output, zip_output]
    )

# Launch application
if __name__ == "__main__":
    print("🚀 Запуск Music Source Separator...")
    print("🔄 Начата автоматическая инициализация моделей...")
    
    # Check system status after a short delay to show initialization progress
    time.sleep(2)
    initial_status = get_initial_status()
    print(f"Статус системы: {initial_status}")
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )