import os
import zipfile
from typing import Dict, List, Optional
import asyncio
import google.generativeai as genai
from google.api_core import exceptions
import gradio as gr
import pkg_resources

# --- Диагностика ---
try:
    version = pkg_resources.get_distribution("google-generativeai").version
    print(f"--- УСТАНОВЛЕНА ВЕРСИЯ google-generativeai: {version} ---")
except pkg_resources.DistributionNotFound:
    print("--- ПРЕДУПРЕЖДЕНИЕ: Библиотека google-generativeai не найдена ---")

# --- 1. Конфигурация API и Моделей ---
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
if not GOOGLE_API_KEY:
    raise gr.Error("Переменная окружения GOOGLE_API_KEY не установлена. Установите ее в секретах вашего Space.")

try:
    genai.configure(api_key=GOOGLE_API_KEY)
except Exception as e:
    raise gr.Error(f"Ошибка при настройке Google Gemini API: {e}.")

def get_available_models() -> List[str]:
    available_models = []
    try:
        for m in genai.list_models():
            if 'generateContent' in m.supported_generation_methods and not any(x in m.name.lower() for x in ['vision', 'tts', 'audio']):
                available_models.append(m.name)
    except Exception as e:
        print(f"Предупреждение: Не удалось получить список моделей: {e}.")
        available_models.extend(['models/gemini-1.5-flash-latest', 'models/gemini-pro'])
    return sorted(list(set(available_models)))

AVAILABLE_MODELS = get_available_models()
if not AVAILABLE_MODELS:
    raise gr.Error("Не найдено моделей, совместимых с 'generateContent'.")

# --- 2. Общие константы и глобальные переменные ---
TITLE = """<h1 align="center">✨ Gemini/Gemma Code Analysis</h1>"""
AVATAR_IMAGES = (None, "https://media.roboflow.com/spaces/gemini-icon.png")
TEXT_EXTENSIONS = [".bat", ".c", ".cfg", ".conf", ".cpp", ".cs", ".css", ".go", ".h", ".html", ".ini", ".java", ".js", ".json", ".jsx", ".md", ".php", ".ps1", ".py", ".rb", ".rs", ".sh", ".toml", ".ts", ".tsx", ".txt", ".xml", ".yaml", ".yml"]
DEFAULT_MODEL = "models/gemini-1.5-flash-latest" if "models/gemini-1.5-flash-latest" in AVAILABLE_MODELS else AVAILABLE_MODELS[0]
DEFAULT_NUM_AI_VARIANTS = 1

EXTRACTED_FILES: Dict[str, str] = {}

# --- 3. CSS для темной темы "Глубокий Космос" ---
custom_css = """
/* --- ИСПРАВЛЕНИЕ МАКЕТА И ПРОКРУТКИ --- */
/* 1. Класс для левой колонки: фиксированная высота, flex-контекст и главное - min-width: 0 */
.chat-left-column {
    height: 80vh; /* Ограничиваем высоту колонки */
    display: flex;
    flex-direction: column;
    min-width: 0; /* <-- КЛЮЧЕВОЕ ИСПРАВЛЕНИЕ: Предотвращает "съезд" дочерних элементов */
}

/* 2. Chatbot сам по себе должен заполнять доступное пространство в колонке */
#chatbot {
    flex-grow: 1; /* Заставляет чат-бот занимать всю доступную высоту */
    overflow: hidden; /* Обрезает все, что выходит за его пределы */
}

/* 3. Внутренний контейнер сообщений чат-бота получает прокрутку */
#chatbot .scroll-hide {
    height: 100%; /* Занимает 100% высоты родителя (#chatbot) */
    overflow-y: auto !important; /* Добавляет вертикальную прокрутку */
    word-wrap: break-word; /* Принудительный перенос слов */
}

/* 4. Принудительный перенос слов для самих сообщений (на всякий случай) */
.gradio-container .message {
    word-wrap: break-word;
    overflow-wrap: break-word;
}

/* --- СТИЛЬ ДЛЯ КОМПАКТНЫХ КНОПОК --- */
.compact-buttons button {
    padding: 4px 8px !important;      /* Уменьшаем внутренние отступы */
    min-height: 30px !important;     /* Уменьшаем минимальную высоту */
    margin-bottom: 8px !important;   /* Уменьшаем отступ снизу */
}
/* --- КОНЕЦ СТИЛЕЙ --- */


:root { --primary-color: #3B82F6; --app-bg-color: #111827; --input-bg-color: #1F2937; --border-color: #4B5563; --text-color-primary: #F3F4F6; --text-color-secondary: #9CA3AF; --label-color: #E5E7EB; }
.gradio-container { background-color: var(--app-bg-color) !important; color: var(--text-color-primary) !important; }
h1, .gr-markdown p { color: var(--text-color-primary) !important; }
.gradio-container label, .gradio-container .gr-info { font-weight: 600 !important; color: var(--label-color) !important; }
.variant-container { background: var(--input-bg-color); border: 1px solid var(--border-color); border-radius: 8px; margin-bottom: 10px; padding: 15px; }
.variant-header { display: flex; justify-content: space-between; align-items: center; margin-bottom: 10px; color: var(--label-color); }
.copy-button { background: transparent; border: 1px solid var(--border-color); border-radius: 5px; cursor: pointer; padding: 5px; }
.copy-button:hover { background: #374151; }
.copy-button svg { stroke: var(--secondary-color); }
.error-message { background-color: #450A0A; color: #F87171; border: 1px solid #7F1D1D; }
"""

