import gradio as gr
from huggingface_hub import InferenceClient
import PyPDF2
from PIL import Image, ImageFont, ImageDraw
import cv2
import numpy as np
from pydub import AudioSegment
from langdetect import detect
from rembg import remove
import torch
import torchvision
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.transforms import functional as F
import tempfile
import time
import requests
import zipfile
import os
import torchaudio
from transformers import pipeline
from googletrans import Translator

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

# Загрузка предобученной модели для обнаружения объектов
model = fasterrcnn_resnet50_fpn(weights="DEFAULT")
model.eval()

# Функция для обработки 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 remove_background(image):
    if image is None:
        return None
    output = remove(image)
    return output

# Функция для определения количества объектов на изображении
def count_objects(image):
    if image is None:
        return "Изображение не загружено."
    
    # Преобразование изображения в формат, подходящий для модели
    img = Image.open(image.name).convert("RGB")
    img_tensor = F.to_tensor(img).unsqueeze(0)
    
    # Обнаружение объектов
    with torch.no_grad():
        predictions = model(img_tensor)
    
    # Подсчет количества объектов
    num_objects = len(predictions[0]['labels'])
    return f"Количество объектов на изображении: {num_objects}"

# Функция для конвертации изображения
def convert_image(image, target_format):
    if image is None:
        return None
    img = Image.open(image.name)
    
    # Создаем временный файл
    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{target_format.lower()}") as tmp_file:
        img.save(tmp_file, format=target_format)
        return tmp_file.name

# Функция для определения языка текста
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}"

# Функция для изменения размера изображения
def resize_image(image, width: int, height: int):
    if image is None:
        return None
    
    img = Image.open(image.name)
    resized_img = img.resize((width, height))
    
    # Сохраняем изображение во временный файл
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
        resized_img.save(tmp_file.name)
        return tmp_file.name

# Функция для конвертации аудио
def convert_audio(file, target_format: str):
    if file is None:
        return "Файл не загружен."
    
    audio = AudioSegment.from_file(file.name)
    
    # Сохраняем аудио в новом формате
    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{target_format.lower()}") as tmp_file:
        audio.export(tmp_file.name, format=target_format)
        return tmp_file.name

# Функция для перевода текста
def translate_text(text: str, target_language: str):
    translator = Translator()
    translation = translator.translate(text, dest=target_language)
    return translation.text

# Функция для обрезки видео
def trim_video(file, start_time: float, end_time: float):
    if file is None:
        return "Файл не загружен."
    
    cap = cv2.VideoCapture(file.name)
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # Создаем объект для записи видео
    output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    # Обрезка видео
    cap.set(cv2.CAP_PROP_POS_MSEC, start_time * 1000)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret or cap.get(cv2.CAP_PROP_POS_MSEC) >= end_time * 1000:
            break
        out.write(frame)
    
    cap.release()
    out.release()
    return output_path

# Функция для объединения PDF-файлов
def merge_pdfs(files):
    if not files:
        return "Файлы не загружены."
    
    merger = PyPDF2.PdfMerger()
    for file in files:
        merger.append(file.name)
    
    # Сохраняем объединенный PDF во временный файл
    output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf").name
    merger.write(output_path)
    merger.close()
    return output_path

