import os
import json
import numpy as np
import uuid
from PIL import Image
import requests
import telebot
import cv2
from tqdm import tqdm
from utils import ImageDataset
from run import run
from stereoimage_generation import create_stereoimages

# Загрузка конфигурации
with open('config.json', 'r') as f:
    config = json.load(f)


def parse_inf_values(d):
    for key, value in d.items():
        if value == "inf":
            d[key] = float('inf')
        elif isinstance(value, dict):
            parse_inf_values(value)


parse_inf_values(config)

telegram_token = config['telegram_token']
image_processing_config = config['image_processing_config']
depth_map_options = config['depth_map_options']

bot = telebot.TeleBot(telegram_token)


class VideoProcessor:
    def __init__(self, video_path, output_folder, frames_per_second):
        self.video_path = video_path
        self.output_folder = output_folder
        self.frames_per_second = frames_per_second

    def extract_frames(self):
        os.makedirs(self.output_folder, exist_ok=True)
        cap = cv2.VideoCapture(self.video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

        for frame_num in tqdm(range(total_frames), desc="Извлечение кадров"):
            ret, frame = cap.read()
            if not ret:
                break

            frame_filename = os.path.join(self.output_folder, f"frame_{frame_num}.jpg")
            cv2.imwrite(frame_filename, frame)

        cap.release()

    def create_video_from_frames(self, input_folder, output_filename):
        frame_files = [f"{input_folder}/{f}" for f in os.listdir(input_folder) if f.endswith('.jpg')]
        frame_files.sort(key=lambda x: int(''.join(filter(str.isdigit, x))))

        img = cv2.imread(frame_files[0])
        height, width, _ = img.shape

        out = cv2.VideoWriter(output_filename, cv2.VideoWriter_fourcc(*'mp4v'), self.frames_per_second, (width, height))

        for frame_file in tqdm(frame_files, desc="Собираем видео из кадров"):
            img = cv2.imread(frame_file)
            out.write(img)

        out.release()


class OptionsV:
    def __init__(self, data_dir, output_dir):
        self.data_dir = data_dir
        self.output_dir = output_dir
        self.patch_netsize = 2 * depth_map_options['net_receptive_field_size']
        # Устанавливаем остальные параметры из depth_map_options
        for key, value in depth_map_options.items():
            setattr(self, key, value)


def invert_depth_map(depth_map):
    return depth_map * -1


def create_stereo_images(input_image, depth_map, output_path, divergence, separation):
    inverted_depth_map = invert_depth_map(depth_map)
    stereo_images = create_stereoimages(
        input_image, inverted_depth_map, divergence=divergence, separation=separation,
        modes=['left-right'], fill_technique='polylines_sharp')

    output_path_jpg = os.path.splitext(output_path)[0] + '.jpg'
    stereo_images[0].save(output_path_jpg, format='JPEG', quality=100)
    return output_path_jpg


def square_image(img):
    width, height = img.size
    if width != height:
        size = max(width, height)
        new_img = Image.new('RGB', (size, size), color=(0, 0, 0))
        new_img.paste(img, ((size - width) // 2, (size - height) // 2))
        return new_img
    else:
        return img


def split_and_square_image(input_image_path):
    img = Image.open(input_image_path)
    width, height = img.size
    half_width = width // 2

    left_half = img.crop((0, 0, half_width, height))
    right_half = img.crop((half_width, 0, width, height))

    left_half_square = square_image(left_half)
    right_half_square = square_image(right_half)

    combined_width = left_half_square.width + right_half_square.width
    combined = Image.new('RGB', (combined_width, left_half_square.height))
    combined.paste(left_half_square, (0, 0))
    combined.paste(right_half_square, (left_half_square.width, 0))

    return combined


def process_frames_for_depth_maps(frame_files, temp_folder):
    dataset = ImageDataset(temp_folder, 'test')
    option = OptionsV(temp_folder, temp_folder)

    run(dataset, option)


def create_stereo_from_depth_maps(frame_files, temp_folder, processed_folder, divergence, separation):
    for frame_file in tqdm(frame_files, desc="Создание стереоизображений"):
        input_image_path = os.path.join(temp_folder, frame_file)
        input_image = Image.open(input_image_path)

        depth_map_path = os.path.join(temp_folder, os.path.splitext(frame_file)[0] + '.png')
        depth_map = Image.open(depth_map_path)

        # Обеспечение одинакового размера карты глубины и исходного изображения
        if input_image.size != depth_map.size:
            depth_map = depth_map.resize(input_image.size, Image.ANTIALIAS)

        depth_map_np = np.array(depth_map)
        stereo_image_path = create_stereo_images(input_image, depth_map_np, depth_map_path, divergence, separation)

        final_image_path = os.path.join(processed_folder, frame_file)
        final_image = split_and_square_image(stereo_image_path)
        final_image.save(final_image_path, format='JPEG', quality=100)


def process_image(input_file_path, output_path, divergence, separation, chat_id, message_id):
    class OptionsI:
        def __init__(self):
            for key, value in depth_map_options.items():
                setattr(self, key, value)
            self.data_dir = image_processing_config['input_dir']
            self.output_dir = output_path
            self.patch_netsize = 2 * depth_map_options['net_receptive_field_size']

    option = OptionsI()
    dataset = ImageDataset(image_processing_config['input_dir'], 'test')

    # Начало создания карты глубины
    bot.edit_message_text("⏳ Начало создания карты глубины...", chat_id=chat_id, message_id=message_id)
    run(dataset, option)

    depth_map_path = os.path.join(output_path, os.path.splitext(os.path.basename(input_file_path))[0] + '.png')
    input_image = Image.open(input_file_path)
    depth_map = np.array(Image.open(depth_map_path))

    # Генерация стереоизображения
    bot.edit_message_text("⏳ Генерация стереоизображения...", chat_id=chat_id, message_id=message_id)
    stereo_image_path = create_stereo_images(input_image, depth_map, depth_map_path, divergence, separation)

    # Преобразование в квадратное изображение
    bot.edit_message_text("⏳ Преобразование в квадратное изображение...", chat_id=chat_id, message_id=message_id)
    final_image = split_and_square_image(stereo_image_path)
    final_image_path = os.path.join(output_path, "final_" + os.path.basename(stereo_image_path))
    final_image.save(final_image_path, format='JPEG', quality=100)

    # Удаляем все временные файлы
    if os.path.exists(input_file_path):
        os.remove(input_file_path)

    if os.path.exists(depth_map_path):
        os.remove(depth_map_path)

    if os.path.exists(stereo_image_path):
        os.remove(stereo_image_path)

    return final_image_path


def process_video(input_video_path, output_dir, divergence, separation, chat_id, message_id):
    temp_folder = os.path.join(output_dir, "video_frames")
    processed_folder = os.path.join(output_dir, "processed_frames")
    os.makedirs(temp_folder, exist_ok=True)
    os.makedirs(processed_folder, exist_ok=True)

    video_processor = VideoProcessor(input_video_path, temp_folder, 24)
    video_processor.extract_frames()

    frame_files = [f for f in os.listdir(temp_folder) if f.endswith('.jpg')]
    frame_files.sort(key=lambda x: int(''.join(filter(str.isdigit, x))))

    # Генерация глубинных карт
    process_frames_for_depth_maps(frame_files, temp_folder)

    # Создание стереоизображений после генерации всех глубинных карт
    create_stereo_from_depth_maps(frame_files, temp_folder, processed_folder, divergence, separation)

    output_video_file = os.path.join(output_dir, f"processed_video.mp4")
    video_processor.create_video_from_frames(processed_folder, output_video_file)

    for f in os.listdir(temp_folder):
        os.remove(os.path.join(temp_folder, f))
    for f in os.listdir(processed_folder):
        os.remove(os.path.join(processed_folder, f))

    os.rmdir(temp_folder)
    os.rmdir(processed_folder)

    return output_video_file


@bot.message_handler(content_types=['video'])
def handle_video(message):
    sent_message = bot.reply_to(message, "📥 Обработка видео...")

    input_dir = image_processing_config['input_dir']
    output_dir = image_processing_config['output_dir']
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    file_info = bot.get_file(message.video.file_id)
    downloaded_file = bot.download_file(file_info.file_path)
    input_file_name = f"{uuid.uuid4()}.mp4"
    input_file_path = os.path.join(input_dir, input_file_name)

    with open(input_file_path, 'wb') as new_file:
        new_file.write(downloaded_file)

    divergence = image_processing_config['divergence']
    separation = image_processing_config['separation']

    try:
        output_video_path = process_video(
            input_file_path, output_dir, divergence, separation, message.chat.id, sent_message.id)
        if os.path.exists(output_video_path):
            with open(output_video_path, 'rb') as video:
                bot.send_video(message.chat.id, video)
            bot.edit_message_text("✅ Готово! Вот твое обработанное видео! 🎥", chat_id=message.chat.id,
                                  message_id=sent_message.id)
            os.remove(output_video_path)
        else:
            raise FileNotFoundError("Результат не найден.")
    except Exception as e:
        bot.edit_message_text(f"❌ Упс! Что-то пошло не так: {str(e)} 🤔 Попробуй ещё раз!", chat_id=message.chat.id,
                              message_id=sent_message.id)


@bot.message_handler(content_types=['photo'])
def handle_photo(message):
    sent_message = bot.reply_to(message, "📥 Обработка изображения...")

    input_dir = image_processing_config['input_dir']
    output_dir = image_processing_config['output_dir']
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    file_info = bot.get_file(message.photo[-1].file_id)
    downloaded_file = bot.download_file(file_info.file_path)
    input_file_name = f"{uuid.uuid4()}.jpg"
    input_file_path = os.path.join(input_dir, input_file_name)

    with open(input_file_path, 'wb') as new_file:
        new_file.write(downloaded_file)

    divergence = image_processing_config['divergence']
    separation = image_processing_config['separation']

    try:
        output_file_path = process_image(
            input_file_path, output_dir, divergence, separation, message.chat.id, sent_message.id)
        if os.path.exists(output_file_path):
            with open(output_file_path, 'rb') as img:
                bot.send_photo(message.chat.id, img)
            bot.edit_message_text("✅ Готово! Вот твое стереоизображение! 📸👓", chat_id=message.chat.id,
                                  message_id=sent_message.id)
            os.remove(output_file_path)  # Удаляем выходной файл после отправки
        else:
            raise FileNotFoundError("Результат не найден.")
    except Exception as e:
        bot.edit_message_text(f"❌ Упс! Что-то пошло не так: {str(e)} 🤔 Попробуй ещё раз!", chat_id=message.chat.id,
                              message_id=sent_message.id)


@bot.message_handler(content_types=['text'])
def handle_text(message):
    url = message.text.strip()
    if url.startswith('http://') or url.startswith('https://'):
        sent_message = bot.reply_to(message, "🌍 Пытаюсь загрузить изображение по ссылке...")

        try:
            response = requests.get(url)
            response.raise_for_status()

            input_file_name = f"{uuid.uuid4()}.jpg"
            input_file_path = os.path.join(image_processing_config['input_dir'], input_file_name)

            with open(input_file_path, 'wb') as f:
                f.write(response.content)

            bot.edit_message_text("🎨 Создание карты глубины с твоей ссылки...", chat_id=message.chat.id,
                                  message_id=sent_message.id)

            divergence = image_processing_config['divergence']
            separation = image_processing_config['separation']
            output_file_path = process_image(
                input_file_path, image_processing_config['output_dir'], divergence, separation, message.chat.id,
                sent_message.id)

            if os.path.exists(output_file_path):
                with open(output_file_path, 'rb') as img:
                    bot.send_photo(message.chat.id, img)
                bot.edit_message_text("✅ Готово! Вот твое стереоизображение! 📸👓", chat_id=message.chat.id,
                                      message_id=sent_message.id)
                os.remove(output_file_path)  # Удаляем выходной файл после отправки
            else:
                raise FileNotFoundError("Результат не найден.")
        except requests.exceptions.RequestException as e:
            bot.edit_message_text(f"❌ Не удалось загрузить изображение. Ошибка: {str(e)}. 🛑", chat_id=message.chat.id,
                                  message_id=sent_message.id)
        except Exception as e:
            bot.edit_message_text(f"❌ Упс! Что-то пошло не так: {str(e)} 🤔 Попробуй ещё раз!", chat_id=message.chat.id,
                                  message_id=sent_message.id)


bot.polling()


Надо сделать что бы при успешном преобразовании удалялось исходное сообщение с отправленным файлом и сообщение ✅ Готово! Вот твое обработанное видео!