import gradio as gr
from huggingface_hub import InferenceClient
import PyPDF2
from PIL import Image
import cv2
import numpy as np
from pydub import AudioSegment
from langdetect import detect

# Инициализация клиента для модели HuggingFaceH4/zephyr-7b-beta
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Функция для обработки PDF
def process_pdf(file):
    pdf_reader = PyPDF2.PdfReader(file)
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

# Функция для обработки изображений
def process_image(file):
    image = Image.open(file)
    return f"Изображение: {image.size[0]}x{image.size[1]} пикселей, формат: {image.format}"

# Функция для обработки видео
def process_video(file):
    cap = cv2.VideoCapture(file.name)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = frame_count / cap.get(cv2.CAP_PROP_FPS)
    cap.release()
    return f"Видео: длительность {duration:.2f} секунд, {frame_count} кадров"

# Функция для обработки аудио
def process_audio(file):
    audio = AudioSegment.from_file(file)
    return f"Аудио: длительность {len(audio) / 1000:.2f} секунд, частота {audio.frame_rate} Гц"

# Функция для обработки текстового файла
def process_txt(file):
    with open(file.name, "r", encoding="utf-8") as f:
        text = f.read()
    return text

# Функция для определения языка текста
def detect_language(text):
    try:
        return detect(text)
    except:
        return "en"  # По умолчанию английский

# Функция для обработки сообщений, истории и файлов
def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    file=None,
):
    # Если загружен файл, обрабатываем его
    if file is not None:
        file_type = file.name.split(".")[-1].lower()
        if file_type == "pdf":
            file_info = process_pdf(file)
        elif file_type in ["jpg", "jpeg", "png", "bmp", "gif"]:
            file_info = process_image(file)
        elif file_type in ["mp4", "avi", "mov"]:
            file_info = process_video(file)
        elif file_type in ["mp3", "wav", "ogg"]:
            file_info = process_audio(file)
        elif file_type == "txt":
            file_info = process_txt(file)
        else:
            file_info = "Неизвестный тип файла"
        message += f"\n[Пользователь загрузил файл: {file.name}]\n{file_info}"

    # Определяем язык сообщения
    language = detect_language(message)

    # Добавляем системное сообщение с учетом языка
    if language == "ru":
        system_message = "Вы дружелюбный чат-бот, который понимает русский язык."
    else:
        system_message = "You are a friendly chatbot."

    # Добавляем системное сообщение
    messages = [{"role": "system", "content": system_message}]

    # Добавляем историю сообщений
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # Добавляем текущее сообщение пользователя
    messages.append({"role": "user", "content": message})

    # Генерация ответа с использованием модели HuggingFaceH4/zephyr-7b-beta
    response = ""
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

# Функция для сброса истории чата
def reset_chat():
    return []

# Функция для анализа текстового файла
def analyze_txt(file):
    if file is None:
        return "Файл не загружен."
    text = process_txt(file)
    return f"Содержимое файла:\n{text}"

# Создание интерфейса
with gr.Blocks() as demo:
    gr.Markdown("# Felguk v0")
    
    # Переменная для управления видимостью
    current_tab = gr.State(value="chat")  # По умолчанию открыт чат

    # Функция для переключения вкладок
    def switch_tab(tab):
        return tab

    # Кнопки для переключения между вкладками
    with gr.Row():
        chat_button = gr.Button("Чат", variant="primary")
        tools_button = gr.Button("Felguk Tools", variant="secondary")

    # Интерфейс чата
    with gr.Column(visible=True) as chat_interface:
        gr.Markdown("## Чат")
        chat_interface = gr.ChatInterface(
            respond,
            additional_inputs=[
                gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
                gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
                gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
                gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top-p (nucleus sampling)",
                ),
                gr.File(label="Загрузите файл (опционально)"),  # Поле для загрузки файлов
            ],
        )
        new_chat_button = gr.Button("Новый чат", variant="secondary")
        new_chat_button.click(fn=reset_chat, outputs=chat_interface.chatbot)

    # Felguk Tools: Txt Analyzer
    with gr.Column(visible=False) as tools_interface:
        gr.Markdown("## Felguk Tools")
        txt_file = gr.File(label="Загрузите txt файл", file_types=[".txt"])
        txt_output = gr.Textbox(label="Содержимое файла", interactive=False)
        analyze_button = gr.Button("Анализировать")
        analyze_button.click(fn=analyze_txt, inputs=txt_file, outputs=txt_output)
        back_to_chat_button = gr.Button("Вернуться в чат", variant="primary")

    # Логика переключения вкладок
    chat_button.click(
        fn=lambda: (gr.Column.update(visible=True), gr.Column.update(visible=False)),
        outputs=[chat_interface, tools_interface],
    )
    tools_button.click(
        fn=lambda: (gr.Column.update(visible=False), gr.Column.update(visible=True)),
        outputs=[chat_interface, tools_interface],
    )
    back_to_chat_button.click(
        fn=lambda: (gr.Column.update(visible=True), gr.Column.update(visible=False)),
        outputs=[chat_interface, tools_interface],
    )

# Запуск интерфейса
if __name__ == "__main__":
    demo.launch()