# ==============================================================================
# Wymuszona instalacja bibliotek
# ==============================================================================
import os
import subprocess
import sys

def install_packages():
    packages = [
        "gradio>=4.0.0", "transformers>=4.38.1", "torch>=2.2.0", 
        "accelerate>=0.27.2", "diffusers>=0.26.3", "fpdf2", "Pillow", 
        "datasets", "gradio_client", "speechbrain", "soundfile", "moviepy==1.0.3"
    ]
    for package in packages:
        try:
            module_name = package.split('>')[0].split('=')[0].split('<')[0]
            __import__(module_name)
            print(f"{package} jest już zainstalowany.")
        except ImportError:
            print(f"Instalowanie {package}...")
            subprocess.check_call([sys.executable, "-m", "pip", "install", package])

if 'PACKAGES_INSTALLED' not in os.environ:
    install_packages()
    os.environ['PACKAGES_INSTALLED'] = 'TRUE'

# ==============================================================================
# Główny kod aplikacji
# ==============================================================================
import gradio as gr
from transformers import pipeline as text_pipeline
from diffusers import DiffusionPipeline
import torch
import re
from fpdf import FPDF
from PIL import Image
import time
from datasets import Dataset, Features, Value, Sequence, Image as HFImage
from gradio_client import Client
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
import soundfile as sf
import numpy as np
from moviepy.editor import ImageClip, concatenate_videoclips, AudioFileClip, CompositeAudioClip
from speechbrain.pretrained import EncoderClassifier

# --- Konfiguracja ---
LLM_MODEL = "HuggingFaceH4/zephyr-7b-beta"
IMAGE_MODEL = "stabilityai/stable-diffusion-2-1-base"
TTS_MODEL = "microsoft/speecht5_tts"
VOCODER_MODEL = "microsoft/speecht5_hifigan"
SPEAKER_EMBEDDING_MODEL = "speechbrain/spkrec-xvect-voxceleb"
DATASET_PATH = "saved_courses"
TEXT_TO_3D_MODEL_SPACE = "stabilityai/TripoSR"
device = "cuda" if torch.cuda.is_available() else "cpu"

# --- Globalne zmienne dla modeli ---
text_generator, image_generator, client_3d, tts_processor, tts_model, vocoder, speaker_embeddings = (None,) * 7
LLM_LOADED, IMAGE_MODEL_LOADED, CLIENT_3D_LOADED, TTS_LOADED = (False,) * 4

# --- Funkcje ładowania modeli ---
def load_all_models():
    global text_generator, image_generator, client_3d, tts_processor, tts_model, vocoder, speaker_embeddings
    global LLM_LOADED, IMAGE_MODEL_LOADED, CLIENT_3D_LOADED, TTS_LOADED
    
    try:
        print("Ładowanie modelu językowego..."); text_generator = text_pipeline("text-generation", model=LLM_MODEL, torch_dtype=torch.bfloat16, device_map="auto"); LLM_LOADED = True; print("Model językowy załadowany.")
    except Exception as e: print(f"Błąd ładowania LLM: {e}")
    try:
        print("Ładowanie modelu obrazkowego..."); image_generator = DiffusionPipeline.from_pretrained(IMAGE_MODEL); image_generator.to(device); IMAGE_MODEL_LOADED = True; print("Model obrazkowy załadowany.")
    except Exception as e: print(f"Błąd ładowania Image Model: {e}")
    try:
        print("Inicjalizacja klienta 3D..."); client_3d = Client(TEXT_TO_3D_MODEL_SPACE); CLIENT_3D_LOADED = True; print("Klient 3D gotowy.")
    except Exception as e: print(f"Błąd inicjalizacji klienta 3D: {e}")
    try:
        print("Ładowanie modeli TTS...")
        tts_processor = SpeechT5Processor.from_pretrained(TTS_MODEL)
        tts_model = SpeechT5ForTextToSpeech.from_pretrained(TTS_MODEL).to(device)
        vocoder = SpeechT5HifiGan.from_pretrained(VOCODER_MODEL).to(device)
        speaker_model = EncoderClassifier.from_hparams(source=SPEAKER_EMBEDDING_MODEL, savedir=os.path.join('/tmp', SPEAKER_EMBEDDING_MODEL))
        dummy_waveform = torch.randn(1, 16000).to(device)
        speaker_embeddings = speaker_model.encode_batch(dummy_waveform).squeeze().to(device)
        TTS_LOADED = True
        print("Modele TTS załadowane.")
    except Exception as e: print(f"Błąd ładowania modeli TTS: {e}")
    
    # Zwracamy informację do UI, że modele są gotowe
    return gr.update(interactive=True), gr.update(interactive=True)

