src/firestore_db.py

from google.cloud import firestore
from google.oauth2 import service_account
from config.settings import Settings
from src.logger import logger
from datetime import datetime

class FirestoreDB:
    def __init__(self):
        # Kiểm tra FIRESTORE_CREDENTIALS từ Settings
        if not Settings.FIRESTORE_CREDENTIALS:
            logger.error("Thiếu FIRESTORE_CREDENTIALS trong cấu hình.")
            raise ValueError("Thiếu FIRESTORE_CREDENTIALS. Vui lòng đặt biến môi trường FIRESTORE_CREDENTIALS với nội dung JSON credentials.")

        try:
            # Khởi tạo credentials từ FIRESTORE_CREDENTIALS
            credentials_dict = Settings.FIRESTORE_CREDENTIALS
            if not isinstance(credentials_dict, dict) or "type" not in credentials_dict:
                raise ValueError("FIRESTORE_CREDENTIALS không phải là JSON hợp lệ.")

            # Sử dụng credentials cho google-cloud-firestore
            credentials_obj = service_account.Credentials.from_service_account_info(credentials_dict)
            self.db = firestore.Client(credentials=credentials_obj, project=credentials_dict.get("project_id"))
            logger.info("Khởi tạo FirestoreDB thành công với FIRESTORE_CREDENTIALS.")
        except Exception as e:
            logger.error(f"Lỗi khi khởi tạo FirestoreDB: {str(e)}")
            raise

        self.documents = self.db.collection("documents")
        self.lottery = self.db.collection("lottery")

    def save_document(self, text, vector):
        if not self.document_exists(text):
            self.documents.add({"text": text, "vector": vector.tolist()})
            logger.info(f"Đã lưu document: {text[:50]}...")

    def document_exists(self, text):
        query = self.documents.where(filter=firestore.FieldFilter("text", "==", text))
        return bool(query.get())

    def get_all_documents(self):
        return [(doc.to_dict()["text"], doc.to_dict()["vector"]) for doc in self.documents.stream()]

    def save_lottery(self, data):
        ngay = data.get("ngay")
        dai = data.get("dai")
        if not all([ngay, dai]):
            logger.error("Dữ liệu thiếu ngay hoặc dai")
            return False
        try:
            # Chuyển đổi ngày từ YYYY-MM-DD sang DD-MM-YYYY
            ngay = datetime.strptime(ngay, "%Y-%m-%d").strftime("%d-%m-%Y")
            data["ngay"] = ngay
        except ValueError as e:
            logger.error(f"Định dạng ngày không hợp lệ: {ngay}, lỗi: {str(e)}")
            return False
        key = f"lottery-{ngay}_{dai}"  # Thêm tiền tố "lottery-" để đồng bộ
        logger.debug(f"Dữ liệu lottery trước khi lưu: {data}")
        self.lottery.document(key).set(data)
        logger.info(f"Đã lưu (hoặc ghi đè) lottery: {key}")
        return True

    def lottery_exists(self, key):
        return bool(self.lottery.document(key).get().exists)

    def get_lottery(self, dai, start_date, end_date):
        try:
            docs = self.lottery.where(filter=firestore.FieldFilter("dai", "==", dai)).stream()
            data = []
            for doc in docs:
                doc_data = doc.to_dict()
                ngay = doc_data.get("ngay")
                if not ngay or not isinstance(ngay, str):
                    logger.warning(f"Bỏ qua bản ghi có 'ngay' không hợp lệ: {doc.id}, ngay: {ngay}")
                    continue
                try:
                    doc_date = datetime.strptime(ngay, "%d-%m-%Y")
                    # Nếu start_date và end_date đều được cung cấp
                    if start_date and end_date:
                        start = datetime.strptime(start_date, "%d-%m-%Y")
                        end = datetime.strptime(end_date, "%d-%m-%Y")
                        if start <= doc_date <= end:
                            data.append(doc_data)
                    # Nếu chỉ có end_date (như trong analyze_position)
                    elif end_date:
                        end = datetime.strptime(end_date, "%d-%m-%Y")
                        if doc_date <= end:
                            data.append(doc_data)
                    # Nếu không có ngày cụ thể, lấy tất cả
                    else:
                        data.append(doc_data)
                except ValueError:
                    logger.warning(f"Ngày không hợp lệ trong bản ghi: {doc.id}, ngay: {ngay}")
                    continue
            logger.debug(f"Dữ liệu lấy từ Firestore cho {dai} từ {start_date} đến {end_date}: {data}")
            return data
        except Exception as e:
            logger.error(f"Lỗi khi lấy dữ liệu xổ số: {str(e)}")
            return []

    def get_lottery_by_date(self, dai, date):
        try:
            if not date or not isinstance(date, str):
                logger.warning(f"Ngày không hợp lệ: {date}")
                return None
            datetime.strptime(date, "%d-%m-%Y")  # Kiểm tra định dạng
            docs = self.lottery.where(filter=firestore.FieldFilter("dai", "==", dai)).where(filter=firestore.FieldFilter("ngay", "==", date)).stream()
            result = [doc.to_dict() for doc in docs if doc.to_dict().get("ngay") and isinstance(doc.to_dict().get("ngay"), str)]
            if len(result) > 1:
                logger.warning(f"Tìm thấy nhiều document khớp với {dai} ngày {date}: {len(result)} document")
            logger.debug(f"Dữ liệu từ get_lottery_by_date cho {dai} ngày {date}: {result}")
            return result[0] if result else None
        except ValueError as e:
            logger.error(f"Lỗi định dạng ngày trong get_lottery_by_date: {str(e)}")
            return None
        except Exception as e:
            logger.error(f"Lỗi khi lấy dữ liệu xổ số theo ngày: {str(e)}")
            return None