import gradio as gr
import torch
import time
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

print("🚀 ЗАГРУЖАЮ ОПТИМИЗИРОВАННУЮ СИСТЕМУ...")

# ================== ИСПОЛЬЗУЕМ ТОЛЬКО TINYLLAMA ==================
# Она быстрая, уже работает, занимает мало памяти
MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

# Глобальные переменные для кэша
_tokenizer = None
_model = None
_translator = None

def load_system():
    """Загружаем только одну модель - TinyLlama"""
    global _tokenizer, _model, _translator
    
    print("📦 Загружаю TinyLlama (самая быстрая для CPU)...")
    
    try:
        # Токенизатор
        _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        _tokenizer.pad_token = _tokenizer.eos_token
        
        # Модель с оптимизациями для CPU
        _model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.float32,  # Используем float32 для CPU
            device_map="auto",
            low_cpu_mem_usage=True,
            offload_folder="./offload"  # Выгружаем на диск если не хватает RAM
        )
        
        print("✅ TinyLlama загружена и оптимизирована для CPU!")
        
        # Переводчик
        print("🌐 Загружаю переводчик...")
        _translator = pipeline(
            "translation_ru_to_en", 
            model="Helsinki-NLP/opus-mt-ru-en",
            device="cpu"
        )
        print("✅ Переводчик готов!")
        
        return True
        
    except Exception as e:
        print(f"❌ Ошибка загрузки: {e}")
        return False

# ================== ПРАВИЛЬНЫЕ ПРОМТ-ШАБЛОНЫ ==================
def get_enhancement_prompt(user_input: str) -> str:
    """Создаем ПРАВИЛЬНЫЙ промт для улучшения, а не для рассказа"""
    return f"""<|system|>
Ты помощник для создания промтов к AI-генератору изображений Dreamlike Photoreal 2.0.
ТВОЯ ЗАДАЧА: Взять короткое описание и добавить КОНКРЕТНЫЕ детали для фотореалистичной генерации.

ДОБАВЬ:
1. Описание внешности (если есть персонаж)
2. Детали окружения
3. Освещение и время суток
4. Атмосферу и настроение
5. Стиль изображения

НЕ пиши рассказ, НЕ добавляй сюжет. Только описание для генератора изображений.

Примеры правильных улучшений:
• "девушка в лесу" → "Молодая девушка с рыжими волосами в белом платье стоит в магическом лесу на закате, золотые лучи солнца пробиваются сквозь листву, атмосферное освещение, фотореалистично, детализированно"
• "город ночью" → "Футуристический город ночью с неоновыми вывесками, мокрые улицы отражают свет, киберпанк стиль, детализированная архитектура, фотореалистично"</s>
<|user|>
Улучши это описание для генерации фотореалистичного изображения: "{user_input}"</s>
<|assistant|>
Улучшенное описание:"""

# ================== БЫСТРАЯ ГЕНЕРАЦИЯ ==================
def enhance_prompt_fast(user_input: str) -> str:
    """Быстро улучшаем промт с TinyLlama"""
    if _model is None or _tokenizer is None:
        return user_input + " [система не загружена]"
    
    try:
        # Создаем правильный промт
        system_prompt = get_enhancement_prompt(user_input)
        
        # Токенизация (ограничиваем длину)
        inputs = _tokenizer(
            system_prompt,
            return_tensors="pt",
            max_length=256,
            truncation=True,
            padding=True
        )
        
        # Быстрая генерация с оптимизированными параметрами
        with torch.no_grad():
            outputs = _model.generate(
                **inputs,
                max_new_tokens=120,  # Не больше 120 новых токенов
                temperature=0.8,
                do_sample=True,
                top_p=0.9,
                repetition_penalty=1.1,
                pad_token_id=_tokenizer.eos_token_id,
                eos_token_id=_tokenizer.eos_token_id,
                no_repeat_ngram_size=3  # Избегаем повторений
            )
        
        # Декодируем
        response = _tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Извлекаем только улучшенную часть
        if "Улучшенное описание:" in response:
            enhanced = response.split("Улучшенное описание:")[-1].strip()
        else:
            # Fallback: удаляем системный промт
            enhanced = response.replace(system_prompt, "").strip()
        
        # Очищаем от метаданных
        for marker in ["<|endoftext|>", "</s>", "<|assistant|>", "<|user|>", "<|system|>"]:
            enhanced = enhanced.split(marker)[0].strip()
        
        # Если результат слишком короткий, добавляем базовые улучшения
        if len(enhanced.split()) < 5:
            enhanced = f"{user_input}, фотореалистично, высокое качество, детализированно, красивое освещение"
        
        return enhanced[:400]  # Ограничиваем длину
        
    except Exception as e:
        print(f"Ошибка генерации: {e}")
        return f"{user_input}, фотореалистично, высокое качество, детализированно"

def generate_negative_prompt_smart(english_prompt: str) -> str:
    """Умный негативный промт"""
    base = [
        "blurry", "ugly", "deformed", "bad anatomy",
        "bad proportions", "extra limbs", "mutated hands",
        "poorly drawn hands", "text", "watermark", "signature",
        "worst quality", "low quality", "jpeg artifacts"
    ]
    
    prompt_lower = english_prompt.lower()
    
    # Добавляем контекстные исключения
    if any(word in prompt_lower for word in ["realistic", "photoreal", "photo"]):
        base.extend(["cartoon", "anime", "painting", "drawing", "illustration", "3d render"])
    
    if any(word in prompt_lower for word in ["bright", "light", "sun"]):
        base.extend(["dark", "dull", "underexposed", "gloomy"])
    
    if any(word in prompt_lower for word in ["portrait", "face", "person"]):
        base.extend(["asymmetric eyes", "cloned face", "malformed limbs"])
    
    # Убираем дубликаты
    unique = []
    for term in base:
        if term not in unique:
            unique.append(term)
    
    return ", ".join(unique)