# --- Funkcje pomocnicze i główne ---
def list_saved_projects():
    if not os.path.exists(DATASET_PATH): return []
    try:
        dataset = Dataset.load_from_disk(DATASET_PATH)
        return dataset["project_name"]
    except Exception as e:
        print(f"Błąd odczytu projektów: {e}"); return []

def parse_course_to_structure(markdown_text):
    course_data = {'title': "Nowy Kurs", 'steps': []}
    title_match = re.search(r"^#\s*(.*)", markdown_text, re.MULTILINE)
    if title_match: course_data['title'] = title_match.group(1).strip()
    step_pattern = re.compile(r"(\d+\.\s*\*\*(.*?)\*\*)\s*([\s\S]*?)(?=\n\d+\.|\Z)", re.MULTILINE)
    steps_found = step_pattern.findall(markdown_text)
    for _, step_title, step_description in steps_found:
        course_data['steps'].append({'title': step_title.strip(), 'description': step_description.strip()})
    return course_data

def generate_course_structure_and_images(topic, progress=gr.Progress(track_tqdm=True)):
    if not LLM_LOADED or not IMAGE_MODEL_LOADED:
        gr.Error("Kluczowe modele AI nie zostały załadowane. Sprawdź logi.")
        return [None, gr.State([])] + [gr.update(visible=False)]*2 + [gr.update(interactive=False)]*3 + [gr.update()]*11

    progress(0, desc="Generowanie tekstu kursu...")
    prompt = f"<|system|>\nJesteś ekspertem w tworzeniu kursów online...</s>\n<|user|>\nTemat kursu: \"{topic}\"</s>\n<|assistant|>"
    response = text_generator(prompt, max_new_tokens=1024, do_sample=True, temperature=0.7, top_p=0.95)
    course_text = response[0]['generated_text'].split('<|assistant|>')[-1].strip()
    course_data = parse_course_to_structure(course_text)
    
    progress(0.2, desc="Generowanie obrazków...")
    generated_images = []
    if course_data['steps']:
        for i, step in enumerate(course_data['steps']):
            progress(0.2 + (i + 1) / len(course_data['steps']) * 0.8, desc=f"Generowanie obrazka dla: {step['title']}")
            image_prompt = f"cinematic photo of '{step['title']}', professional photography, high detail, 8k"
            image = image_generator(image_prompt, num_inference_steps=25).images[0]
            generated_images.append(image)

    updates = [gr.update(value=course_data['title'])]
    for i in range(5):
        if i < len(course_data['steps']):
            updates.extend([gr.update(value=course_data['steps'][i]['title']), gr.update(value=course_data['steps'][i]['description'])])
        else:
            updates.extend(["", ""])

    return generated_images, gr.State(generated_images), gr.update(visible=True), gr.update(visible=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), *updates

def load_project(project_name):
    if not project_name:
        gr.Warning("Wybierz projekt do wczytania.")
        return [None, gr.State([])] + [gr.update(visible=False)]*2 + [gr.update(interactive=False)]*3 + [gr.update()]*11

    dataset = Dataset.load_from_disk(DATASET_PATH)
    project_data = dataset.filter(lambda example: example["project_name"] == project_name)[0]
    
    loaded_images = project_data["images"]
    
    updates = [gr.update(value=project_data['course_title'])]
    for i in range(5):
        if i < len(project_data['steps']):
            step = project_data['steps'][i]
            updates.extend([gr.update(value=step['title']), gr.update(value=step['description'])])
        else:
            updates.extend(["", ""])

    return loaded_images, gr.State(loaded_images), gr.update(visible=True), gr.update(visible=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), *updates

def update_preview(course_title, *args):
    markdown_preview = f"# {course_title}\n\n"
    step_data = list(args)
    for i in range(0, len(step_data), 2):
        step_title, step_description = step_data[i], step_data[i+1]
        if step_title:
            markdown_preview += f"### Krok {i//2 + 1}: {step_title}\n{step_description}\n\n"
    return markdown_preview

