# app.py
import os
import zipfile
import rarfile
from PyPDF2 import PdfReader, PdfWriter
from pydub import AudioSegment
from PIL import Image
import pdf2image
import streamlit as st
import ffmpeg

def convert_file(input_path, output_path):
    if input_path.endswith('.pdf'):
        # Convert PDF to WEBP
        images = pdf2image.convert_from_path(input_path)
        for i, image in enumerate(images):
            image.save(f"{os.path.splitext(output_path)[0]}_page_{i+1}.webp", 'WEBP')

        # Convert WEBP back to PDF
        webp_images = [Image.open(f"{os.path.splitext(output_path)[0]}_page_{i+1}.webp") for i in range(len(images))]
        webp_images[0].save(output_path, 'PDF', resolution=100.0, save_all=True, append_images=webp_images[1:])

    elif input_path.endswith('.mp4'):
        # Convert MP4 to MKV using ffmpeg
        ffmpeg.input(input_path).output(output_path, codec='libx264').run()

    elif input_path.endswith('.wav'):
        # Convert WAV to MP3
        audio = AudioSegment.from_wav(input_path)
        audio.export(output_path, format='mp3')

    elif input_path.endswith(('.png', '.jpg')):
        # Convert PNG/JPG to WEBP
        image = Image.open(input_path)
        image.save(output_path, 'WEBP')

def extract_and_convert(file_path):
    if file_path.endswith('.zip'):
        with zipfile.ZipFile(file_path, 'r') as zip_ref:
            zip_ref.extractall(os.path.splitext(file_path)[0])
            extracted_dir = os.path.splitext(file_path)[0]
    elif file_path.endswith('.rar'):
        with rarfile.RarFile(file_path, 'r') as rar_ref:
            rar_ref.extractall(os.path.splitext(file_path)[0])
            extracted_dir = os.path.splitext(file_path)[0]
    else:
        return

    for root, _, files in os.walk(extracted_dir):
        for file in files:
            input_path = os.path.join(root, file)
            output_path = os.path.join(root, os.path.splitext(file)[0] + get_output_extension(file))
            convert_file(input_path, output_path)

def get_output_extension(file):
    if file.endswith('.pdf'):
        return '.pdf'
    elif file.endswith('.mp4'):
        return '.mkv'
    elif file.endswith('.wav'):
        return '.mp3'
    elif file.endswith(('.png', '.jpg')):
        return '.webp'
    return ''

# Streamlit UI
st.title("File Converter Space")

# Language selection
language = st.selectbox("Select Language / انتخاب زبان", ["English", "فارسی"])

if language == "English":
    description = """
    This space allows you to convert files between various formats.
    Supported conversions:
    - PDF to WEBP and back to PDF
    - MP4 to MKV
    - WAV to MP3
    - PNG/JPG to WEBP
    """
    upload_label = "Choose a file"
    convert_button = "Convert Files"
    success_message = "Files converted successfully!"
    converted_files_label = "Converted Files:"
else:
    description = """
    این اسپیس به شما امکان می‌دهد فایل‌ها را بین فرمت‌های مختلف تبدیل کنید.
    تبدیل‌های پشتیبانی شده:
    - PDF به WEBP و بازگشت به PDF
    - MP4 به MKV
    - WAV به MP3
    - PNG/JPG به WEBP
    """
    upload_label = "یک فایل انتخاب کنید"
    convert_button = "تبدیل فایل‌ها"
    success_message = "فایل‌ها با موفقیت تبدیل شدند!"
    converted_files_label = "فایل‌های تبدیل شده:"

st.write(description)

uploaded_file = st.file_uploader(upload_label, type=['zip', 'rar'])

if uploaded_file is not None:
    file_path = f"temp_{uploaded_file.name}"
    with open(file_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    st.success("File uploaded successfully!")

    if st.button(convert_button):
        extract_and_convert(file_path)
        st.success(success_message)

        # Display converted files
        st.write(converted_files_label)
        extracted_dir = os.path.splitext(file_path)[0]
        for root, _, files in os.walk(extracted_dir):
            for file in files:
                if file.endswith(('.pdf', '.mkv', '.mp3', '.webp')):
                    st.write(f"- {file}")