# Создание интерфейса
with gr.Blocks() as demo:
    # Заголовок с изображением
    gr.HTML("""
    <div style="text-align: center;">
        <img src="https://huggingface.co/spaces/Felguk/Felguk-v0/resolve/main/hd_crop_4c480c6a2c7e176289b0dfcb64a30603_67753ddec8355.png" alt="Felguk Logo" style="width: 300px;">
    </div>
    """)
    
    gr.Markdown("Чат-бот Felguk v0. Отвечает на том же языке, на котором вы написали. Задавайте вопросы и загружайте файлы (PDF, изображения, видео, аудио, txt)!")
    
    # Кнопка "Новый чат"
    with gr.Row():
        new_chat_button = gr.Button("Новый чат", variant="secondary")

    # Felguk Tools: Отдельные инструменты
    with gr.Tab("Felguk Tools"):
        # Вкладка Txt Analyzer
        with gr.Tab("Анализатор текста"):
            gr.Markdown("## Анализатор текста")
            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)

        # Вкладка Remove Background
        with gr.Tab("Удаление фона"):
            gr.Markdown("## Удаление фона с изображения")
            image_input = gr.Image(label="Загрузите изображение", type="pil")
            image_output = gr.Image(label="Результат (без фона)", type="pil")
            remove_bg_button = gr.Button("Удалить фон")
            remove_bg_button.click(fn=remove_background, inputs=image_input, outputs=image_output)

        # Вкладка Numage
        with gr.Tab("Numage"):
            gr.Markdown("## Numage: Подсчет объектов на изображении")
            numage_input = gr.File(label="Загрузите изображение", file_types=["image"])
            numage_output = gr.Textbox(label="Результат", interactive=False)
            numage_button = gr.Button("Определить количество объектов")
            numage_button.click(fn=count_objects, inputs=numage_input, outputs=numage_output)

        # Вкладка ConverjerIMG
        with gr.Tab("Конвертер изображений"):
            gr.Markdown("## Конвертер изображений")
            img_input = gr.File(label="Загрузите изображение", file_types=["image"])
            img_format = gr.Dropdown(
                choices=["JPEG", "PNG", "BMP", "GIF", "TIFF"],
                label="Выберите формат для конвертации",
                value="JPEG"
            )
            img_output = gr.File(label="Результат конвертации")
            convert_button = gr.Button("Конвертировать")
            convert_button.click(fn=convert_image, inputs=[img_input, img_format], outputs=img_output)

        # Вкладка Felguk-ImageResizer
        with gr.Tab("Felguk-ImageResizer"):
            gr.Markdown("## Felguk-ImageResizer: Изменение размера изображения")
            image_input = gr.File(label="Загрузите изображение", file_types=["image"])
            width_input = gr.Number(label="Ширина", value=300)
            height_input = gr.Number(label="Высота", value=300)
            resize_button = gr.Button("Изменить размер")
            resized_image_output = gr.Image(label="Результат", type="filepath")
            resize_button.click(fn=resize_image, inputs=[image_input, width_input, height_input], outputs=resized_image_output)

        # Вкладка Felguk-AudioConverter
        with gr.Tab("Felguk-AudioConverter"):
            gr.Markdown("## Felguk-AudioConverter: Конвертация аудиофайлов")
            audio_input = gr.File(label="Загрузите аудиофайл", file_types=[".mp3", ".wav"])
            target_format_input = gr.Dropdown(
                choices=["mp3", "wav", "ogg"],
                label="Выберите формат для конвертации",
                value="mp3"
            )
            convert_button = gr.Button("Конвертировать")
            converted_audio_output = gr.Audio(label="Результат")
            convert_button.click(fn=convert_audio, inputs=[audio_input, target_format_input], outputs=converted_audio_output)

        # Вкладка Felguk-TextTranslator
        with gr.Tab("Felguk-TextTranslator"):
            gr.Markdown("## Felguk-TextTranslator: Перевод текста")
            text_input = gr.Textbox(label="Введите текст для перевода", lines=5)
            target_language_input = gr.Dropdown(
                choices=["en", "ru", "es", "fr", "de"],
                label="Выберите язык перевода",
                value="en"
            )
            translate_button = gr.Button("Перевести")
            translated_text_output = gr.Textbox(label="Результат", interactive=False)
            translate_button.click(fn=translate_text, inputs=[text_input, target_language_input], outputs=translated_text_output)

        # Вкладка Felguk-VideoTrimmer
        with gr.Tab("Felguk-VideoTrimmer"):
            gr.Markdown("## Felguk-VideoTrimmer: Обрезка видеофайлов")
            video_input = gr.File(label="Загрузите видеофайл", file_types=[".mp4", ".avi"])
            start_time_input = gr.Number(label="Начало обрезки (сек)", value=0)
            end_time_input = gr.Number(label="Конец обрезки (сек)", value=10)
            trim_button = gr.Button("Обрезать видео")
            trimmed_video_output = gr.Video(label="Результат")
            trim_button.click(fn=trim_video, inputs=[video_input, start_time_input, end_time_input], outputs=trimmed_video_output)

        # Вкладка Felguk-PDFMerger
        with gr.Tab("Felguk-PDFMerger"):
            gr.Markdown("## Felguk-PDFMerger: Объединение PDF-файлов")
            pdf_files_input = gr.Files(label="Загрузите PDF-файлы", file_types=[".pdf"])
            merge_button = gr.Button("Объединить PDF")
            merged_pdf_output = gr.File(label="Результат")
            merge_button.click(fn=merge_pdfs, inputs=pdf_files_input, outputs=merged_pdf_output)

    # Интерфейс чата
    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.click(fn=reset_chat, outputs=chat_interface.chatbot)

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