def move_step(direction, step_index, images_state, *text_fields):
    if direction == "up":
        if step_index == 0: return [images_state] + [gr.update() for _ in text_fields]
        swap_with_index = step_index - 1
    else:
        if step_index >= 4 or not text_fields[(step_index + 1) * 2]: return [images_state] + [gr.update() for _ in text_fields]
        swap_with_index = step_index + 1

    images_state[step_index], images_state[swap_with_index] = images_state[swap_with_index], images_state[step_index]
    
    current_title, current_desc = text_fields[step_index * 2], text_fields[step_index * 2 + 1]
    swap_title, swap_desc = text_fields[swap_with_index * 2], text_fields[swap_with_index * 2 + 1]

    updates = list(text_fields)
    updates[step_index * 2], updates[step_index * 2 + 1] = swap_title, swap_desc
    updates[swap_with_index * 2], updates[swap_with_index * 2 + 1] = current_title, current_desc
    
    final_updates = [gr.update(value=v) for v in updates]
    return [images_state] + final_updates

def generate_3d_model(prompt_text, progress=gr.Progress(track_tqdm=True)):
    if not CLIENT_3D_LOADED:
        gr.Warning("Klient do generowania modeli 3D nie jest dostępny.")
        return None
    if not prompt_text:
        gr.Info("Wpisz tytuł kroku, aby wygenerować model 3D.")
        return None
    progress(0, desc="Wysyłanie zapytania do modelu 3D... (to może potrwać kilka minut)")
    result_path = client_3d.predict(prompt_text, None, True, 0.5, 1024, 256, api_name="/generate")
    progress(1, desc="Model 3D gotowy!")
    return result_path

def generate_audio(text, progress=gr.Progress(track_tqdm=True)):
    if not TTS_LOADED:
        gr.Warning("Modele Text-to-Speech nie są dostępne.")
        return None
    if not text:
        gr.Info("Brak tekstu do wygenerowania audio.")
        return None
    progress(0, desc="Generowanie audio...")
    inputs = tts_processor(text=text, return_tensors="pt")
    speech = tts_model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.unsqueeze(0), vocoder=vocoder)
    output_path = f"/tmp/audio_{int(time.time())}.wav"
    sf.write(output_path, speech.cpu().numpy(), samplerate=16000)
    progress(1, desc="Audio gotowe!")
    return output_path

def generate_video_from_ui(images_state, course_title, *args, progress=gr.Progress(track_tqdm=True)):
    if not images_state or not course_title:
        gr.Warning("Najpierw wygeneruj kurs, aby stworzyć wideo.")
        return None
    step_data = list(args)
    audio_clips, video_clips, temp_audio_files = [], [], []
    total_steps = len(images_state)
    for i in range(total_steps):
        progress((i + 1) / (total_steps + 1), desc=f"Tworzenie audio dla kroku {i+1}/{total_steps}")
        step_description = step_data[i*2+1]
        if not step_description: continue
        audio_path = generate_audio(step_description)
        if audio_path:
            audio_clip = AudioFileClip(audio_path)
            audio_clips.append(audio_clip)
            temp_audio_files.append(audio_path)
            video_clip = ImageClip(np.array(images_state[i])).set_duration(audio_clip.duration)
            video_clips.append(video_clip)
    if not video_clips:
        gr.Warning("Nie udało się stworzyć żadnych klipów wideo.")
        return None
    progress(0.9, desc="Montowanie finalnego wideo...")
    final_video = concatenate_videoclips(video_clips, method="compose")
    final_video.audio = CompositeAudioClip(audio_clips)
    video_output_path = f"/tmp/final_video_{int(time.time())}.mp4"
    final_video.write_videofile(video_output_path, fps=24, codec="libx24")
    for f in temp_audio_files: os.remove(f)
    progress(1, desc="Wideo gotowe!")
    return video_output_path
    
def save_project(project_name, images_state, course_title, *args):
    if not project_name: return "⚠️ Proszę podać nazwę projektu."
    steps_list = []
    step_data = list(args)
    for i in range(0, len(step_data), 2):
        if step_data[i]: steps_list.append({'title': step_data[i], 'description': step_data[i+1]})
    new_item = {'project_name': project_name, 'course_title': course_title, 'steps': steps_list, 'images': images_state}
    features = Features({'project_name': Value('string'), 'course_title': Value('string'), 'steps': Sequence({'title': Value('string'), 'description': Value('string')}), 'images': Sequence(HFImage())})
    if os.path.exists(DATASET_PATH):
        dataset = Dataset.load_from_disk(DATASET_PATH)
    else:
        dataset = Dataset.from_dict({key: [] for key in features.keys()}, features=features)
    dataset = dataset.add_item(new_item)
    dataset.save_to_disk(DATASET_PATH)
    return f"✅ Projekt '{project_name}' został pomyślnie zapisany!"

