import fitz
import os


def process_flatten(source, dest, selected_files=None, progress_callback=None):
    os.makedirs(dest, exist_ok=True)
    total = 0
    processed = 0
    errors = []

    file_list = []
    for root, dirs, files in os.walk(source):
        for filename in files:
            if filename.lower().endswith('.pdf'):
                rel_path = os.path.relpath(os.path.join(root, filename), source)
                if selected_files is None or rel_path in selected_files:
                    file_list.append((root, filename))

    total = len(file_list)

    for root, filename in file_list:
        if progress_callback and progress_callback.is_cancelled():
            break

        doc_path = os.path.join(root, filename)
        rel_dir = os.path.relpath(root, source)
        display_path = filename if rel_dir == "." else os.path.join(rel_dir, filename)
        target_dir = os.path.join(dest, rel_dir)
        out_path = os.path.join(target_dir, filename)

        try:
            os.makedirs(target_dir, exist_ok=True)
            doc = fitz.open(doc_path)
            doc_aplati = fitz.open()

            for page in doc:
                rect_page = page.rect
                nouvelle_page = doc_aplati.new_page(width=rect_page.width, height=rect_page.height)
                zoom = 4
                mat = fitz.Matrix(zoom, zoom)
                pix = page.get_pixmap(matrix=mat, alpha=False)
                nouvelle_page.insert_image(rect_page, pixmap=pix)

            doc_aplati.save(out_path, garbage=4, deflate=True)
            doc.close()
            doc_aplati.close()

            processed += 1
            if progress_callback:
                progress_callback.on_file_done(display_path, True, None)
        except Exception as e:
            errors.append((display_path, str(e)))
            if progress_callback:
                progress_callback.on_file_done(display_path, False, str(e))
        finally:
            if progress_callback:
                progress_callback.on_progress(processed, total)

    return processed, total, errors


if __name__ == '__main__':
    import sys
    source = sys.argv[1] if len(sys.argv) > 1 else r"D:\ANNEXES CORRIGEES"
    dest = sys.argv[2] if len(sys.argv) > 2 else r"D:\ANNEXES POUR IMPRESSION"
    p, t, e = process_flatten(source, dest)
    print(f"Traités : {p}/{t}, Erreurs : {len(e)}")
    for f, msg in e:
        print(f"  ERREUR {f} : {msg}")