db/mongodb.py

from pymongo import MongoClient
import gridfs
from bson import ObjectId
import os
from io import BytesIO
import magic

def connect_mongodb(db_name, collection_name):
    client = MongoClient("mongodb://localhost:27017")
    db = client[db_name]
    fs = gridfs.GridFS(db, collection=collection_name)
    return fs



def save_file_to_mongodb(uploaded_file, db_name="ppt", collection_name="root_file", file_name=None, file_tail=".pptx"):
    """
    Lưu file PowerPoint (pptx) vào MongoDB bằng GridFS
    nhưng không lưu nếu tên file đã tồn tại.
    
    :param uploaded_file: đối tượng UploadedFile từ Streamlit
    :param db_name: Tên database trong MongoDB
    :param collection_name: Tên collection GridFS
    :param file_name: Tên file muốn lưu (không cần .pptx). Nếu để None, lấy tên gốc.
    :return: file_id nếu lưu thành công, None nếu file đã tồn tại
    """
    client = MongoClient("mongodb://localhost:27017/")
    db = client[db_name]
    fs = gridfs.GridFS(db, collection=collection_name)

    # Xác định tên file
    if not file_name:
        # Lấy tên file từ uploaded_file (VD: "slide.pptx")
        file_name = uploaded_file.name
    else:
        # Nếu người dùng chỉ truyền tên, thêm .pptx nếu chưa có
        if not file_name.endswith(file_tail):
            file_name = file_name + file_tail

    # Kiểm tra file đã tồn tại trong MongoDB chưa
    existing_file = fs.find_one({"filename": file_name})
    if existing_file:
        print(f"⚠️ File '{file_name}' đã tồn tại trong MongoDB. Không lưu lại. Xin vui lòng đổi tên.")
        client.close()
        return None

    # Đảm bảo con trỏ file đang ở đầu
    uploaded_file.seek(0)
    file_bytes = uploaded_file.read()

    # Lưu nội dung file (bytes) vào MongoDB
    file_id = fs.put(file_bytes, filename=file_name)
    print(f"✅ File '{file_name}' đã được lưu vào '{collection_name}' với ID: {file_id}")
    client.close()
    return file_id

def delete_pptx_from_mongodb(file_id, db_name="ppt", collection_name="root_file"):
    """
    Xóa file PowerPoint khỏi MongoDB theo ID.

    :param file_id: ID của file cần xóa (chuỗi hoặc ObjectId)
    :param db_name: Tên database trong MongoDB
    :param collection_name: Tên collection GridFS
    """
    # Kết nối đến MongoDB
    client = MongoClient("mongodb://localhost:27017/")
    db = client[db_name]
    fs = gridfs.GridFS(db, collection=collection_name)

    try:
        # Chuyển đổi ID nếu cần
        if not isinstance(file_id, ObjectId):
            file_id = ObjectId(file_id)

        # Kiểm tra file có tồn tại không
        if fs.exists(file_id):
            fs.delete(file_id)
            print(f"✅ Đã xóa file với ID: {file_id}")
        else:
            print(f"⚠️ Không tìm thấy file với ID: {file_id}")
    except Exception as e:
        print(f"❌ Lỗi khi xóa file: {e}")

    client.close()