def generate_pdf_from_ui(images_state, course_title, *args):
    if not course_title or not images_state: return None
    pdf = FPDF()
    pdf.add_page()
    try:
        pdf.add_font('DejaVu', '', 'DejaVuSans.ttf', uni=True); pdf.set_font('DejaVu', '', 12)
    except RuntimeError:
        print("Brak czcionki DejaVu, używam 'Arial'."); pdf.set_font('Arial', '', 12)
    pdf.set_font_size(24)
    pdf.multi_cell(0, 10, course_title.encode('latin-1', 'replace').decode('latin-1')); pdf.ln(10)
    temp_image_paths = []
    for i, img in enumerate(images_state):
        path = f"/tmp/temp_image_{i}.png"; img.save(path); temp_image_paths.append(path)
    step_data = list(args); image_counter = 0
    for i in range(0, len(step_data), 2):
        step_title, step_description = step_data[i], step_data[i+1]
        if step_title:
            pdf.set_font_size(16)
            pdf.multi_cell(0, 8, f"Krok {i//2 + 1}: {step_title}".encode('latin-1', 'replace').decode('latin-1')); pdf.ln(4)
            if image_counter < len(temp_image_paths):
                pdf.image(temp_image_paths[image_counter], x=None, y=None, w=180); pdf.ln(5); image_counter += 1
            pdf.set_font_size(12)
            pdf.multi_cell(0, 6, step_description.encode('latin-1', 'replace').decode('latin-1')); pdf.ln(10)
    pdf_output_path = f"/tmp/kurs_{int(time.time())}.pdf"; pdf.output(pdf_output_path)
    for path in temp_image_paths: os.remove(path)
    return pdf_output_path

# --- Budowa Interfejsu Gradio ---
custom_css = "body, #root { height: 100vh; margin: 0; padding: 0; } .gradio-container { max-width: 100% !important; } .main, .wrap, .gap { height: 100%; } .app-title { text-align: center; font-size: 2.5em; color: #2c3e50; font-weight: bold; margin-bottom: 0px !important; } .app-subtitle { text-align: center; font-size: 1.1em; color: #576574; margin-top: 5px; margin-bottom: 20px; } .section-title { font-size: 1.5em; font-weight: bold; color: #2980b9; border-bottom: 2px solid #3498db; padding-bottom: 5px; margin-top: 10px; } footer { display: none !important; }"