# --- 4. Функции для работы с файлами ---
def extract_text_from_zip(zip_file_path: str) -> Dict[str, str]:
    text_contents = {}
    try:
        with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
            for file_info in zip_ref.infolist():
                if file_info.is_dir(): continue
                if any(file_info.filename.lower().endswith(ext) for ext in TEXT_EXTENSIONS):
                    with zip_ref.open(file_info) as file:
                        text_contents[file_info.filename] = file.read().decode("utf-8", errors="replace")
    except Exception as e:
        gr.Warning(f"Не удалось прочитать ZIP-архив: {e}")
    return text_contents

def extract_text_from_single_file(file_path: str) -> Dict[str, str]:
    text_contents = {}
    filename = os.path.basename(file_path)
    if any(filename.lower().endswith(ext) for ext in TEXT_EXTENSIONS):
        try:
            with open(file_path, "r", encoding="utf-8", errors="replace") as file:
                text_contents[filename] = file.read()
        except Exception as e:
            gr.Warning(f"Не удалось прочитать файл {filename}: {e}")
    return text_contents

def upload_files(files: List[gr.File], chatbot: List[Dict[str, str]]):
    global EXTRACTED_FILES
    EXTRACTED_FILES = {}
    all_extracted_files = {}
    for file_obj in files:
        if file_obj.name.lower().endswith(".zip"): extracted = extract_text_from_zip(file_obj.name)
        else: extracted = extract_text_from_single_file(file_obj.name)
        all_extracted_files.update(extracted)

    if not all_extracted_files:
        gr.Warning(f"Загружено {len(files)} файл(ов), но не найдено поддерживаемых текстовых файлов.")
        return chatbot, ""
    
    EXTRACTED_FILES = all_extracted_files
    file_list_md = "#### Загруженные файлы:\n" + "\n".join([f"- `{name}`" for name in EXTRACTED_FILES.keys()])
    chatbot.append({"role": "assistant", "content": f"Файлы ({len(EXTRACTED_FILES)} шт.) загружены. Теперь задавай свой вопрос."})
    return chatbot, file_list_md

# --- 5. Логика генерации ответов ---
def format_history_for_gemini(history: List[Dict[str, str]]):
    gemini_history = []
    for msg in history[:-1]: 
        role = "user" if msg['role'] == "user" else "model"
        gemini_history.append({'role': role, 'parts': [{'text': msg['content']}]})
    return gemini_history

def format_variants_html(variants: List[str]) -> str:
    if not variants: return ""
    html_outputs = []
    for i, variant_text in enumerate(variants):
        js_safe_text = variant_text.replace('`', '\\`').replace('\n', '\\n').replace("'", "\\'")
        copy_button_html = f"""<button onclick="navigator.clipboard.writeText(`{js_safe_text}`)" class="copy-button" title="Копировать текст"><svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect><path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path></svg></button>"""
        if "<strong>Ошибка" in variant_text:
            html_outputs.append(f'<div class="variant-container error-message">{variant_text}</div>')
        else:
            header = f"<strong>Вариант {i + 1}</strong>" if len(variants) > 1 else ""
            html_outputs.append(f'<div class="variant-container"><div class="variant-header">{header}{copy_button_html}</div><div class="variant-text">{variant_text}</div></div>')
    return "".join(html_outputs)

async def generate_single_variant_async(full_context: List[Dict], model_name: str, temperature: float):
    try:
        model = genai.GenerativeModel(model_name=model_name)
        response = await model.generate_content_async(contents=full_context, generation_config={"temperature": temperature})
        return response.text.strip() if response.text else "*(AI не вернул текст)*"
    except exceptions.PermissionDenied as e:
        return f"<strong>Ошибка: Отказано в доступе.</strong><br>Пожалуйста, проверьте ваш `GOOGLE_API_KEY`."
    except Exception as e:
        return f"<strong>Ошибка генерации:</strong><br>{e}"

