import os
import zipfile
import shutil
from utils.utils import unzip_office_file, translate_text, preprocess_text, postprocess_text
from powerpoint.xml_handling import *
from pymongo import MongoClient
import gridfs
from io import BytesIO

def create_pptx_and_store_in_mongodb(temp_dir, pptx_filename):
    """
    Tạo file PPTX từ thư mục chứa nội dung đã giải nén và lưu vào MongoDB mà không lưu file trên ổ cứng.
    """
    pptx_buffer = BytesIO()  

    with zipfile.ZipFile(pptx_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root_dir, _, files in os.walk(temp_dir):
            for file in files:
                file_path = os.path.join(root_dir, file)
                arcname = os.path.relpath(file_path, temp_dir)
                zipf.write(file_path, arcname)

    pptx_buffer.seek(0)  

    client = MongoClient("mongodb+srv://admin:1highbar456@cluster0.equkm.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0")
    db = client['pptx']
    fs = gridfs.GridFS(db, collection='final_file')

    file_id = fs.put(pptx_buffer, filename=pptx_filename)

    print(f"PPTX đã được lưu vào MongoDB với ID: {file_id}")
    client.close()

    return file_id

def translate_and_replace_pptx(xml_folder, file_name, source_lang='vn', target_lang='en', slides_per_batch=5):
    slides_dir = os.path.join(xml_folder, "ppt/slides")
    all_slides = sorted([f for f in os.listdir(slides_dir)
                         if f.startswith("slide") and f.endswith(".xml")],
                        key=lambda x: int(x[5:-4]))

    # Xử lý theo từng batch slide
    for i in range(0, len(all_slides), slides_per_batch):
        batch_slides = all_slides[i:i + slides_per_batch]
        slide_text_mapping = {}  
        smartart_text_mapping = {} 

        for slide_file in batch_slides:
            slide_index = int(slide_file[5:-4])
            slide_path = os.path.join(slides_dir, slide_file)
            slide_text_mapping[slide_index] = extract_text_from_slide(slide_path)  # Lấy list các tuple (text, rPr)

            # Xử lý SmartArt qua file .rels của slide
            rels_file = os.path.join(xml_folder, "ppt/slides/_rels", slide_file + ".rels")
            base_path = os.path.join(xml_folder, "ppt")
            smartart_data_path = get_smartart_data_file(rels_file, base_path)
            if smartart_data_path:
                smartart_text_mapping[slide_index] = extract_text_from_smartart(smartart_data_path) # Lấy list các tuple (text, rPr)


        # Gộp text để dịch theo batch, giữ lại rPr
        combined_slide_text_list = []
        for slide_index in sorted(slide_text_mapping.keys()):
            combined_slide_text_list.extend(slide_text_mapping[slide_index])

        combined_smartart_text_list = []
        for slide_index in sorted(smartart_text_mapping.keys()):
            combined_smartart_text_list.extend(smartart_text_mapping[slide_index])
        
        # Tách text ra khỏi tuple để dịch
        slide_texts_to_translate = [text for text, _ in combined_slide_text_list]
        smartart_texts_to_translate = [text for text, _ in combined_smartart_text_list]

        # Dịch văn bản slide và SmartArt
        combined_slide_text_string = preprocess_text(slide_texts_to_translate)
        combined_smartart_text_string = preprocess_text(smartart_texts_to_translate)

        translated_slide_string = translate_text(combined_slide_text_string, source_lang, target_lang)
        translated_smartart_string = translate_text(combined_smartart_text_string, source_lang, target_lang)

        # Postprocess để có list các văn bản đã dịch
        translated_slide_texts = postprocess_text(translated_slide_string)
        translated_smartart_texts = postprocess_text(translated_smartart_string)

        # **Quan trọng:** Tạo danh sách tuple (translated_text, rPr)
        translated_slide_data = []
        for i, (original_text, rPr) in enumerate(combined_slide_text_list):
            if i < len(translated_slide_texts):
                translated_slide_data.append((translated_slide_texts[i], rPr))
            else:
                translated_slide_data.append(("", rPr)) # Trường hợp không đủ translated text
        
        translated_smartart_data = []
        for i, (original_text, rPr) in enumerate(combined_smartart_text_list):
            if i < len(translated_smartart_texts):
                translated_smartart_data.append((translated_smartart_texts[i], rPr))
            else:
                translated_smartart_data.append(("", rPr))  # Trường hợp không đủ translated text

        # Thay thế văn bản trong slide
        slide_index = 0
        for slide_index in sorted(slide_text_mapping.keys()):
            slide_file = f"slide{slide_index}.xml"
            slide_path = os.path.join(slides_dir, slide_file)
            num_texts = len(slide_text_mapping[slide_index])
            replace_data = translated_slide_data[:num_texts]
            replace_text_in_slide(slide_path, replace_data) # truyền vào danh sách (translated_text, rPr)
            translated_slide_data = translated_slide_data[num_texts:]  # Cập nhật danh sách cho slide tiếp theo

        # Thay thế văn bản trong SmartArt
        for slide_index in sorted(smartart_text_mapping.keys()):
            rels_file = os.path.join(xml_folder, "ppt/slides/_rels", f"slide{slide_index}.xml.rels")
            base_path = os.path.join(xml_folder, "ppt")
            smartart_data_path = get_smartart_data_file(rels_file, base_path)
            if smartart_data_path:
                num_texts = len(smartart_text_mapping[slide_index])
                replace_data = translated_smartart_data[:num_texts]
                replace_text_in_smartart(smartart_data_path, replace_data, None) # truyền vào danh sách (translated_text, rPr)
                translated_smartart_data = translated_smartart_data[num_texts:]  # Cập nhật danh sách cho slide tiếp theo

    file_id = create_pptx_and_store_in_mongodb(xml_folder, file_name)
    return file_id

def translate_pptx(pptx_id, file_name, source_lang='vn', target_lang='en', slides_per_batch=5):
    client = MongoClient("mongodb+srv://admin:1highbar456@cluster0.equkm.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0")
    db = client['pptx']
    fs = gridfs.GridFS(db, collection='root_file')
    
    ppt_file = fs.get(pptx_id)
    prs = BytesIO(ppt_file.read())
    
    xml_folder = unzip_office_file(prs)
    file_id =  translate_and_replace_pptx(xml_folder, file_name, source_lang, target_lang, slides_per_batch=slides_per_batch)
    shutil.rmtree(xml_folder)

    return file_id