def download_pptx_from_mongodb(file_id, save_path, save_name, db_name="ppt", collection_name="root_file"):
    """
    Tải file PowerPoint từ MongoDB GridFS và lưu về máy.
    
    :param file_id:       ID của file cần tải (dạng chuỗi hoặc ObjectId)
    :param save_path:     Đường dẫn đến thư mục sẽ lưu file (VD: 'D:/output')
    :param save_name:     Tên file khi lưu (VD: 'my_presentation.pptx')
    :param db_name:       Tên database trong MongoDB (mặc định: 'ppt')
    :param collection_name: Tên collection GridFS (mặc định: 'root_file')
    """
    # Đảm bảo thư mục lưu file tồn tại
    os.makedirs(save_path, exist_ok=True)

    # Tạo đường dẫn đầy đủ cho file
    full_file_path = os.path.join(save_path, save_name)

    # Kết nối đến MongoDB
    client = MongoClient("mongodb://localhost:27017/")
    db = client[db_name]
    fs = gridfs.GridFS(db, collection=collection_name)

    try:
        # Chuyển đổi ID nếu cần
        if not isinstance(file_id, ObjectId):
            file_id = ObjectId(file_id)

        # Lấy dữ liệu file từ GridFS
        file_data = fs.get(file_id)
        
        # Ghi dữ liệu ra file
        with open(full_file_path, "wb") as f:
            f.write(file_data.read())

        print(f"✅ File đã được tải về: {full_file_path}")
    except Exception as e:
        print(f"❌ Lỗi khi tải file: {e}")
    finally:
        client.close()

def save_xml_to_gridfs(xml_content, file_name, db_name="ppt", collection_name="original_xml"):
    """
    Lưu XML vào MongoDB GridFS.
    
    :param xml_content: Chuỗi XML cần lưu
    :param file_name: Tên file XML
    :param db_name: Tên database MongoDB
    :param collection_name: Tên collection GridFS
    """
    client = MongoClient("mongodb://localhost:27017/")
    db = client[db_name]
    fs = gridfs.GridFS(db, collection=collection_name)

    # Kiểm tra file đã tồn tại chưa
    existing_file = fs.find_one({"filename": file_name})
    if existing_file:
        print(f"⚠️ File '{file_name}' đã tồn tại trong GridFS. Không lưu lại.")
        return

    # Chuyển đổi chuỗi XML thành bytes và lưu vào GridFS
    file_id = fs.put(xml_content.encode("utf-8"), filename=file_name)
    print(f"✅ XML '{file_name}' đã được lưu vào GridFS với ID: {file_id}")

def fetch_file_from_mongodb(db_name, collection_name, file_id):
    client = MongoClient("mongodb://localhost:27017/")  # Cập nhật nếu cần
    db = client[db_name]
    fs = gridfs.GridFS(db, collection_name)

    try:
        file_data = fs.get(file_id)
        pptx_io = BytesIO(file_data.read())
        pptx_io.seek(0)  # Đặt lại vị trí đầu file
        return pptx_io, file_data.filename
    except Exception as e:
        print(f"Lỗi khi lấy file từ MongoDB: {e}")
        return None, None
    
def detect_file_type(uploaded_file):
    if uploaded_file is not None:
        try:
            file_bytes = uploaded_file.read(4096)  # Đọc nhiều bytes hơn để nhận diện MIME
            mime = magic.Magic(mime=True)
            file_type = mime.from_buffer(file_bytes)
        except Exception as e:
            print(f"Error detecting file type: {e}")
            file_type = "Unknown"
        
        # Danh sách MIME types phổ biến
        mime_types = {
            "application/pdf": "PDF",
            "application/vnd.openxmlformats-officedocument.presentationml.presentation": "PPTX",
            "application/vnd.ms-powerpoint": "PPTX",
            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "Excel",
            "application/vnd.ms-excel": "Excel",
            "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "Word",
            "application/msword": "Word",
            "text/csv": "CSV",
            "text/plain": "CSV"  # Một số file CSV có thể nhận diện là text/plain
        }
        
        detected_type = mime_types.get(file_type, "Unknown")
        
        # Nếu vẫn không chắc, kiểm tra phần mở rộng file
        if detected_type == "Unknown":
            ext = os.path.splitext(uploaded_file.name)[1].lower()
            ext_mapping = {".csv": "CSV", ".docx": "Word", ".doc": "Word", ".xlsx": "Excel", ".pptx": "PPTX", ".pdf": "PDF"}
            detected_type = ext_mapping.get(ext, "Unknown")
        
        return detected_type
    return None