def process_pipeline(user_input: str):
    """Полный пайплайн обработки"""
    start_time = time.time()
    
    if not user_input.strip():
        return "Введите промт", "", "", "⏳ Введите промт"
    
    print(f"⚡ Обработка: {user_input}")
    
    # 1. Улучшаем промт (максимум 5 секунд)
    gen_start = time.time()
    enhanced_ru = enhance_prompt_fast(user_input)
    gen_time = time.time() - gen_start
    
    # 2. Переводим
    trans_start = time.time()
    try:
        if _translator:
            translated = _translator(enhanced_ru[:300])[0]['translation_text']
        else:
            translated = enhanced_ru
    except:
        translated = enhanced_ru
    trans_time = time.time() - trans_start
    
    # 3. Улучшаем английский для Dreamlike
    dreamlike_keywords = ", photorealistic, high quality, detailed, 8K, professional photography"
    if "photorealistic" not in translated.lower():
        translated += dreamlike_keywords
    
    # 4. Негативный промт
    negative_start = time.time()
    negative_en = generate_negative_prompt_smart(translated)
    negative_time = time.time() - negative_start
    
    # 5. Статистика
    total_time = time.time() - start_time
    
    info = f"""
    ⚡ БЫСТРАЯ ОБРАБОТКА:
    • Генерация: {gen_time:.1f} сек
    • Перевод: {trans_time:.1f} сек
    • Негативный промт: {negative_time:.1f} сек
    • Всего: {total_time:.1f} сек
    
    📊 СТАТИСТИКА:
    • Модель: TinyLlama-1.1B (оптимизирована для CPU)
    • Длина промта: {len(translated.split())} слов
    • Статус: ✅ Готово к использованию в Dreamlike
    """
    
    return enhanced_ru, translated, negative_en, info.strip()

# ================== ПРОСТОЙ ИНТЕРФЕЙС ==================
with gr.Blocks(title="AI Улучшитель промтов для Dreamlike", theme=gr.themes.Soft()) as demo:
    
    # Заголовок
    gr.Markdown("""
    # 🚀 УЛЬТРА-БЫСТРЫЙ AI УЛУЧШИТЕЛЬ ПРОМТОВ
    ### Оптимизирован для CPU • TinyLlama • Мгновенная обработка
    """)
    
    # Статус загрузки
    status = gr.Textbox(
        label="🔄 Статус системы",
        value="Загружаю систему...",
        interactive=False
    )
    
    # Основной интерфейс
    with gr.Column(visible=False) as main_ui:
        # Ввод
        input_prompt = gr.Textbox(
            label="🎯 Введите промт на русском",
            placeholder="Пример: девушка в магическом лесу на закате",
            lines=2
        )
        
        # Кнопка
        process_btn = gr.Button("⚡ УЛУЧШИТЬ ПРОМТ", variant="primary", size="lg")
        
        # Результаты
        gr.Markdown("### 📝 УЛУЧШЕННЫЙ ПРОМТ (русский)")
        output_ru = gr.Textbox(
            label="",
            lines=3,
            interactive=True
        )
        
        gr.Markdown("### 🌐 ДЛЯ DREAMLIKE (английский)")
        with gr.Row():
            output_en = gr.Textbox(
                label="🎨 Основной промт",
                lines=3,
                interactive=True
            )
            
            output_neg = gr.Textbox(
                label="🚫 Negative Prompt",
                lines=3,
                interactive=True
            )
        
        # Информация
        output_info = gr.Textbox(
            label="📊 СТАТИСТИКА",
            lines=4,
            interactive=False
        )
        
        # Примеры
        gr.Markdown("### 🧪 БЫСТРЫЕ ПРИМЕРЫ")
        examples = gr.Examples(
            examples=[
                ["девушка в магическом лесу на закате"],
                ["киберпанк улица ночью с неоновыми вывесками"],
                ["космонавт играет на гитаре на луне"],
                ["кошка в шляпе читает газету в кафе"],
                ["дракон на вершине горы"]
            ],
            inputs=[input_prompt],
            outputs=[output_ru, output_en, output_neg, output_info],
            fn=process_pipeline,
            label="Кликните для теста:"
        )
        
        # Обработчики
        def process_and_update(prompt):
            return process_pipeline(prompt)
        
        process_btn.click(
            fn=process_and_update,
            inputs=[input_prompt],
            outputs=[output_ru, output_en, output_neg, output_info]
        )
        
        input_prompt.submit(
            fn=process_and_update,
            inputs=[input_prompt],
            outputs=[output_ru, output_en, output_neg, output_info]
        )
    
    # Функция инициализации
    def initialize_system():
        if load_system():
            return "✅ СИСТЕМА ГОТОВА! TinyLlama загружена и оптимизирована.\n⚡ Введите промт и нажмите кнопку.", gr.update(visible=True)
        else:
            return "❌ Ошибка загрузки. Попробуйте перезапустить Space.", gr.update(visible=False)
    
    # Загружаем при старте
    demo.load(
        fn=initialize_system,
        outputs=[status, main_ui]
    )

# ================== ЗАПУСК ==================
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860
    )