with gr.Blocks(theme=gr.themes.Soft(primary_hue="sky", secondary_hue="blue"), title="Kreator Kursów AI", css=custom_css) as demo:
    
    gr.HTML("<h1 class='app-title'>🎨 Kreator Kursów AI 🚀</h1>")
    gr.HTML("<p class='app-subtitle'>Twoje centrum dowodzenia do tworzenia multimedialnych kursów DIY</p>")
    
    images_state = gr.State([])

    with gr.Row(equal_height=True):
        with gr.Column(scale=2):
            gr.Markdown("<h2 class='section-title'>⚙️ Panel Sterowania</h2>")
            with gr.Group():
                gr.Markdown("### 1. Wygeneruj lub Wczytaj Projekt")
                topic_input = gr.Textbox(label="Wpisz temat nowego kursu", placeholder="np. Jak zbudować domek dla jeży?", lines=3, interactive=False)
                generate_button = gr.Button("🚀 Wygeneruj Nowy Kurs!", variant="primary", interactive=False)
            with gr.Accordion("📚 Biblioteka Projektów", open=True):
                saved_projects_dropdown = gr.Dropdown(label="Wybierz zapisany projekt", choices=list_saved_projects(), interactive=True)
                load_project_button = gr.Button("📂 Wczytaj Projekt")
            with gr.Group():
                gr.Markdown("### 3. Zapisz Postępy")
                project_name_input = gr.Textbox(label="Nazwa Projektu", placeholder="Mój kurs o jeżach...")
                save_button = gr.Button("💾 Zapisz Projekt", interactive=False)
                save_confirmation_output = gr.Markdown()

        with gr.Column(scale=4):
            gr.Markdown("<h2 class='section-title'>✍️ Główny Edytor</h2>")
            with gr.Group(visible=False) as editor_group:
                course_title_editor = gr.Textbox(label="Tytuł Główny Kursu", interactive=True)
                step_editors = []
                step_text_editors = []
                for i in range(5):
                    with gr.Accordion(f"Krok {i+1}", open=True):
                        with gr.Row():
                            with gr.Column(scale=10):
                                step_title = gr.Textbox(label="Tytuł kroku", interactive=True)
                                step_description = gr.Textbox(label="Opis kroku", lines=5, interactive=True)
                            with gr.Column(scale=1, min_width=80):
                                up_button = gr.Button("🔼")
                                down_button = gr.Button("🔽")
                        with gr.Row():
                            generate_3d_btn = gr.Button("🧊 3D")
                            generate_audio_btn = gr.Button("🎵 Audio")
                        step_editors.extend([step_title, step_description, generate_3d_btn, generate_audio_btn, up_button, down_button])
                        step_text_editors.extend([step_title, step_description])

        with gr.Column(scale=3):
            gr.Markdown("<h2 class='section-title'>👁️ Podgląd i Eksport</h2>")
            with gr.Group(visible=False) as preview_group:
                with gr.Tabs():
                    with gr.TabItem("Tekst"):
                        preview_display = gr.Markdown("Podgląd kursu pojawi się tutaj...")
                    with gr.TabItem("Multimedia"):
                        image_gallery = gr.Gallery(label="Ilustracje", show_label=False, columns=3, height="auto")
                        model_3d_display = gr.Model3D(label="Model 3D", show_label=False)
                        audio_player = gr.Audio(label="Lektor", show_label=False)
                gr.Markdown("### 4. Pobierz Gotowy Kurs")
                with gr.Row():
                    pdf_button = gr.Button("📄 PDF", interactive=False)
                    video_button = gr.Button("🎬 Wideo", interactive=False)
                file_output = gr.File(label="Pobierz PDF")
                video_output = gr.Video(label="Pobierz Wideo", format="mp4")

    # --- Logika Akcji ---
    editor_text_fields = [course_title_editor] + step_text_editors
    all_outputs_for_load_and_generate = [
        image_gallery, images_state, editor_group, preview_group,
        pdf_button, save_button, video_button
    ] + editor_text_fields
    
    demo.load(
        fn=load_all_models,
        inputs=None,
        outputs=[generate_button, topic_input]
    )

    generate_button.click(
        fn=generate_course_structure_and_images,
        inputs=[topic_input],
        outputs=all_outputs_for_load_and_generate
    ).then(fn=update_preview, inputs=editor_text_fields, outputs=preview_display)
    
    load_project_button.click(
        fn=load_project,
        inputs=[saved_projects_dropdown],
        outputs=all_outputs_for_load_and_generate
    ).then(fn=update_preview, inputs=editor_text_fields, outputs=preview_display)
    
    for editor_component in editor_text_fields:
        editor_component.change(fn=update_preview, inputs=editor_text_fields, outputs=preview_display)
    
    pdf_button.click(fn=generate_pdf_from_ui, inputs=[images_state] + editor_text_fields, outputs=file_output)
    save_button.click(fn=save_project, inputs=[project_name_input, images_state] + editor_text_fields, outputs=[save_confirmation_output])

    for i in range(5):
        step_title_comp, step_desc_comp, btn_3d_comp, btn_audio_comp, up_btn, down_btn = step_editors[i*6 : (i+1)*6]
        
        btn_3d_comp.click(fn=generate_3d_model, inputs=[step_title_comp], outputs=[model_3d_display])
        btn_audio_comp.click(fn=generate_audio, inputs=[step_desc_comp], outputs=[audio_player])
        
        up_btn.click(
            fn=move_step,
            inputs=[gr.State("up"), gr.State(i), images_state] + step_text_editors,
            outputs=[images_state] + step_text_editors
        ).then(fn=update_preview, inputs=editor_text_fields, outputs=preview_display)
        
        down_btn.click(
            fn=move_step,
            inputs=[gr.State("down"), gr.State(i), images_state] + step_text_editors,
            outputs=[images_state] + step_text_editors
        ).then(fn=update_preview, inputs=editor_text_fields, outputs=preview_display)
        
    video_button.click(
        fn=generate_video_from_ui,
        inputs=[images_state] + editor_text_fields,
        outputs=[video_output]
    )

demo.queue().launch(show_api=False)