async def generate_response_flow(chatbot: List[Dict[str, str]], model_name: str, temperature: float, num_variants: int):
    if not chatbot or chatbot[-1]['role'] == "assistant":
        yield chatbot
        return

    chatbot.append({"role": "assistant", "content": "..."})
    yield chatbot

    gemini_history = format_history_for_gemini(chatbot)
    user_prompt = chatbot[-2]['content']
    
    user_parts = []
    if EXTRACTED_FILES:
        file_context = "\n\n".join([f"### File: {name}\n```\n{content}\n```" for name, content in EXTRACTED_FILES.items()])
        user_parts.append({'text': f"КОНТЕКСТ ФАЙЛОВ:\n{file_context}"})
    
    user_parts.append({'text': user_prompt})
    
    full_context = gemini_history + [{'role': 'user', 'parts': user_parts}]

    tasks = [generate_single_variant_async(full_context, model_name, temperature) for _ in range(int(num_variants))]
    results = await asyncio.gather(*tasks)

    final_html = format_variants_html(results)
    chatbot[-1]['content'] = final_html
    yield chatbot

# --- 6. Вспомогательные функции для UI ---
def user(text_prompt: str, chatbot: List[Dict[str, str]]):
    return "", chatbot + [{"role": "user", "content": text_prompt}] if text_prompt else chatbot

def prepare_for_regeneration(chatbot: List[Dict[str, str]]):
    if len(chatbot) >= 2 and chatbot[-1]['role'] == "assistant" and chatbot[-2]['role'] == "user":
        return chatbot[:-1]
    return chatbot

def reset_app():
    global EXTRACTED_FILES
    EXTRACTED_FILES = {}
    initial_message = [{"role": "assistant", "content": "Выполнен сброс. Загрузить — это для файлов, остальное — понятно."}]
    return initial_message, ""

# --- 7. Создание интерфейса Gradio ---
with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.HTML(TITLE)

    # Создаем одну главную строку, как в первом примере
    with gr.Row(equal_height=False):

        # --- Левая колонка: Чат, отображение файлов и ввод ---
        # Оставляем ее почти без изменений, только убираем класс chat-left-column,
        # если он больше не нужен для специфичных CSS-правил
        with gr.Column(scale=3):
            chatbot_component = gr.Chatbot(
                value=[{"role": "assistant", "content": "Ку Ами~! Загрузить — это для файлов, остальное — понятно."}],
                label="Gemini Code Analysis",
                avatar_images=AVATAR_IMAGES,
                type='messages',
                show_copy_button=True,
                elem_id="chatbot",
                # Рекомендуется задать высоту, чтобы колонка не "прыгала"
                height=500
            )
            file_display_component = gr.Markdown(label="Загруженные Файлы")
            text_prompt_component = gr.Textbox(
                placeholder="Введи свой запрос...", 
                show_label=False, 
                container=False, 
                lines=2, # Начальное количество строк
                max_lines=10 # Максимальное количество строк, после которого появится прокрутка
            )

        # --- Правая колонка: Все настройки и кнопки ---
        # Создаем вторую колонку в той же строке
        with gr.Column(scale=1, min_width=250): # min_width для удобства на маленьких экранах
            # Переносим сюда все элементы управления
            model_selector = gr.Dropdown(choices=AVAILABLE_MODELS, value=DEFAULT_MODEL, label="Выбирай Модель", interactive=True)
            temperature_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.9, label="Температура")
            num_variants_slider = gr.Slider(minimum=1, maximum=5, step=1, value=DEFAULT_NUM_AI_VARIANTS, label="Количество Вариантов Ответа")

            gr.Markdown("---") # Визуальный разделитель

            # Группируем кнопки
            send_button_component = gr.Button("Спросить (Shift+Enter)", variant="primary")
            regenerate_button_component = gr.Button("🔄 Переспросить", variant="secondary")
            upload_button_component = gr.UploadButton("📤 Загрузить", file_count="multiple", file_types=[".zip"] + TEXT_EXTENSIONS, variant="secondary")
            reset_button_component = gr.Button("🗑️ Сбросить", variant="stop")

    # --- 8. Логика обработчиков событий ---
    generation_inputs = [chatbot_component, model_selector, temperature_slider, num_variants_slider]

    text_prompt_component.submit(
        fn=user, inputs=[text_prompt_component, chatbot_component],
        outputs=[text_prompt_component, chatbot_component], queue=False
    ).then(
        fn=generate_response_flow, inputs=generation_inputs, outputs=[chatbot_component]
    )
    send_button_component.click(
        fn=user, inputs=[text_prompt_component, chatbot_component],
        outputs=[text_prompt_component, chatbot_component], queue=False
    ).then(
        fn=generate_response_flow, inputs=generation_inputs, outputs=[chatbot_component]
    )

    upload_button_component.upload(fn=upload_files, inputs=[upload_button_component, chatbot_component], outputs=[chatbot_component, file_display_component])

    regenerate_button_component.click(fn=prepare_for_regeneration, inputs=[chatbot_component], outputs=[chatbot_component], queue=False).then(
        fn=generate_response_flow, inputs=generation_inputs, outputs=[chatbot_component]
    )

    reset_button_component.click(fn=reset_app, outputs=[chatbot_component, file_display_component], queue=False)

demo.queue().launch()